feat(path): init

Author: maedoc

vbjax/app/path/README.md

@@ -0,0 +1,77 @@
+# **JAX-based DeltaNet Model for Pathfinder Task**
+
+## **1\. Project Overview**
+
+This project provides a JAX-based implementation of a recurrent neural network inspired by the DeltaNet architecture. The model is specifically designed to solve the Long Range Arena (LRA) Pathfinder task. This task, which requires determining if two points on a 64x64 grid are connected by a continuous path, serves as an ideal testbed. It directly challenges a model's ability to capture long-range spatial dependencies and route information selectively, capabilities that are central to the implemented architecture.  
+The model's design is guided by a theoretical framework that views cognition not as a single, passive feed-forward pass, but as an active, iterative process of belief updating. This perspective is more aligned with biological cognition. The core of the model is a recurrent update rule that mathematically embodies this process of Bayesian belief propagation through cycles of prediction and error correction.
+
+## **2\. Theoretical Foundations**
+
+The model synthesizes several key concepts from computational neuroscience and machine learning to create a unified system.
+
+* **Active Inference & Belief Updating**: Unlike standard Vision Transformers which process an entire image in one go, this model processes information sequentially, token-by-token. Each new token (an image patch) is used to update an internal belief state. This iterative refinement is computationally analogous to how an organism might scan a scene, gathering evidence over time. This approach has implications for biological plausibility and offers potential efficiency gains in complex tasks where only a subset of information is relevant at any given moment.  
+* **The DeltaNet Update Rule**: The heart of the recurrence is the update to a state matrix S at each step t. This matrix represents the accumulated knowledge or beliefs of the model. The update rule has two equivalent mathematical forms:  
+  1. Multiplicative Form: S\_t \= S\_{t-1} @ (I \- B \* k\_t @ k\_t.T) \+ v\_t @ k\_t.T  
+     This form highlights how the prior belief S\_{t-1} is first adjusted by "explaining away" the predictable parts of the input, and then new information is added.  
+  2. Prediction Error Form: S\_t \= S\_{t-1} \+ (v\_t \- B \* S\_{t-1} @ k\_t) @ k\_t.T  
+     This is the form directly implemented in the code as it makes the learning signal explicit. It elegantly decomposes the update into intuitive computational steps:  
+     * v\_pred \= B \* S\_{t-1} @ k\_t: The model uses its prior belief (S\_{t-1}) to predict the current value vector (v\_t) based on the current context or key vector (k\_t).  
+     * error \= v\_t \- v\_pred: The model calculates the prediction error—the "surprise"—which quantifies how much the actual input deviates from the prediction. This error signal is the primary driver of learning and belief change.  
+     * The outer product of the error and the k\_t vector creates the final update matrix dS, ensuring that the change to the belief state S is precisely targeted based on the context that produced the error.  
+     * B is a learnable, per-head scalar parameter that modulates the strength of the prediction, allowing each head to learn how confidently to trust its own prior beliefs.  
+* **Criticality & Precision Amplification**: The framework hypothesizes that effective cognitive systems operate near a critical state, where information can propagate without dying out or exploding. This is characterized by "avalanches" of neural activity that follow scale-free power laws. To model this, the system incorporates a precision amplification mechanism tied to prediction error:  
+  * The magnitude of the prediction error ||error|| is calculated. A large error signifies a highly surprising event.  
+  * This magnitude is used to compute a temperature T \= exp(-alpha \* ||error||). A larger error leads to a smaller temperature.  
+  * This temperature dynamically sharpens the softmax output for the efferent (outgoing) signal: softmax((S\_t @ q\_t) / T).  
+  * The consequence is that more "surprising" events (large errors) trigger higher-precision (lower temperature) outputs. This amplifies the signal, ensuring that significant, belief-violating information is propagated forcefully through the network, leading to rapid adaptation.  
+* **Normalization**: Input tokens are z-scored before being projected into Q, K, and V vectors. This serves as a simple homeostatic mechanism to stabilize the network's internal dynamics. By ensuring each token has a mean of zero and unit variance, it prevents the recurrent updates from causing the signals to either explode or vanish over long sequences, keeping the network in a responsive and trainable regime.
+
+## **3\. Model Architecture & Implementation**
+
+### **3.1 Data Pipeline**
+
+* **Input**:  
+  * Images (ims.npy): A NumPy array of shape (N, 64, 64).  
+  * Labels (labels.npy): A NumPy array of binary labels (N,).  
+* **Preprocessing**:  
+  * Each 64x64 image is divided into a sequence of non-overlapping 4x4 patches. This tokenization transforms the spatial problem into a sequential one.  
+  * Each patch is flattened into a 16-dimensional vector (token), resulting in a sequence length of (64/4)^2 \= 256 tokens.  
+  * To retain crucial spatial information lost during flattening, two types of embeddings are added to each token: positional embeddings for the patch's original (x, y) grid location and sequential embeddings for the token's index in the 1D sequence.
+
+### **3.2 Core Logic (delta\_net\_layer\_forward)**
+
+The model consists of several stacked DeltaNet layers. Within each layer, a jax.lax.scan function iterates through the sequence, which is the functional equivalent of a for loop in JAX and the engine for the temporal recurrence. For each token:
+
+1. The input token is **z-scored** for normalization.  
+2. The normalized token is projected into per-head Q, K, and V vectors.  
+3. Q and K vectors are passed through a sigmoid nonlinearity to constrain their values.  
+4. The per-head state matrix S is updated using the computationally explicit prediction error formulation of the DeltaNet rule.  
+5. A precision-amplified softmax is applied to generate an efferent signal, with the temperature determined by the prediction error.  
+6. The signals from all heads are concatenated, projected back to the embedding dimension, and added to the original input via a **skip connection**. This is critical for training deep networks as it provides a direct path for gradients to flow, mitigating the vanishing gradient problem.  
+7. A final layer normalization is applied to stabilize the output of the entire layer before it is passed to the next.
+
+### **3.3 JAX Implementation Details**
+
+* **Framework**: The model is built entirely in JAX, using optax for the Adam optimizer and its rich ecosystem of optimizers and schedulers.  
+* **Dataclasses**:  
+  * TrainingConfig: Stores all global hyperparameters like learning rate, batch size, and model dimensions (embed\_dim, num\_heads, num\_layers).  
+  * ModelParams & DeltaLayerParams: Flax-style dataclasses hold all trainable parameters. This is essential as it organizes the parameters into a nested structure that JAX recognizes as a pytree, allowing functions like grad and jit to work seamlessly on the entire model. The trainable B parameter, for instance, lives inside DeltaLayerParams with a shape of (num\_heads,).  
+* **Vectorization**: vmap is used extensively to parallelize computations. It allows the same function to be efficiently run over the batch dimension and across the different attention heads within each layer, leading to clean and highly performant code.  
+* **Compilation**: jit is used on the train\_step function. This compiles the entire device-side computation graph for the loss calculation, backpropagation, and optimizer update, resulting in significant speedups on accelerators like GPUs.
+
+### **3.4 Code Structure**
+
+* main.py: A self-contained script with all the code.  
+  * **Dataclasses**: ModelParams, DeltaLayerParams, TrainingConfig.  
+  * init\_model(): Initializes all model parameters with random values according to the specified configuration.  
+  * delta\_net\_layer\_forward(): Contains the core recurrent logic for a single DeltaNet layer.  
+  * forward\_pass(): Defines the full pass through all layers, from input embedding to final logit.  
+  * loss\_fn(): Calculates the sigmoid binary cross-entropy loss from the logits and labels.  
+  * train\_step(): A JIT-compiled function that executes one full training step.  
+  * train(): The main training loop that handles data loading, batching, and orchestrates the training process.
+
+## **4\. Next Steps & Future Directions**
+
+* **Analysis**: Implement logging within the train\_step to record the average magnitude of the prediction error ||error|| for each batch. This data can be used to test the criticality hypothesis by plotting the distribution of error magnitudes over time and checking for power-law characteristics (e.g., via a log-log plot).  
+* **Optimization**: Experiment with more advanced optax learning rate schedules, such as warmup\_cosine\_decay\_schedule, which can lead to more stable training and better final performance.  
+* **Architecture**: Explore alternative architectural choices. For example, instead of averaging all final token outputs for the readout, one could use only the output of the very last token in the sequence. Different nonlinearities for the Q and K vectors (e.g., tanh) could also be investigated.

vbjax/app/path/__init__.py

@@ -0,0 +1,3 @@
+"""
+
+"""

vbjax/app/path/model.py

@@ -0,0 +1,262 @@
+from functools import partial
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+from jax import random, grad, jit, vmap
+import jax_dataclasses as jdc
+
+# Set random seed for reproducibility
+key = random.PRNGKey(0)
+
+# setup model structure using jax dataclass
+@jdc.pytree_dataclass
+class Model:
+    """
+    DeltaNet-like model for solving pathfinder task
+    """
+
+
+
+# --- 1. Define Data Structures ---
+
+@dataclass
+class ModelParams:
+    # Input Embeddings
+    pos_embeds: jnp.ndarray      # (16, 16, embed_dim)
+    seq_embeds: jnp.ndarray      # (seq_len, embed_dim)
+    
+    # DeltaNet Layers
+    delta_layers: list           # List of layer-specific parameters
+    
+    # Output Head
+    output_w: jnp.ndarray        # (embed_dim, 1)
+    output_b: jnp.ndarray
+
+@dataclass
+class DeltaLayerParams:
+    wq: jnp.ndarray              # (num_heads, embed_dim, head_dim)
+    wk: jnp.ndarray              # (num_heads, embed_dim, head_dim)
+    wv: jnp.ndarray              # (num_heads, embed_dim, head_dim)
+    wo: jnp.ndarray              # (num_heads * head_dim, embed_dim)
+    norm: jnp.ndarray
+    B: jnp.ndarray               # (num_heads,) - NEW: Learnable per-head parameter
+
+@dataclass
+class TrainingConfig:
+    learning_rate: float
+    batch_size: int
+    num_iterations: int
+    # Model dimensions
+    embed_dim: int
+    num_heads: int
+    num_layers: int
+    patch_size: int
+    image_size: int
+    # Precision amplification factor
+    alpha: float
+
+# --- 2. Model Initialization ---
+
+def init_model(key, config):
+    head_dim = config.embed_dim // config.num_heads
+    seq_len = (config.image_size // config.patch_size)**2
+    
+    keys = random.split(key, 2 + config.num_layers * 6) # Adjusted for new B param
+    
+    delta_layers = []
+    for i in range(config.num_layers):
+        layer_params = DeltaLayerParams(
+            wq=random.normal(keys[i*6+0], (config.num_heads, config.embed_dim, head_dim)),
+            wk=random.normal(keys[i*6+1], (config.num_heads, config.embed_dim, head_dim)),
+            wv=random.normal(keys[i*6+2], (config.num_heads, config.embed_dim, head_dim)),
+            wo=random.normal(keys[i*6+3], (config.num_heads * head_dim, config.embed_dim)),
+            norm=jnp.ones(config.embed_dim),
+            B=jnp.ones(config.num_heads) # NEW: Initialize B for each head
+        )
+        delta_layers.append(layer_params)
+        
+    return ModelParams(
+        pos_embeds=random.normal(keys[-4], (config.image_size // config.patch_size, config.image_size // config.patch_size, config.embed_dim)),
+        seq_embeds=random.normal(keys[-3], (seq_len, config.embed_dim)),
+        delta_layers=delta_layers,
+        output_w=random.normal(keys[-2], (config.embed_dim, 1)),
+        output_b=random.normal(keys[-1], (1,))
+    )
+
+# --- 3. Core Model Logic ---
+
+def z_score(x, axis=-1, eps=1e-5):
+    mean = jnp.mean(x, axis=axis, keepdims=True)
+    std = jnp.std(x, axis=axis, keepdims=True)
+    return (x - mean) / (std + eps)
+
+def delta_net_layer_forward(params: DeltaLayerParams, x: jnp.ndarray, S: jnp.ndarray, alpha: float):
+    
+    def scan_fn(carry, token_x):
+        S_prev = carry
+        
+        # Z-score input token
+        token_x_norm = z_score(token_x)
+        
+        # Project to Q, K, V for each head
+        q_h = jnp.einsum('d,hdm->hm', token_x_norm, params.wq)
+        k_h = jnp.einsum('d,hdm->hm', token_x_norm, params.wk)
+        v_h = jnp.einsum('d,hdm->hm', token_x_norm, params.wv)
+        
+        # Sigmoid nonlinearity
+        q_h, k_h = jax.nn.sigmoid(q_h), jax.nn.sigmoid(k_h)
+        
+        # --- Per-head recurrent update (vmapped) ---
+        def head_update(S_h, B_h, q_token, k_token, v_token):
+            # NEW: Prediction includes learnable parameter B
+            v_pred = B_h * (S_h @ k_token)
+            error = v_token - v_pred
+            
+            # State update
+            dS = jnp.outer(error, k_token)
+            S_curr = S_h + dS
+            
+            # --- Mathematical Equivalence Comment ---
+            # The update S_curr = S_prev + outer(v - B * (S_prev @ k), k)
+            # expands to: S_curr = S_prev + v @ k.T - B * S_prev @ k @ k.T
+            # which rearranges to: S_curr = S_prev @ (I - B * k @ k.T) + v @ k.T
+            # This is the DeltaNet update rule.
+            # -----------------------------------------
+            
+            # Precision amplification
+            temp = jnp.exp(-alpha * jnp.linalg.norm(error)) + 1e-6
+            
+            # Efferent signal
+            efferent = jax.nn.softmax((S_curr @ q_token) / temp)
+            return S_curr, efferent
+
+        S_next, efferent_h = vmap(head_update)(S_prev, params.B, q_h, k_h, v_h)
+        
+        # Concatenate heads and project out
+        efferent_flat = efferent_h.flatten()
+        output_token = efferent_flat @ params.wo
+        
+        # Skip connection
+        final_token = token_x + output_token
+        
+        return S_next, final_token
+
+    final_S, output_sequence = jax.lax.scan(scan_fn, S, x)
+    return output_sequence, final_S
+
+def forward_pass(params: ModelParams, batch_patches: jnp.ndarray, config: TrainingConfig):
+    # Add positional and sequential embeddings
+    batch_size, seq_len, _ = batch_patches.shape
+    pos_embeds_flat = params.pos_embeds.reshape(seq_len, config.embed_dim)
+    x = batch_patches + pos_embeds_flat[jnp.newaxis, :, :] + params.seq_embeds[jnp.newaxis, :, :]
+
+    # Initialize recurrent states (one per layer)
+    head_dim = config.embed_dim // config.num_heads
+    initial_S = jnp.zeros((config.num_layers, config.num_heads, head_dim, head_dim))
+    
+    # Propagate through layers
+    current_x = x
+    for i in range(config.num_layers):
+        layer_params = params.delta_layers[i]
+        S_layer = initial_S[i]
+        
+        # Vmap across the batch dimension
+        batch_layer_fwd = vmap(delta_net_layer_forward, in_axes=(None, 0, None, None))
+        current_x, _ = batch_layer_fwd(layer_params, current_x, S_layer, config.alpha)
+        current_x = jax.nn.layer_norm(current_x, use_scale=False) * layer_params.norm
+        
+    # Readout from the mean of the final sequence
+    final_representation = jnp.mean(current_x, axis=1) # (batch, embed_dim)
+    logits = final_representation @ params.output_w + params.output_b
+    return jnp.squeeze(logits)
+
+# --- 4. Loss and Training Functions ---
+
+def loss_fn(params: ModelParams, batch_patches, batch_labels, config: TrainingConfig):
+    logits = forward_pass(params, batch_patches, config)
+    loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, batch_labels))
+    return loss
+
+@partial(jit, static_argnums=(3,))
+def train_step(params: ModelParams, opt_state, batch, config: TrainingConfig):
+    batch_patches, batch_labels = batch
+    loss_val, grads = grad(loss_fn, has_aux=False)(params, batch_patches, batch_labels, config)
+    updates, opt_state = optimizer.update(grads, opt_state, params)
+    params = optax.apply_updates(params, updates)
+    return params, opt_state, loss_val
+
+def train(config: TrainingConfig, model_params: ModelParams, train_images, train_labels):
+    print("--- Starting Training ---")
+    opt_state = optimizer.init(model_params)
+    key = random.PRNGKey(42)
+
+    # Data preprocessing function
+    def preprocess_images(images):
+        p = config.patch_size
+        b, h, w = images.shape
+        # (B, H/p, W/p, p, p)
+        patches = images.reshape(b, h // p, p, w // p, p).transpose(0, 1, 3, 2, 4)
+        # (B, H/p * W/p, p*p)
+        patches = patches.reshape(b, (h // p) * (w // p), p * p)
+        return patches / 255.0
+
+    for i in range(config.num_iterations):
+        # Create a mini-batch
+        key, subkey = random.split(key)
+        indices = random.permutation(subkey, train_images.shape[0])[:config.batch_size]
+        batch_images = train_images[indices]
+        batch_labels = train_labels[indices]
+        
+        batch_patches = preprocess_images(batch_images)
+        
+        model_params, opt_state, loss_val = train_step(model_params, opt_state, (batch_patches, batch_labels), config)
+        
+        if i % 10 == 0:
+            print(f"Iteration {i:03d} | Loss: {loss_val:.4f}")
+            
+    print("--- Training Finished ---")
+    return model_params
+
+
+# --- 5. Main Execution ---
+if __name__ == '__main__':
+    # Create dummy data files for demonstration
+    print("Creating dummy data files 'ims.npy' and 'labels.npy'...")
+    dummy_ims = np.random.randint(0, 256, size=(100, 64, 64), dtype=np.uint8)
+    dummy_labels = np.random.randint(0, 2, size=(100,), dtype=np.int32)
+    np.save('ims.npy', dummy_ims)
+    np.save('labels.npy', dummy_labels)
+    print("Dummy data created.")
+
+    # Configuration
+    config = TrainingConfig(
+        learning_rate=1e-4,
+        batch_size=8,
+        num_iterations=101,
+        embed_dim=64,
+        num_heads=4,
+        num_layers=3,
+        patch_size=4,
+        image_size=64,
+        alpha=1.0,
+    )
+
+    # Initialize optimizer
+    optimizer = optax.adam(config.learning_rate)
+
+    # Initialize model
+    key = random.PRNGKey(0)
+    model_params = init_model(key, config)
+
+    # Load data
+    try:
+        images_data = np.load('ims.npy')
+        labels_data = np.load('labels.npy')
+        print(f"Loaded data: images shape {images_data.shape}, labels shape {labels_data.shape}")
+    except FileNotFoundError:
+        print("Error: Could not find 'ims.npy' or 'labels.npy'.")
+        exit()
+
+    # Run training
+    trained_params = train(config, model_params, jnp.array(images_data, dtype=jnp.float32), jnp.array(labels_data))