mmhdc 0.1.0__tar.gz

mmhdc-0.1.0/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/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,40 @@
+# 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/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mmhdc"
+version = "0.1.0"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { file = "LICENSE" }
+authors = [
+  { name = "Nikita Zeulin" },
+]
+dependencies = [
+  "torch",
+]
+
+[project.optional-dependencies]
+cpp = ["ninja"]
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+mmhdc = ["cpp/*.cpp"]

mmhdc-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mmhdc-0.1.0/src/mmhdc/__init__.py

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

mmhdc-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/mmhdc/utils/__init__.py

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

mmhdc-0.1.0/src/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/src/mmhdc.egg-info/PKG-INFO

@@ -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/src/mmhdc.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+LICENSE
+README.md
+pyproject.toml
+src/mmhdc/__init__.py
+src/mmhdc/_cpp_backend.py
+src/mmhdc/model.py
+src/mmhdc.egg-info/PKG-INFO
+src/mmhdc.egg-info/SOURCES.txt
+src/mmhdc.egg-info/dependency_links.txt
+src/mmhdc.egg-info/requires.txt
+src/mmhdc.egg-info/top_level.txt
+src/mmhdc/cpp/mmhdc.cpp
+src/mmhdc/utils/__init__.py
+src/mmhdc/utils/transform.py
+tests/test_step.py

mmhdc-0.1.0/src/mmhdc.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mmhdc-0.1.0/src/mmhdc.egg-info/requires.txt

@@ -0,0 +1,4 @@
+torch
+
+[cpp]
+ninja

mmhdc-0.1.0/src/mmhdc.egg-info/top_level.txt

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

mmhdc-0.1.0/tests/test_step.py

@@ -0,0 +1,165 @@
+import os
+import sys
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from example.mnist_config import get_config as _get_main_config
+from data import load_mnist
+from mmhdc import MultiMMHDC
+from mmhdc.utils import HDTransform
+
+
+_MAIN_CFG = _get_main_config()
+_MODEL_CFG = _MAIN_CFG.model
+_DEVICE = _MAIN_CFG.device
+_TRANSFORM_DTYPE = _MODEL_CFG.get("transform_dtype", None) or torch.float32
+_TRANSFORM_BATCH_SIZE = _MODEL_CFG.get("transform_batch_size", None)
+
+
+def _get_batch():
+    X_raw, y_raw, _, _ = load_mnist("mnist")
+    X = X_raw.to(dtype=torch.float32)
+    y = y_raw.to(dtype=torch.int64)
+
+    X_flat = X.reshape(X.shape[0], -1)
+    transform = HDTransform(
+        in_channels=X_flat.shape[1],
+        out_channels=_MAIN_CFG.dataset.model_dim,
+        seed=0,
+        batch_size=_TRANSFORM_BATCH_SIZE,
+        normalize=bool(_MODEL_CFG.get("normalize", True)),
+        device=_DEVICE,
+        dtype=_TRANSFORM_DTYPE,
+    )
+
+    X_hd = transform(X_flat)
+    batch_size = _MAIN_CFG.training.batch_size
+    X_batch = X_hd[:batch_size].to(_DEVICE)
+    y_batch = y[:batch_size].to(_DEVICE)
+    return X_hd, y, X_batch, y_batch
+
+
+def _make_model(backend: str, init_prototypes: torch.Tensor):
+    model = MultiMMHDC(
+        num_classes=_MAIN_CFG.dataset.num_classes,
+        out_channels=_MAIN_CFG.dataset.model_dim,
+        lr=float(_MODEL_CFG.learning_rate),
+        C=float(_MODEL_CFG.C),
+        backend=backend,
+        device=_DEVICE,
+        dtype=torch.float32,
+    )
+    model.prototypes.data = init_prototypes.detach().clone().to(_DEVICE)
+    return model
+
+
+def test_step_and_gradient_descent_losses_match(output_plot: bool = False):
+    X_hd, y_all, X_batch, y_batch = _get_batch()
+
+    # Build one common initialization so every method starts identically.
+    init_model = MultiMMHDC(
+        num_classes=_MAIN_CFG.dataset.num_classes,
+        out_channels=_MAIN_CFG.dataset.model_dim,
+        lr=float(_MODEL_CFG.learning_rate),
+        C=float(_MODEL_CFG.C),
+        backend="python",
+        device=_DEVICE,
+        dtype=torch.float32,
+    )
+
+    init_model.initialize(X_hd, y_all)
+    init_prototypes = init_model.prototypes.detach().clone()
+
+    model_cpp = _make_model("cpp", init_prototypes)
+    model_py_opt = _make_model("python", init_prototypes)
+    model_py_ref = _make_model("python", init_prototypes)
+
+    model_gd = _make_model("python", init_prototypes)
+    model_gd.prototypes = torch.nn.Parameter(
+        init_prototypes.detach().clone().to(_DEVICE),
+        requires_grad=True,
+    )
+    optimizer = torch.optim.SGD([model_gd.prototypes], lr=float(_MODEL_CFG.learning_rate))
+
+    num_steps = 50
+    losses_cpp = []
+    losses_py_opt = []
+    losses_py_ref = []
+    losses_gd = []
+
+    for _ in range(num_steps):
+        model_cpp.step(X_batch, y_batch)
+        losses_cpp.append(model_cpp.loss(X_batch, y_batch).detach())
+
+        model_py_opt._py_step(X_batch, y_batch, optimized=True)
+        losses_py_opt.append(model_py_opt.loss(X_batch, y_batch).detach())
+
+        model_py_ref._py_step(X_batch, y_batch, optimized=False)
+        losses_py_ref.append(model_py_ref.loss(X_batch, y_batch).detach())
+
+        optimizer.zero_grad()
+        model_gd.loss(X_batch, y_batch).backward()
+        optimizer.step()
+        losses_gd.append(model_gd.loss(X_batch, y_batch).detach())
+
+    losses_cpp_t = torch.stack(losses_cpp)
+    losses_py_opt_t = torch.stack(losses_py_opt)
+    losses_py_ref_t = torch.stack(losses_py_ref)
+    losses_gd_t = torch.stack(losses_gd)
+
+    diff_cpp_t = (losses_cpp_t - losses_gd_t).abs()
+    diff_py_opt_t = (losses_py_opt_t - losses_gd_t).abs()
+    diff_py_ref_t = (losses_py_ref_t - losses_gd_t).abs()
+
+    print("\n[step-loss-table]")
+    print("step |          sgd |          cpp |       py_opt |       py_ref |   |cpp-sgd| | |py_opt-sgd| | |py_ref-sgd|")
+    print("-----+--------------+--------------+--------------+--------------+-------------+---------------+--------------")
+    for step_idx in range(num_steps):
+        print(
+            f"{step_idx + 1:>4} | "
+            f"{losses_gd_t[step_idx].item():>12.6f} | "
+            f"{losses_cpp_t[step_idx].item():>12.6f} | "
+            f"{losses_py_opt_t[step_idx].item():>12.6f} | "
+            f"{losses_py_ref_t[step_idx].item():>12.6f} | "
+            f"{diff_cpp_t[step_idx].item():>11.6f} | "
+            f"{diff_py_opt_t[step_idx].item():>13.6f} | "
+            f"{diff_py_ref_t[step_idx].item():>12.6f}"
+        )
+
+    if output_plot:
+        import matplotlib
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+
+        output_dir = os.path.join(os.path.dirname(__file__), "output")
+        os.makedirs(output_dir, exist_ok=True)
+        plot_path = os.path.join(output_dir, "step_vs_gd_losses.png")
+
+        steps = range(1, num_steps + 1)
+        plt.figure(figsize=(9, 5))
+        plt.plot(steps, losses_gd_t.cpu().numpy(), label="SGD", marker="o")
+        plt.plot(steps, losses_cpp_t.cpu().numpy(), label="C++", marker="s")
+        plt.plot(steps, losses_py_opt_t.cpu().numpy(), label="Python optimized", marker="^")
+        plt.plot(steps, losses_py_ref_t.cpu().numpy(), label="Python reference", marker="d")
+        plt.xlabel("Step")
+        plt.ylabel("Loss")
+        plt.title("Step Procedure Loss Comparison")
+        plt.legend()
+        plt.tight_layout()
+        plt.savefig(plot_path, dpi=150)
+        plt.close()
+        print(f"[step-loss-plot] saved: {plot_path}")
+
+    assert torch.allclose(losses_cpp_t, losses_gd_t, atol=1e-4, rtol=1e-4), (
+        "C++ step losses diverge from SGD. "
+        f"max_abs_diff={(losses_cpp_t - losses_gd_t).abs().max().item():.3e}"
+    )
+    assert torch.allclose(losses_py_opt_t, losses_gd_t, atol=1e-4, rtol=1e-4), (
+        "Python optimized step losses diverge from SGD. "
+        f"max_abs_diff={(losses_py_opt_t - losses_gd_t).abs().max().item():.3e}"
+    )
+    assert torch.allclose(losses_py_ref_t, losses_gd_t, atol=1e-4, rtol=1e-4), (
+        "Python reference step losses diverge from SGD. "
+        f"max_abs_diff={(losses_py_ref_t - losses_gd_t).abs().max().item():.3e}"
+    )