Performance: fold cache, rank-1 MSTDP, cached MSTDPET decays

- topology.py: cache the folded (a_eff, b_sum) whenever every pipeline
  feature is static (is_static flag; Probability, MeanField, adaptation
  features, and sub-features are dynamic). Invalidated on learning
  updates, normalize, reset, and device/dtype moves. Weight*Mask*
  Intensity at 100x20000: 4506 -> 66 us/step CPU, 321 -> 58 us GPU.
- topology_features.py: Weight's per-time-step normalization is deferred
  by the connection until after the fold consumes the pre-norm value,
  removing the full-matrix clone every step. Fix the reduction assertion
  (isinstance(x, callable) raised TypeError for any custom reduction).
- MCC_learning.py: MSTDP with the default batch reductions applies the
  reward-weighted eligibility as two rank-1 addmm_ updates instead of
  materializing [batch, src, tgt] every step (5077 -> 899 us CPU,
  834 -> 332 us GPU at 100x20000; exact same numbers, verified to 1e-7).
  Dense path kept for averaging buffers, custom reductions, and sparse
  weights. Skip the decay multiply when decay is the no-op 1.0. Give
  MSTDPET the same (dt, device)-keyed decay/learning-rate cache as MSTDP
  (was 3 exp() + 2 host-to-device copies per step).
- New tests: rank-1 MSTDP == dense eligibility (batch 1/4, tensor
  rewards), fold-cache storage + invalidation, deferred time-step norm.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Hananel-Hazan

bindsnet/datasets/torchvision_wrapper.py

@@ -16,18 +16,13 @@ def create_torchvision_dataset_wrapper(ds_type):
         ds_type = getattr(torchDB, ds_type)
 
     class TorchvisionDatasetWrapper(ds_type):
-        __doc__ = (
-            """BindsNET torchvision dataset wrapper for:
+        __doc__ = """BindsNET torchvision dataset wrapper for:
 
         The core difference is the output of __getitem__ is no longer
         (image, label) rather a dictionary containing the image, label,
         and their encoded versions if encoders were provided.
 
-            \n\n"""
-            + str(ds_type)
-            if ds_type.__doc__ is None
-            else ds_type.__doc__
-        )
+            \n\n""" + str(ds_type) if ds_type.__doc__ is None else ds_type.__doc__
 
         def __init__(
             self,

bindsnet/learning/MCC_learning.py

@@ -89,8 +89,9 @@ def update(self, **kwargs) -> None:
         Abstract method for a learning rule update.
         """
 
-        # Implement decay.
-        if self.decay:
+        # Implement decay (self.decay == 1.0 is the no-op default; skip the
+        # full-matrix multiply in that case).
+        if self.decay != 1.0:
             self.feature_value *= self.decay
 
         # Enforce polarities
@@ -497,11 +498,6 @@ def _connection_update(self, **kwargs) -> None:
                 self.target.n,
                 device=self.target.s.device,
             )
-        if not hasattr(self, "eligibility"):
-            self.eligibility = torch.zeros(
-                batch_size, *self.feature_value.shape, device=self.feature_value.device
-            )
-
         # Reshape pre- and post-synaptic spikes.
         source_s = self.source.s.view(batch_size, -1).float()
         target_s = self.target.s.view(batch_size, -1).float()
@@ -532,38 +528,82 @@ def _connection_update(self, **kwargs) -> None:
             else torch.as_tensor(a_minus, device=self.feature_value.device)
         )
 
-        # Compute weight update based on the eligibility value of the past timestep.
-        update = reward * self.eligibility
+        # With no averaging and the standard batch reductions (squeeze for
+        # batch size 1, sum otherwise), the [batch, src, tgt] eligibility never
+        # needs to be materialized: its batch reduction is a sum of two matrix
+        # products applied to the weights directly. The eligibility of the past
+        # timestep is outer(p_plus, target_s) + outer(source_s, p_minus), with
+        # p_plus/p_minus still holding their previous-step values here.
+        fast = (
+            self.average_update == 0
+            and self.reduction in (torch.squeeze, torch.sum)
+            and not self.feature_value.is_sparse
+        )
+        if fast:
+            if hasattr(self, "_prev_target_s"):
+                if isinstance(reward, torch.Tensor):
+                    # Keep reward on-device (no host sync for tensor rewards).
+                    update = (
+                        self.p_plus.t() @ self._prev_target_s
+                        + self._prev_source_s.t() @ self.p_minus
+                    )
+                    self.feature_value += (self.nu[0].to(dev) * reward) * update
+                else:
+                    alpha = float(self.nu[0]) * reward
+                    if alpha != 0.0:
+                        self.feature_value.addmm_(
+                            self.p_plus.t(), self._prev_target_s, alpha=alpha
+                        )
+                        self.feature_value.addmm_(
+                            self._prev_source_s.t(), self.p_minus, alpha=alpha
+                        )
+        else:
+            # Dense-eligibility path: averaging buffers, custom reductions, or
+            # sparse weights.
+            if not hasattr(self, "eligibility"):
+                self.eligibility = torch.zeros(
+                    batch_size,
+                    *self.feature_value.shape,
+                    device=self.feature_value.device,
+                )
 
-        if self.average_update > 0:
-            self.average_buffer[self.average_buffer_index] = self.reduction(
-                update, dim=0
-            )
-            self.average_buffer_index = (
-                self.average_buffer_index + 1
-            ) % self.average_update
+            # Compute weight update based on the eligibility value of the past timestep.
+            update = reward * self.eligibility
 
-            if self.continues_update or self.average_buffer_index == 0:
-                update = self.nu[0] * torch.mean(self.average_buffer, dim=0)
+            if self.average_update > 0:
+                self.average_buffer[self.average_buffer_index] = self.reduction(
+                    update, dim=0
+                )
+                self.average_buffer_index = (
+                    self.average_buffer_index + 1
+                ) % self.average_update
+
+                if self.continues_update or self.average_buffer_index == 0:
+                    update = self.nu[0] * torch.mean(self.average_buffer, dim=0)
+                    if self.feature_value.is_sparse:
+                        update = update.to_sparse()
+                    self.feature_value += update
+            else:
+                update = self.nu[0] * self.reduction(update, dim=0)
                 if self.feature_value.is_sparse:
                     update = update.to_sparse()
                 self.feature_value += update
-        else:
-            update = self.nu[0] * self.reduction(update, dim=0)
-            if self.feature_value.is_sparse:
-                update = update.to_sparse()
-            self.feature_value += update
 
         # Update P^+ and P^- values.
         self.p_plus *= self._decay_plus
         self.p_plus += a_plus * source_s
         self.p_minus *= self._decay_minus
         self.p_minus += a_minus * target_s
 
-        # Calculate point eligibility value.
-        self.eligibility = torch.bmm(
-            self.p_plus.unsqueeze(2), target_s.unsqueeze(1)
-        ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1))
+        if fast:
+            # Keep this step's spikes for the next step's rank-1 update.
+            self._prev_source_s = source_s.clone()
+            self._prev_target_s = target_s.clone()
+        else:
+            # Calculate point eligibility value.
+            self.eligibility = torch.bmm(
+                self.p_plus.unsqueeze(2), target_s.unsqueeze(1)
+            ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1))
 
         super().update()
 
@@ -688,26 +728,36 @@ def _connection_update(self, **kwargs) -> None:
 
         # Parse keyword arguments.
         reward = kwargs["reward"]
-        a_plus = kwargs.get("a_plus", 1.0)
-        # if isinstance(a_plus, dict):
-        #     for k, v in a_plus.items():
-        #         a_plus[k] = torch.tensor(v, device=self.feature_value.device)
-        # else:
-        a_plus = torch.tensor(a_plus, device=self.feature_value.device)
-        a_minus = kwargs.get("a_minus", -1.0)
-        # if isinstance(a_minus, dict):
-        #     for k, v in a_minus.items():
-        #         a_minus[k] = torch.tensor(v, device=self.feature_value.device)
-        # else:
-        a_minus = torch.tensor(a_minus, device=self.feature_value.device)
+
+        # Build learning-rate and decay tensors, cached on (dt, device) so a
+        # change to either (network.dt edits, .to(device) moves) recomputes them.
+        dev = self.feature_value.device
+        dt = float(self.connection.dt)
+        if getattr(self, "_decay_key", None) != (dt, dev):
+            self._decay_key = (dt, dev)
+            self._a_plus_default = torch.tensor(1.0, device=dev)
+            self._a_minus_default = torch.tensor(-1.0, device=dev)
+            self._decay_plus = torch.exp(-dt / self.tc_plus.to(dev))
+            self._decay_minus = torch.exp(-dt / self.tc_minus.to(dev))
+            self._decay_e_trace = torch.exp(-dt / self.tc_e_trace.to(dev))
+        a_plus = kwargs.get("a_plus", None)
+        a_plus = (
+            self._a_plus_default
+            if a_plus is None
+            else torch.as_tensor(a_plus, device=dev)
+        )
+        a_minus = kwargs.get("a_minus", None)
+        a_minus = (
+            self._a_minus_default
+            if a_minus is None
+            else torch.as_tensor(a_minus, device=dev)
+        )
 
         # Calculate value of eligibility trace based on the value
         # of the point eligibility value of the past timestep.
         # Note: eligibility = [source.n, target.n] > 0 where source and target spiked
         # Note: high negs. ->
-        self.eligibility_trace *= torch.exp(
-            -self.connection.dt / self.tc_e_trace
-        )  # Decay
+        self.eligibility_trace *= self._decay_e_trace  # Decay
         self.eligibility_trace += self.eligibility / self.tc_e_trace  # Additive changes
         # ^ Also effected by delay in last step
 
@@ -733,9 +783,9 @@ def _connection_update(self, **kwargs) -> None:
             self.feature_value += update
 
         # Update P^+ and P^- values.
-        self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus)  # Decay
+        self.p_plus *= self._decay_plus  # Decay
         self.p_plus += a_plus * source_s  # Scaled source spikes
-        self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus)  # Decay
+        self.p_minus *= self._decay_minus  # Decay
         self.p_minus += a_minus * target_s  # Scaled target spikes
 
         # Notes:

bindsnet/network/topology.py

@@ -442,6 +442,15 @@ def __init__(
 
         self.sparse_compute = sparse_compute
 
+        # Cached (a_eff, b_sum) for pipelines whose features are all static
+        # (see AbstractFeature.is_static). Invalidated whenever a feature can
+        # change: learning updates, normalize, reset, device/dtype moves.
+        self._fold_cache = None
+
+    def _apply(self, fn, recurse=True):
+        self._fold_cache = None
+        return super()._apply(fn, recurse)
+
     def compute(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
         """
@@ -463,34 +472,50 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
         """
         s = s.view(s.size(0), self.source.n)
 
-        # running product of multiplicative factors, [source.n, target.n]
-        a_eff = None
-        # running additive offset, [source.n, target.n]; None while still zero
-        b_eff = None
-        for f in self.pipeline:
-            factor = f.compute(s)
-            if factor is None:
-                # Side-effect-only pipeline entries (sub-features) fold as identity.
-                continue
-            if isinstance(factor, torch.Tensor) and factor.is_sparse:
-                # Sparse feature values carry a leading batch dim ([1, src, tgt])
-                # from prime_feature; densify to the fold's [src, tgt] shape.
-                factor = factor.to_dense().view(self.source.n, self.target.n)
-            op = getattr(f, "op", "mul")
-            if op == "mul":
-                a_eff = factor if a_eff is None else a_eff * factor
-                if b_eff is not None:
-                    b_eff = b_eff * factor
-            else:  # additive contribution: "add" -> +factor, "sub" -> -factor
-                term = factor if op == "add" else -factor
-                b_eff = term if b_eff is None else b_eff + term
-
-        if a_eff is None:
-            # Degenerate pipeline with no multiplicative feature: every source
-            # neuron contributes with unit weight.
-            a_eff = torch.ones(self.source.n, self.target.n, device=s.device)
-        if not torch.is_floating_point(a_eff):
-            a_eff = a_eff.float()
+        deferred = None
+        if self._fold_cache is not None:
+            a_eff, b_sum = self._fold_cache
+        else:
+            # running product of multiplicative factors, [source.n, target.n]
+            a_eff = None
+            # running additive offset, [source.n, target.n]; None while still zero
+            b_eff = None
+            for f in self.pipeline:
+                factor = f.compute(s)
+                # Compute-time side effects (e.g. per-time-step weight
+                # normalization) run after the fold has consumed this value.
+                d = getattr(f, "defer", None)
+                if d is not None:
+                    deferred = [d] if deferred is None else deferred + [d]
+                if factor is None:
+                    # Side-effect-only pipeline entries (sub-features) fold as identity.
+                    continue
+                if isinstance(factor, torch.Tensor) and factor.is_sparse:
+                    # Sparse feature values carry a leading batch dim ([1, src, tgt])
+                    # from prime_feature; densify to the fold's [src, tgt] shape.
+                    factor = factor.to_dense().view(self.source.n, self.target.n)
+                op = getattr(f, "op", "mul")
+                if op == "mul":
+                    a_eff = factor if a_eff is None else a_eff * factor
+                    if b_eff is not None:
+                        b_eff = b_eff * factor
+                else:  # additive contribution: "add" -> +factor, "sub" -> -factor
+                    term = factor if op == "add" else -factor
+                    b_eff = term if b_eff is None else b_eff + term
+
+            if a_eff is None:
+                # Degenerate pipeline with no multiplicative feature: every source
+                # neuron contributes with unit weight.
+                a_eff = torch.ones(self.source.n, self.target.n, device=s.device)
+            if not torch.is_floating_point(a_eff):
+                a_eff = a_eff.float()
+
+            # Additive terms apply to every synapse regardless of spikes, so
+            # their contribution is the source-sum, a constant [target.n] row.
+            b_sum = b_eff.sum(dim=0) if b_eff is not None else None
+
+            if all(getattr(f, "is_static", True) for f in self.pipeline):
+                self._fold_cache = (a_eff, b_sum)
 
         # The gather pays for its ``nonzero()`` only where that sync is expensive
         # relative to the matmul: on CUDA it needs a large weight matrix to win;
@@ -511,11 +536,12 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
         else:
             out = s.to(a_eff.dtype) @ a_eff
 
-        # Additive terms apply to every synapse regardless of spikes, so sum over
-        # all source rows (the closed form of the [batch, source.n, target.n]
-        # source-sum).
-        if b_eff is not None:
-            out = out + b_eff.sum(dim=0)
+        if b_sum is not None:
+            out = out + b_sum
+
+        if deferred is not None:
+            for fn in deferred:
+                fn()
 
         if self.traces:
             self.activity = out
@@ -559,6 +585,7 @@ def update(self, **kwargs) -> None:
         learning = kwargs.get("learning", False)
         if learning and not self.manual_update:
             # Pipeline learning
+            self._fold_cache = None
             for f in self.pipeline:
                 f.update(**kwargs)
 
@@ -568,6 +595,7 @@ def normalize(self) -> None:
         Normalize all features in the connection.
         """
         # Normalize pipeline features
+        self._fold_cache = None
         for f in self.pipeline:
             f.normalize()
 
@@ -578,6 +606,7 @@ def reset_state_variables(self) -> None:
         """
         super().reset_state_variables()
 
+        self._fold_cache = None
         for f in self.pipeline:
             f.reset_state_variables()