finished

Author: Kevin Chang

bindsnet/network/topology.py

@@ -2052,133 +2052,3 @@ def reset_state_variables(self) -> None:
 
 
 
-# designed for ANN
-# the weights are not a true connection
-# gradient 
-
-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()

bindsnet/network/topology_features.py

@@ -955,23 +955,38 @@ class ArctangentSurrogateFeature(AbstractFeature):
 
     def __init__(
         self,
+        name: str,
         spike_threshold: float = 1.0,
         alpha: float = 2.0,
         dt: float = 1.0,
         reset_mechanism: str = "subtract",
+        value: Union[torch.Tensor, float, int] = None,
+        range: Optional[Union[list, tuple]] = None,
         **kwargs
     ):
         """
         Initialize arctangent surrogate feature.
         
         Args:
+            name: Name of the feature
             spike_threshold: Voltage threshold for spike generation
             alpha: Steepness parameter for surrogate gradient (higher = steeper)
             dt: Integration time step
             reset_mechanism: Post-spike reset ("subtract", "zero", "none")
+            value: Initial membrane potential values (optional)
+            range: Range of acceptable values for membrane potential
             **kwargs: Additional arguments for AbstractFeature
         """
-        super().__init__(**kwargs)
+        # Set default range if not provided
+        if range is None:
+            range = [-10.0, 10.0]  # Reasonable range for membrane potentials
+            
+        super().__init__(
+            name=name,
+            value=value,
+            range=range,
+            **kwargs
+        )
 
         self.spike_threshold = spike_threshold
         self.alpha = alpha
@@ -983,26 +998,32 @@ def __init__(
         self.batch_size = None
         self.target_size = None
         self.initialized = False
+        self.connection = None
 
-    def compute(
-        self, 
-        connection: 'MulticompartmentConnection',
-        source_s: torch.Tensor,
-        **kwargs
-    ) -> torch.Tensor:
+    def compute(self, conn_spikes) -> torch.Tensor:
         """
         Compute forward pass with arctangent surrogate gradients.
         
         Args:
-            connection: Parent MulticompartmentConnection
-            source_s: Source layer spikes [batch_size, source_neurons]
-            **kwargs: Additional arguments
+            conn_spikes: Connection spikes tensor [batch_size, source_neurons * target_neurons]
             
         Returns:
-            Target spikes with differentiable surrogate gradients [batch_size, target_neurons]
+            Target spikes with differentiable surrogate gradients [batch_size, source_neurons * target_neurons]
         """
-        # Step 1: Compute synaptic input
-        synaptic_input = torch.mm(source_s.float(), connection.w)
+        # Ensure connection is available
+        if self.connection is None:
+            raise RuntimeError("ArctangentSurrogateFeature not properly initialized. Call prime_feature first.")
+            
+        # Reshape conn_spikes to [batch_size, source_neurons, target_neurons]
+        batch_size = conn_spikes.size(0)
+        source_n = self.connection.source.n
+        target_n = self.connection.target.n
+        
+        # Reshape connection spikes to matrix form
+        conn_spikes_matrix = conn_spikes.view(batch_size, source_n, target_n)
+        
+        # Step 1: Compute synaptic input (sum over source neurons)
+        synaptic_input = conn_spikes_matrix.sum(dim=1)  # [batch_size, target_neurons]
 
         # Step 2: Initialize membrane potential if needed
         if not self.initialized:
@@ -1021,7 +1042,15 @@ def compute(
         # Step 5: Apply reset mechanism
         self._apply_reset(spikes)
 
-        return spikes
+        # Step 6: Broadcast spikes back to connection format
+        # Each target spike affects all connections to that target
+        spikes_broadcast = spikes.unsqueeze(1).expand(batch_size, source_n, target_n)
+        
+        # Apply spikes to incoming connections
+        output_spikes = conn_spikes_matrix * spikes_broadcast
+        
+        # Reshape back to original format
+        return output_spikes.view(batch_size, source_n * target_n)
 
     def arctangent_surrogate_spike(
         self, 
@@ -1048,7 +1077,13 @@ def arctangent_surrogate_spike(
     def _initialize_state(self, reference_tensor: torch.Tensor):
         """Initialize state tensors based on input dimensions."""
         self.batch_size, self.target_size = reference_tensor.shape
-        self.v_membrane = torch.zeros_like(reference_tensor)
+        # Initialize membrane potential to match batch size and target neurons
+        if self.v_membrane is None:
+            self.v_membrane = torch.zeros_like(reference_tensor)
+        else:
+            # Expand existing membrane potential to match batch size
+            if self.v_membrane.size(0) != self.batch_size:
+                self.v_membrane = self.v_membrane.expand(self.batch_size, -1)
         self.initialized = True
 
     def _apply_reset(self, spikes: torch.Tensor):
@@ -1067,10 +1102,39 @@ def _apply_reset(self, spikes: torch.Tensor):
 
     def reset_state_variables(self):
         """Reset all internal state variables."""
+        super().reset_state_variables()
         self.v_membrane = None
         self.batch_size = None
         self.target_size = None
         self.initialized = False
+        
+    def prime_feature(self, connection, device, **kwargs) -> None:
+        """
+        Prime the feature for use in a connection.
+        
+        Args:
+            connection: Parent connection object
+            device: Device to run on
+            **kwargs: Additional arguments
+        """
+        # Store connection reference
+        self.connection = connection
+        
+        # Call parent prime_feature
+        super().prime_feature(connection, device, **kwargs)
+        
+    def initialize_value(self):
+        """
+        Initialize default membrane potential values.
+        
+        Returns:
+            Zero membrane potentials for all target neurons
+        """
+        if self.connection is None:
+            raise RuntimeError("Connection not set. Call prime_feature first.")
+            
+        # Initialize with zeros - membrane potentials start at rest
+        return torch.zeros(1, self.connection.target.n)
 
     def get_membrane_potential(self) -> Optional[torch.Tensor]:
         """Get current membrane potential."""
@@ -1101,6 +1165,7 @@ def __repr__(self) -> str:
         """String representation."""
         return (
             f"ArctangentSurrogateFeature("
+            f"name='{self.name}', "
             f"spike_threshold={self.spike_threshold}, "
             f"alpha={self.alpha}, "
             f"dt={self.dt}, "
@@ -1168,11 +1233,61 @@ def backward(ctx, grad_output):
         # threshold and alpha gradients are None (not optimized)
         return grad_input, None, None
 
+class GoodnessScore(AbstractSubFeature):
+    """
+    SubFeature to compute the goodness score (sum of spikes over time) for all layers in a BindsNET network.
+    """
 
+    def __init__(
+        self,
+        name: str,
+        parent_feature: AbstractFeature = None,
+        network=None,  # <-- Add this argument
+        time: int = 250,
+        input_layer: str = "X",
+    ) -> None:
+        super().__init__(name, parent_feature)
+        self.time = time
+        self.input_layer = input_layer
+        self.network = network  # <-- Store the network
+
+    def compute(self, sample: torch.Tensor) -> dict:
+        # Use self.network if provided, else fall back to parent_feature's connection
+        if self.network is not None:
+            network = self.network
+        else:
+            if not hasattr(self.parent, "connection") or self.parent.connection is None:
+                raise RuntimeError("Parent feature must have a valid connection attribute.")
+            if not hasattr(self.parent.connection, "network") or self.parent.connection.network is None:
+                raise RuntimeError("Connection must have a valid network attribute.")
+            network = self.parent.connection.network
+
+        network.reset_state_variables()
+        inputs = {self.input_layer: sample.unsqueeze(0) if sample.dim() == 1 else sample}
+        spike_record = {layer_name: [] for layer_name in network.layers}
+
+        for t in range(self.time):
+            network.run(inputs, time=1)
+            for layer_name, layer in network.layers.items():
+                spike_record[layer_name].append(layer.s.clone().detach())
+
+        goodness_per_layer = {}
+        for layer_name, spikes_list in spike_record.items():
+            spikes = torch.stack(spikes_list, dim=0)
+            goodness = spikes.sum(dim=0).sum(dim=0)
+            goodness_per_layer[layer_name] = goodness
+
+        total_goodness = sum([v.sum() for v in goodness_per_layer.values()])
+        goodness_per_layer["total_goodness"] = total_goodness
+
+        return goodness_per_layer
+    
+    
 # Helper function for easy creation
 def create_arctangent_surrogate_connection(
     source,
     target,
+    name: str = "arctangent_surrogate",
     w: Optional[torch.Tensor] = None,
     spike_threshold: float = 1.0,
     alpha: float = 2.0,
@@ -1186,6 +1301,7 @@ def create_arctangent_surrogate_connection(
     Args:
         source: Source population
         target: Target population
+        name: Name for the surrogate feature
         w: Weight matrix (initialized randomly if None)
         spike_threshold: Spike threshold
         alpha: Surrogate gradient steepness
@@ -1204,6 +1320,7 @@ def create_arctangent_surrogate_connection(
 
     # Create arctangent surrogate feature
     surrogate_feature = ArctangentSurrogateFeature(
+        name=name,
         spike_threshold=spike_threshold,
         alpha=alpha,
         dt=dt,
@@ -1215,7 +1332,7 @@ def create_arctangent_surrogate_connection(
         source=source,
         target=target,
         w=w,
-        features=[surrogate_feature],
+        pipeline=[surrogate_feature],
         **mcc_kwargs
     )