mmhdc 0.1.0__py3-none-any.whl

mmhdc/__init__.py

@@ -0,0 +1,3 @@
+from .model import MultiMMHDC
+
+__all__ = ["MultiMMHDC"]

mmhdc/_cpp_backend.py

@@ -0,0 +1,22 @@
+from functools import lru_cache
+from pathlib import Path
+
+
+@lru_cache(maxsize=1)
+def get_mmhdc_cpp():
+    try:
+        from torch.utils.cpp_extension import load
+
+        return load(
+            name="mmhdc_cpp",
+            extra_cflags=["-O3"],
+            is_python_module=True,
+            sources=[str(Path(__file__).resolve().parent / "cpp" / "mmhdc.cpp")],
+        )
+    except Exception as exc:
+        raise RuntimeError(
+            "Failed to build or load the MM-HDC C++ backend. "
+            "The 'cpp' backend requires a working C++ toolchain and PyTorch C++ "
+            "extension build support. If this environment cannot compile the "
+            "extension, use backend='python' instead."
+        ) from exc

mmhdc/cpp/mmhdc.cpp

@@ -0,0 +1,51 @@
+#include <torch/extension.h>
+
+/*
+Optimised step procedure — global GEMM reformulation:
+
+  All per-class loops are eliminated. The entire batch is handled with
+  exactly two GEMMs regardless of the number of classes K:
+
+    1. scores   = x @ prototypes.T          (N, K)  — one GEMM
+    2. W        built from the boolean margin-violation mask (N, K)
+    3. update   = W.T @ x                   (K, D)  — one GEMM
+
+  This maximises GPU utilisation (two large, dense kernels instead of K
+  small ones) and removes all GPU→CPU syncs (.item() calls).
+*/
+
+torch::Tensor step(torch::Tensor &x, torch::Tensor &y, torch::Tensor &prototypes, float lr, float C) {
+    // scores[i, k] = x_i · w_k  for every sample and every class.
+    auto scores = torch::mm(x, prototypes.t());                    // (N, K)
+
+    // correct_scores[i] = x_i · w_{y_i}
+    auto correct_scores = scores.gather(1, y.unsqueeze(1));        // (N, 1)
+
+    // margin_gap[i, k] = x_i · w_{y_i} - x_i · w_k
+    //                  = correct_score_i  - score_{i,k}
+    // Margin is violated when gap < 2  (and k ≠ y_i).
+    auto violated = (correct_scores - scores) < 2;                 // (N, K) bool
+
+    // Zero out the diagonal: a sample cannot violate a margin against its own class.
+    violated.scatter_(1, y.unsqueeze(1),
+                      torch::zeros_like(y.unsqueeze(1), torch::kBool));
+
+    // Build weight matrix W  (N, K):
+    //   W[i, k]   = -1  if sample i violates margin against class k
+    //   W[i, y_i] = +m_i where m_i is the number of violated margins for sample i
+    //               (matches Python: each sample contributes once per violated class).
+    auto W = -violated.to(x.scalar_type());                        // (N, K)
+    W.scatter_add_(1, y.unsqueeze(1),
+                   violated.sum(1, /*keepdim=*/true).to(x.scalar_type()));
+
+    // prototypes_update[k] = Σ_i  W[i,k] * x_i  — single GEMM.
+    auto prototypes_update = torch::mm(W.t(), x);                  // (K, D)
+
+    // Regularised update.
+    prototypes = (1 - lr / C) * prototypes + lr * prototypes_update;
+    return prototypes;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("step", &step, "MM-HDC prototype update function");
+}

mmhdc/model.py

@@ -0,0 +1,110 @@
+import torch
+from torch.nn.functional import one_hot, relu
+from ._cpp_backend import get_mmhdc_cpp
+
+class MultiMMHDC(torch.nn.Module):
+    def __init__(self, num_classes: int, 
+                 out_channels: int, 
+                 lr: float = 1e-2, 
+                 C: float = 1.0, 
+                 device: str = 'cpu',
+                 backend: str = 'python',
+                 dtype: torch.dtype = torch.float32):
+        
+        super().__init__()
+        self.num_classes = num_classes
+        self.out_channels = out_channels
+        self.lr = lr
+        self.device = device
+        self.C = C
+        self.backend = backend
+        self.prototypes = torch.nn.parameter.Parameter(
+            data=torch.zeros(num_classes, out_channels, dtype=dtype, device=device), 
+            requires_grad=False
+        )
+        self.dtype = dtype
+
+    def forward(self, x: torch.Tensor):
+        return torch.argmax(x @ self.prototypes.T, dim=1)
+    
+    @torch.no_grad()
+    def initialize(self, x: torch.Tensor, y: torch.Tensor):
+        y = y.reshape(-1).to(dtype=torch.int64, device=x.device)
+        x = x.to(dtype=self.prototypes.dtype)
+
+        prototype_sums = torch.zeros(
+            self.num_classes,
+            self.out_channels,
+            dtype=self.prototypes.dtype,
+            device=x.device,
+        )
+        prototype_sums.index_add_(0, y, x)
+
+        class_counts = torch.bincount(y, minlength=self.num_classes)
+        class_counts = class_counts.unsqueeze(1).to(dtype=self.prototypes.dtype, device=x.device)
+
+        prototypes = prototype_sums / class_counts.clamp_min_(1)
+        prototype_norms = torch.norm(prototypes, dim=1, keepdim=True)
+        prototypes = prototypes / prototype_norms.clamp_min_(1e-8)
+
+        self.prototypes.copy_(prototypes.to(device=self.prototypes.device))
+
+    def loss(self, X: torch.Tensor, y: torch.Tensor):
+        y = y.reshape(-1)
+        scores = X @ self.prototypes.T
+        correct_scores = scores.gather(1, y.unsqueeze(1))
+        margins = relu(2 - (correct_scores - scores))
+        true_class_mask = one_hot(y, num_classes=self.num_classes).to(dtype=torch.bool)
+        margins = margins.masked_fill(true_class_mask, 0.0)
+
+        regularizer = self.prototypes.square().sum() / (2 * self.C)
+        return regularizer + margins.sum()
+        
+    def step(self, x: torch.Tensor, y: torch.Tensor):
+        if self.backend == 'cpp':
+            with torch.no_grad():
+                mmhdc_cpp = get_mmhdc_cpp()
+                updated_prototypes = mmhdc_cpp.step(x, y, self.prototypes, self.lr, self.C)
+                self.prototypes.copy_(updated_prototypes)
+        elif self.backend == 'python':
+            self._py_step(x, y)
+        else:
+            raise ValueError(f"Unsupported backend '{self.backend}'. Expected 'cpp' or 'python'.")
+
+    @torch.no_grad()
+    def _py_step(self, x: torch.Tensor, y: torch.Tensor, optimized=True):
+        # This implementation is for reference purpose only. Use matrix-based optimized C++ and Python implementations
+        def _py_step_reference(x: torch.Tensor, y: torch.Tensor):
+            prototypes_update = torch.zeros_like(self.prototypes, dtype=self.dtype)
+            for cls in y.unique():
+                rolled_prototypes = torch.roll(self.prototypes, -cls.item(), dims=0)
+                x_cls = x[y == cls]
+
+                dot = x_cls @ (rolled_prototypes[0] - rolled_prototypes[1:]).T
+                hinge_loss = relu(2 - dot)
+
+                exceeding_margin = hinge_loss > 0
+                num_violations = exceeding_margin.sum(dim=1, dtype=x_cls.dtype)
+                prototypes_update[cls] += (x_cls * num_violations.unsqueeze(1)).sum(0)
+
+                y_true_all = exceeding_margin.any(0).nonzero().flatten()
+                for y_true in y_true_all:
+                    prototypes_update[(y_true + 1 + cls) % self.num_classes] -= x_cls[exceeding_margin[:, y_true]].sum(0)
+            
+            self.prototypes.data = (1 - self.lr / self.C) * self.prototypes.data + self.lr * prototypes_update
+        
+        # Optimized code based on the C++ implementation logic
+        def _py_step_optimized(x: torch.Tensor, y: torch.Tensor):
+            scores = x @ self.prototypes.T
+            correct_scores = scores.gather(1, y.unsqueeze(1))
+
+            violated = (correct_scores - scores) < 2
+            violated.scatter_(1, y.unsqueeze(1), False)
+
+            W = -violated.to(x.dtype)
+            W.scatter_add_(1, y.unsqueeze(1), violated.sum(dim=1, keepdim=True).to(x.dtype))
+
+            prototypes_update = W.T @ x
+            self.prototypes.data = (1 - self.lr / self.C) * self.prototypes.data + self.lr * prototypes_update
+
+        _py_step_reference(x, y) if not optimized else _py_step_optimized(x, y)

mmhdc/utils/__init__.py

@@ -0,0 +1,3 @@
+from .transform import HDTransform
+
+__all__ = ["HDTransform"]

mmhdc/utils/transform.py

@@ -0,0 +1,67 @@
+import torch
+import math
+from typing import Optional
+
+class HDTransform(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        seed: int = 0,
+        batch_size: Optional[int] = None,
+        normalize: bool = True,
+        device: str = 'cpu',
+        dtype: torch.dtype = torch.float32,):
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.batch_size = batch_size
+        self.normalize = normalize
+        self.dtype = dtype
+
+        self._rng = torch.Generator(device="cpu")
+        self._rng.manual_seed(seed)
+        two_pi = torch.tensor(2.0 * math.pi, dtype=self.dtype)
+        self.register_buffer(
+            "_eps",
+            torch.tensor(1e-8, dtype=self.dtype, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_G",
+            torch.randn(in_channels, out_channels, generator=self._rng, dtype=self.dtype).to(device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_phi",
+            torch.rand(out_channels, generator=self._rng, dtype=self.dtype).mul(two_pi).to(device),
+            persistent=False,
+        )
+
+    def _transform(self, x: torch.Tensor):
+        projection_matrix = self._G.unsqueeze(0).expand(x.size(0), -1, -1)
+
+        # OnlineHD feature mapping
+        buf1 = torch.bmm(x.unsqueeze(1), projection_matrix).squeeze(1)
+        buf2 = buf1 + self._phi
+        return buf1.sin_() * buf2.cos_()
+
+    def forward(self, data: torch.Tensor):
+        with torch.no_grad():
+            data = data.to(device=self._G.device, dtype=self.dtype)
+            if self.normalize:
+                norms = torch.norm(data, dim=1, keepdim=True)
+                data = data / torch.maximum(norms, self._eps)
+
+            if self.batch_size is None or self.batch_size >= data.size(0):
+                x_new = self._transform(data)
+                return x_new.cpu() if x_new.device.type != 'cpu' else x_new
+
+            x_new = torch.empty(data.size(0), self.out_channels, dtype=self.dtype)
+            for start in range(0, data.size(0), self.batch_size):
+                end = start + self.batch_size
+                batch = self._transform(data[start:end])
+                x_new[start:end].copy_(batch.cpu() if batch.device.type != 'cpu' else batch)
+
+            return x_new

mmhdc-0.1.0.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: mmhdc
+Version: 0.1.0
+Author: Nikita Zeulin
+License: MIT License
+        
+        Copyright (c) 2026 Nikita Zeulin
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Provides-Extra: cpp
+Requires-Dist: ninja; extra == "cpp"
+Dynamic: license-file
+
+# MM-HDC: Maximum-Margin Hyperdimensional Computing
+
+This repository provides an implementation of the multi-class maximum-margin hyperdimensional computing (MM-HDC) classifier that adopts the optimization problem formulation of multi-class Weston-Watkins SVM to HDC.
+
+The algorithm is implemented in the `MultiMMHDC` class, which can be imported as:
+
+```python
+from mmhdc import MultiMMHDC
+from mmhdc.utils import HDTransform
+```
+
+You can run the MNIST example from the GitHub repo as follows:
+
+```bash
+git clone https://github.com/nzeulin/mmhdc.git && cd mmhdc
+python -m pip install --upgrade pip && python -m pip install mmhdc
+python -m pip install -r example/requirements.txt
+python example/example.py --config example/mnist_config.py
+```
+
+C++ backend is highly recommended to use, as it can significantly accelerate the model training. In this case, use `python -m pip install mmhdc[cpp]`. 
+
+**NOTE:** If you don't have `gcc` installed in your system (required for compiling C++ module), you can install it as a Conda package into your venv: `conda install -c conda-forge gxx_linux-64`.
+
+### Current features
+- C++ backend (`libtorch`) to enable fast MM-HDC training.
+- Support of floating-point prototypes and hypervectors.
+ 
+### Citation
+
+If you use this repository in your research, please cite it as software (paper in progress):
+
+```bibtex
+@software{zeulin_2026_mm_hdc,
+	author = {Zeulin, Nikita},
+	title = {MM-HDC: Maximum-Margin Hyperdimensional Computing},
+	year = {2026},
+	url = {https://github.com/nzeulin/mm-hdc}
+}
+```

mmhdc-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+mmhdc/__init__.py,sha256=qqdxFdDJteLA5dCynf05KmwGNRNYatBepXfcxZvnggg,56
+mmhdc/_cpp_backend.py,sha256=Op6fvnlRmGNYFbTSs62V6cCF0OLGH82MsybSFhYV3p4,730
+mmhdc/model.py,sha256=mXev-BcVWySuKDJZ9PPusuhnzdqKgUCcPlon99hItLU,4755
+mmhdc/cpp/mmhdc.cpp,sha256=q2ysINLkEbbAqxzNYYxNNrjxIKlv6OYWeLjNa5ls3ZM,2222
+mmhdc/utils/__init__.py,sha256=OkPk7PiAuSSc8DvMJfCEon2Uzd1HuvWeFZsde6dcgqc,62
+mmhdc/utils/transform.py,sha256=QgBqg4NGlCQgwGMIwVbOg4grlGiXwjMNiZrf8SlxPT0,2405
+mmhdc-0.1.0.dist-info/licenses/LICENSE,sha256=gPQaAqyDO7WVlNm-korHqkViDn7Z8S1xDqcq74Rs-G8,1070
+mmhdc-0.1.0.dist-info/METADATA,sha256=it4IUEDuxRH-WOAIDOqi95kgqrysl2tUBodCFa7FC5A,3041
+mmhdc-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+mmhdc-0.1.0.dist-info/top_level.txt,sha256=IJcDmUlMuHSA7ub1Ge-093KZECZMQ4fFwF5DpL4X_sk,6
+mmhdc-0.1.0.dist-info/RECORD,,

mmhdc-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

mmhdc-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Nikita Zeulin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mmhdc-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+mmhdc