Create model.py

model.py

@@ -0,0 +1,225 @@
+"""
+Spatial Context Networks (SCN)
+Geometric Semantic Routing in Neural Architectures
+
+Author: Furkan Nar
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class GeometricActivation(nn.Module):
+    """
+    Geometric activation function based on normalized Euclidean distance.
+    
+    Each neuron acts as a point-mass with a learnable centroid in d-dimensional space.
+    Activation is inversely proportional to the normalized distance from the centroid:
+    
+        f(v) = 1 / (||v - mu||_2 / sqrt(d) + epsilon)
+    
+    Args:
+        n_neurons (int): Number of neurons (centroids) in this layer.
+        dim (int): Dimensionality of the input semantic space.
+        stability_factor (float): SF in the paper; epsilon = 1/SF. Default: 10.0
+    """
+
+    def __init__(self, n_neurons: int, dim: int, stability_factor: float = 10.0):
+        super().__init__()
+        self.n_neurons = n_neurons
+        self.dim = dim
+        self.epsilon = 1.0 / stability_factor
+
+        # Learnable centroids: shape (n_neurons, dim)
+        self.centroids = nn.Parameter(torch.randn(n_neurons, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor of shape (batch_size, dim)
+        Returns:
+            activations: Tensor of shape (batch_size, n_neurons)
+        """
+        # x: (B, dim) -> (B, 1, dim)
+        # centroids: (n_neurons, dim) -> (1, n_neurons, dim)
+        diff = x.unsqueeze(1) - self.centroids.unsqueeze(0)  # (B, n_neurons, dim)
+        dist = torch.norm(diff, dim=-1)                       # (B, n_neurons)
+        normalized_dist = dist / math.sqrt(self.dim)
+        activations = 1.0 / (normalized_dist + self.epsilon)
+        return activations
+
+
+class SemanticRoutingLayer(nn.Module):
+    """
+    Semantic routing layer that selectively activates neurons based on
+    geometric affinity to the input.
+
+    Active set: S = { n_i | f_i(q) > tau }
+    Binary mask: M_ij = I[ f_j(v_i) > tau ]
+
+    Args:
+        n_neurons (int): Number of neurons.
+        dim (int): Input dimensionality.
+        routing_threshold (float): Activation threshold tau. Default: 0.5
+        stability_factor (float): Passed to GeometricActivation. Default: 10.0
+    """
+
+    def __init__(
+        self,
+        n_neurons: int,
+        dim: int,
+        routing_threshold: float = 0.5,
+        stability_factor: float = 10.0,
+    ):
+        super().__init__()
+        self.routing_threshold = routing_threshold
+        self.geo_activation = GeometricActivation(n_neurons, dim, stability_factor)
+
+    def forward(self, x: torch.Tensor):
+        """
+        Args:
+            x: Input tensor of shape (batch_size, dim)
+        Returns:
+            activations: Raw activations, shape (batch_size, n_neurons)
+            mask: Binary routing mask, shape (batch_size, n_neurons)
+        """
+        activations = self.geo_activation(x)
+        mask = (activations > self.routing_threshold).float()
+        return activations, mask
+
+
+class ConnectionDensityLayer(nn.Module):
+    """
+    Connection density weighting with adaptive scaling and explosion control.
+
+    C = sum_{i in S} w_i / (alpha / z)
+
+    where alpha = total neurons, z = |S| (active neurons).
+    When C > tau_exp, square-root damping is applied: C_stable = sqrt(C).
+
+    Args:
+        n_neurons (int): Total number of neurons (alpha).
+        explosion_threshold (float): tau_exp. Default: 2.0
+    """
+
+    def __init__(self, n_neurons: int, explosion_threshold: float = 2.0):
+        super().__init__()
+        self.n_neurons = n_neurons
+        self.explosion_threshold = explosion_threshold
+
+        # Learnable per-neuron connection weights
+        self.connection_weights = nn.Parameter(torch.randn(n_neurons))
+
+    def forward(self, activations: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            activations: Shape (batch_size, n_neurons)
+            mask: Binary mask, shape (batch_size, n_neurons)
+        Returns:
+            context: Scalar context score per sample, shape (batch_size, 1)
+        """
+        z = mask.sum(dim=-1, keepdim=True).clamp(min=1.0)  # (B, 1)
+        alpha = float(self.n_neurons)
+
+        # Weighted masked activations
+        weighted = activations * mask * self.connection_weights.unsqueeze(0)  # (B, n)
+        context = weighted.sum(dim=-1, keepdim=True) / (alpha / z)            # (B, 1)
+
+        # Explosion control: sqrt damping
+        context = torch.where(
+            context > self.explosion_threshold,
+            torch.sqrt(context.abs() + 1e-8) * context.sign(),
+            context,
+        )
+        return context
+
+
+class SpatialContextNetwork(nn.Module):
+    """
+    Spatial Context Network (SCN).
+
+    Full architecture:
+        1. SemanticRoutingLayer  — geometric activation + binary routing mask
+        2. ConnectionDensityLayer — adaptive normalization + explosion control
+        3. Linear projection     — map context score to output space
+        4. Pattern distribution  — element-wise multiply by softmax(pattern_weights)
+
+    Args:
+        input_dim (int): Dimensionality of input features.
+        n_neurons (int): Number of hidden geometric neurons. Default: 32
+        output_dim (int): Number of output classes/dimensions. Default: 4
+        routing_threshold (float): Routing threshold tau. Default: 0.5
+        stability_factor (float): Controls epsilon = 1/SF. Default: 10.0
+        explosion_threshold (float): Threshold for sqrt damping. Default: 2.0
+
+    Example::
+
+        model = SpatialContextNetwork(input_dim=10, n_neurons=32, output_dim=4)
+        x = torch.randn(8, 10)
+        output = model(x)  # (8, 4)
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 10,
+        n_neurons: int = 32,
+        output_dim: int = 4,
+        routing_threshold: float = 0.5,
+        stability_factor: float = 10.0,
+        explosion_threshold: float = 2.0,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.n_neurons = n_neurons
+        self.output_dim = output_dim
+
+        self.routing = SemanticRoutingLayer(
+            n_neurons, input_dim, routing_threshold, stability_factor
+        )
+        self.density = ConnectionDensityLayer(n_neurons, explosion_threshold)
+        self.projection = nn.Linear(1, output_dim)
+
+        # Pattern prior weights (learnable)
+        self.pattern_weights = nn.Parameter(torch.zeros(output_dim))
+
+        # Initialise pattern weights to approximate the priors from the paper
+        # [Mathematics=0.38, Language=0.25, Vision=0.22, Reasoning=0.15]
+        with torch.no_grad():
+            prior = torch.tensor([0.38, 0.25, 0.22, 0.15])
+            if output_dim == 4:
+                self.pattern_weights.copy_(torch.log(prior + 1e-8))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor of shape (batch_size, input_dim)
+        Returns:
+            output: Tensor of shape (batch_size, output_dim)
+        """
+        activations, mask = self.routing(x)
+        context = self.density(activations, mask)
+        hidden = self.projection(context)                   # (B, output_dim)
+        output = hidden * F.softmax(self.pattern_weights, dim=-1)
+        return output
+
+    def get_network_stats(self, x: torch.Tensor) -> dict:
+        """
+        Returns diagnostic statistics for a batch of inputs.
+
+        Returns:
+            dict with keys: mean_active_neurons, network_efficiency,
+                            mean_context_score, activations, mask
+        """
+        with torch.no_grad():
+            activations, mask = self.routing(x)
+            context = self.density(activations, mask)
+            active = mask.sum(dim=-1)
+            return {
+                "mean_active_neurons": active.mean().item(),
+                "network_efficiency": (active / self.n_neurons).mean().item(),
+                "mean_context_score": context.mean().item(),
+                "activations": activations,
+                "mask": mask,
+            }