ffssn

Author: Kevin Chang

bindsnet/network/topology.py

@@ -2049,3 +2049,131 @@ def reset_state_variables(self) -> None:
         Contains resetting logic for the connection.
         """
         super().reset_state_variables()
+
+
+class ForwardForwardConnection(AbstractConnection):
+    """
+    Connection class specifically designed for Forward-Forward training with arctangent surrogate gradients.
+    """
+    
+    def __init__(
+        self,
+        source: Nodes,
+        target: Nodes,
+        nu: Optional[Union[float, Sequence[float], Sequence[torch.Tensor]]] = None,
+        weight_decay: float = 0.0,
+        spike_threshold: float = 1.0,
+        alpha: float = 2.0,  # α parameter for arctangent surrogate
+        **kwargs,
+    ) -> None:
+        super().__init__(source, target, nu, weight_decay, **kwargs)
+        
+        # Initialize weights with gradient support
+        w = kwargs.get("w", None)
+        if w is None:
+            if (self.wmin == -np.inf).any() or (self.wmax == np.inf).any():
+                w = torch.clamp(torch.randn(source.n, target.n) * 0.1, self.wmin, self.wmax)
+            else:
+                w = self.wmin + (torch.randn(source.n, target.n) * 0.1) * (self.wmax - self.wmin)
+        else:
+            if (self.wmin == -np.inf).any() or (self.wmax == np.inf).any():
+                w = torch.clamp(w, self.wmin, self.wmax)
+        
+        # CRITICAL: Enable gradients for Forward-Forward training
+        self.w = Parameter(w, requires_grad=True)
+        
+        # Surrogate gradient parameters
+        self.spike_threshold = spike_threshold
+        self.alpha = alpha
+        
+        # Track membrane potential for surrogate gradients
+        self.membrane_potential = None
+        
+    def atan_surrogate_spike(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Arctangent surrogate gradient function.
+        
+        Forward pass: Heaviside step function shifted by threshold
+        Backward pass: Gradient of shifted arc-tan function with parameter α
+        """
+        class AtanSurrogate(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, input, threshold, alpha):
+                ctx.save_for_backward(input)
+                ctx.threshold = threshold
+                ctx.alpha = alpha
+                # Forward: Heaviside step function shifted by threshold
+                return (input > threshold).float()
+            
+            @staticmethod
+            def backward(ctx, grad_output):
+                input, = ctx.saved_tensors
+                grad_input = grad_output.clone()
+                # Backward: Gradient of shifted arc-tan function
+                # surrogate = 1 / (α * |input - threshold| + 1)
+                surrogate_grad = 1.0 / (ctx.alpha * torch.abs(input - ctx.threshold) + 1.0)
+                return grad_input * surrogate_grad, None, None
+        
+        return AtanSurrogate.apply(x, self.spike_threshold, self.alpha)
+
+    def compute_with_surrogate(self, s: torch.Tensor) -> torch.Tensor:
+        """
+        Compute pre-activations with arctangent surrogate gradients.
+        
+        :param s: Incoming spikes [batch_size, source_neurons]
+        :return: Output spikes with surrogate gradients [batch_size, target_neurons]
+        """
+        batch_size = s.shape[0]
+        
+        # Initialize membrane potential if needed
+        if self.membrane_potential is None or self.membrane_potential.shape != (batch_size, self.target.n):
+            self.membrane_potential = torch.zeros(batch_size, self.target.n, device=s.device)
+        
+        # Synaptic input: spikes @ weights
+        synaptic_input = torch.mm(s.float(), self.w)
+        
+        # Simple LIF dynamics with decay (you can customize this)
+        decay_factor = 0.9  # Can be made configurable
+        self.membrane_potential = decay_factor * self.membrane_potential + synaptic_input
+        
+        # Generate spikes with arctangent surrogate gradients
+        spikes = self.atan_surrogate_spike(self.membrane_potential)
+        
+        # Reset mechanism: subtract threshold from membrane potential where spikes occurred
+        self.membrane_potential = self.membrane_potential - spikes * self.spike_threshold
+        
+        return spikes
+
+    def compute(self, s: torch.Tensor) -> torch.Tensor:
+        """
+        Standard compute method (calls compute_with_surrogate for FF training).
+        """
+        return self.compute_with_surrogate(s)
+
+    def reset_membrane_potential(self):
+        """Reset membrane potential (call between samples/batches)."""
+        self.membrane_potential = None
+
+    def update(self, **kwargs) -> None:
+        """
+        Override standard BindsNET update - FF uses PyTorch optimizers.
+        """
+        # Forward-Forward training uses PyTorch optimizers for weight updates
+        # So we don't need the standard BindsNET learning rule updates
+        pass
+
+    def normalize(self) -> None:
+        """
+        Normalize weights so each target neuron has sum of connection weights equal to self.norm.
+        """
+        if self.norm is not None:
+            w_abs_sum = self.w.abs().sum(0).unsqueeze(0)
+            w_abs_sum[w_abs_sum == 0] = 1.0
+            self.w.data *= self.norm / w_abs_sum
+
+    def reset_state_variables(self) -> None:
+        """
+        Contains resetting logic for the connection.
+        """
+        super().reset_state_variables()
+        self.reset_membrane_potential()