Multihead support for equinox

README.md

@@ -78,15 +78,15 @@ uv sync --extra train --extra flax
 
 ## Equinox Quickstart
 ```python
-from memax.equinox.train_utils import get_residual_memory_model
+from memax.equinox.train_utils import build_named_model
 import jax
 import jax.numpy as jnp
 from equinox import filter_jit, filter_vmap
 from memax.equinox.train_utils import add_batch_dim
 
 T, F = 5, 6 # time and feature dim
 
-model = get_residual_memory_model(
+model = build_named_model(
     model_name="LRU", input=F, hidden=8, output=1, num_layers=2,
     key=jax.random.key(0)
 )

memax/datasets/mnist_listops.py

@@ -4,8 +4,7 @@
 import jax.numpy as jnp
 from datasets import load_dataset  # huggingface datasets
 
-from memax.train_utils import get_residual_memory_models
-
+from memax.equinox.train_utils import build_model
 
 NUM_EPOCHS = 100
 BATCH_SIZE = 32
@@ -158,4 +157,4 @@ def make_dataset(dataset_size=3, num_terms=5, key=jax.random.key(0), batch_size=
 
 
 key = jax.random.key(SEED)
-models = get_residual_memory_models(input=1, hidden=256, output=NUM_LABELS, key=key)
+models = build_model(input=1, hidden=256, output=NUM_LABELS, key=key)

memax/equinox/gras.py

@@ -1,7 +1,6 @@
-from beartype.typing import Callable, Optional, Tuple
-
 import equinox as eqx
 import jax
+from beartype.typing import Callable, Optional, Tuple
 from jaxtyping import PRNGKeyArray, Shaped
 
 from memax.equinox.groups import BinaryAlgebra, Module
@@ -13,9 +12,9 @@ class GRAS(Module):
 
     # Generalized Recurrent Algebraic Structure (GRAS)
 
-    A GRAS contains a **set action** $(H, Z, \bullet)$, an initial state $h_0 \in H$, and two maps/functions 
-    
-    $f$ maps input features and a boolean start flag to the action space $Z$. 
+    A GRAS contains a **set action** $(H, Z, \bullet)$, an initial state $h_0 \in H$, and two maps/functions
+
+    $f$ maps input features and a boolean start flag to the action space $Z$.
 
     $f: X^n \times \\{0, 1\\}^n \mapsto Z^n$
 
@@ -42,7 +41,7 @@ class GRAS(Module):
 
     $ a \bullet (b \bullet c) = (a \bullet b) \bullet c $.
 
-    This enables us to execute $\bullet$ via a parallel scan, which is much more efficient than a sequential scan. 
+    This enables us to execute $\bullet$ via a parallel scan, which is much more efficient than a sequential scan.
     The semigroup GRAS therefore contains the same maps $f$ and $g$ as above, but $\bullet$ is now a semigroup operation.
     Furthermore, in a semigroup, the action and recurrent state spaces are identical, i.e., $Z = H$
 
@@ -54,6 +53,9 @@ class GRAS(Module):
     ```
     """
 
+    readout_dim: int
+    """Feature dimension returned by ``backward_map`` before trunk mixing."""
+
     algebra: BinaryAlgebra
     scan: Callable[
         [
@@ -68,7 +70,7 @@ def forward_map(
         self, x: Input, key: Optional[Shaped[PRNGKeyArray, ""]] = None
     ) -> RecurrentState:
         """Maps inputs to the recurrent space.
-        
+
         `(feature, start) -> H`
         """
         raise NotImplementedError
@@ -79,9 +81,11 @@ def backward_map(
         x: Input,
         key: Optional[Shaped[PRNGKeyArray, ""]] = None,
     ) -> OutputEmbedding:
-        """Maps the recurrent space to the output space.
+        """Maps recurrent state and inputs to readout features of size ``readout_dim``.
 
-        `(h, (feature, start)) -> Y`
+        Trunk models (:class:`~memax.equinox.models.residual.ResidualModel`,
+        :class:`~memax.equinox.models.multihead_residual.MultiHeadResidualModel`)
+        apply :class:`~memax.equinox.models.layer_mixer.LayerMixer` after this.
         """
         raise NotImplementedError
 

memax/equinox/models/layer_mixer.py

@@ -0,0 +1,53 @@
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from beartype.typing import Callable, Tuple
+from equinox import nn
+from jaxtyping import Array, PRNGKeyArray, Shaped
+
+
+class LayerMixer(eqx.Module):
+    """Per-head linears from readout features to trunk width, then norm and activation.
+
+    Applies ``H`` independent ``Linear(in_features, head_dim)`` maps, concatenates to
+    ``out_features``, then LayerNorm and activation.
+    """
+
+    num_heads: int
+    head_dim: int
+    in_features: int
+    out_features: int
+    heads: Tuple[nn.Linear, ...]
+    norm: nn.LayerNorm
+    activation: eqx.Module
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        num_heads: int,
+        activation: Callable[[Array], Array] = jax.nn.leaky_relu,
+        *,
+        key: Shaped[PRNGKeyArray, ""],
+    ):
+        if out_features % num_heads != 0:
+            raise ValueError(
+                f"out_features ({out_features}) must be divisible by "
+                f"num_heads ({num_heads})"
+            )
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_heads = num_heads
+        self.head_dim = out_features // num_heads
+        keys = jax.random.split(key, num_heads)
+        self.heads = tuple(
+            nn.Linear(in_features, self.head_dim, key=head_key) for head_key in keys
+        )
+        self.norm = nn.LayerNorm((out_features,), use_weight=False, use_bias=False)
+        self.activation = nn.Lambda(activation)
+
+    def __call__(self, z: Array) -> Array:
+        ys = jnp.stack([head(z) for head in self.heads])
+        y = ys.reshape(-1)
+        y = self.norm(y)
+        return self.activation(y)

memax/equinox/models/multihead_residual.py

@@ -0,0 +1,89 @@
+import jax
+from beartype.typing import List, Optional, Tuple
+from equinox import filter_vmap, nn
+from jaxtyping import PRNGKeyArray, Shaped
+
+from memax.equinox.groups import Module
+from memax.equinox.models.layer_mixer import LayerMixer
+from memax.mtypes import Input, ResetRecurrentState
+
+
+class MultiHeadResidualModel(Module):
+    """A residual stack where per-layer ``ff`` is replaced by :class:`LayerMixer`."""
+
+    layers: List[Module]
+    mixers: List[LayerMixer]
+    map_in: nn.Linear
+    map_out: nn.Linear
+    recurrent_size: int
+    num_heads: int
+
+    def __init__(
+        self,
+        make_layer_fn,
+        input_size,
+        output_size,
+        recurrent_size,
+        num_heads,
+        num_layers=2,
+        activation=jax.nn.leaky_relu,
+        *,
+        key,
+    ):
+        if recurrent_size % num_heads != 0:
+            raise ValueError(
+                f"recurrent_size ({recurrent_size}) must be divisible by "
+                f"num_heads ({num_heads})"
+            )
+        self.recurrent_size = recurrent_size
+        self.num_heads = num_heads
+        self.layers = []
+        self.mixers = []
+        keys = jax.random.split(key, 3)
+        self.map_in = nn.Linear(input_size, recurrent_size, key=keys[0])
+        self.map_out = nn.Linear(recurrent_size, output_size, key=keys[1])
+        key = keys[2]
+        for _ in range(num_layers):
+            key, layer_key, mixer_key = jax.random.split(key, 3)
+            layer = make_layer_fn(recurrent_size=recurrent_size, key=layer_key)
+            self.layers.append(layer)
+            self.mixers.append(
+                LayerMixer(
+                    layer.readout_dim,
+                    recurrent_size,
+                    num_heads=num_heads,
+                    activation=activation,
+                    key=mixer_key,
+                )
+            )
+
+    def __call__(
+        self, h: ResetRecurrentState, x: Input, key: Optional[PRNGKeyArray] = None
+    ) -> Tuple[ResetRecurrentState, ...]:
+        emb, start = x
+        emb = filter_vmap(self.map_in)(emb)
+        layer_in = (emb, start)
+        h_out = []
+        for mixer, recurrent_layer, h_i in zip(self.mixers, self.layers, h):
+            if key is None:
+                key, rkey = None, None
+            else:
+                key, rkey = jax.random.split(key)
+            tmp, feat = recurrent_layer(h_i, layer_in, key=rkey)
+            h_out.append(tmp)
+            z = filter_vmap(mixer)(feat)
+            layer_in = (z, start)
+        out = filter_vmap(self.map_out)(layer_in[0])
+        return tuple(h_out), out
+
+    def initialize_carry(
+        self, key: Optional[Shaped[PRNGKeyArray, ""]] = None
+    ) -> Tuple[ResetRecurrentState, ...]:
+        if key is None:
+            keys = tuple(None for _ in range(len(self.layers)))
+        else:
+            keys = jax.random.split(key, len(self.layers))
+        return tuple(l.initialize_carry(k) for l, k in zip(self.layers, keys))
+
+    def latest_recurrent_state(self, h: ResetRecurrentState) -> ResetRecurrentState:
+        return tuple(l.latest_recurrent_state(h_i) for l, h_i in zip(self.layers, h))

memax/equinox/models/residual.py

@@ -1,10 +1,10 @@
-from beartype.typing import List, Optional, Tuple
-
 import jax
+from beartype.typing import List, Optional, Tuple
 from equinox import filter_vmap, nn
 from jaxtyping import PRNGKeyArray, Shaped
 
 from memax.equinox.groups import Module
+from memax.equinox.models.layer_mixer import LayerMixer
 from memax.mtypes import Input, ResetRecurrentState
 
 
@@ -15,9 +15,10 @@ class ResidualModel(Module):
     There is a nonlinearity between network layers."""
 
     layers: List[Module]
-    ff: List[nn.Sequential]
+    mixers: List[LayerMixer]
     map_in: nn.Linear
     map_out: nn.Linear
+    recurrent_size: int
 
     def __init__(
         self,
@@ -28,26 +29,26 @@ def __init__(
         num_layers=2,
         activation=jax.nn.leaky_relu,
         *,
-        key
+        key,
     ):
+        self.recurrent_size = recurrent_size
         self.layers = []
-        self.ff = []
+        self.mixers = []
         keys = jax.random.split(key, 3)
         self.map_in = nn.Linear(input_size, recurrent_size, key=keys[0])
         self.map_out = nn.Linear(recurrent_size, output_size, key=keys[1])
         key = keys[2]
         for _ in range(num_layers):
-            key, ff_key = jax.random.split(key)
-            self.layers.append(make_layer_fn(recurrent_size=recurrent_size, key=key))
-            self.ff.append(
-                nn.Sequential(
-                    [
-                        nn.Linear(recurrent_size, recurrent_size, key=ff_key),
-                        nn.LayerNorm(
-                            (recurrent_size,), use_weight=False, use_bias=False
-                        ),
-                        nn.Lambda(activation),
-                    ]
+            key, layer_key, mixer_key = jax.random.split(key, 3)
+            layer = make_layer_fn(recurrent_size=recurrent_size, key=layer_key)
+            self.layers.append(layer)
+            self.mixers.append(
+                LayerMixer(
+                    layer.readout_dim,
+                    recurrent_size,
+                    num_heads=1,
+                    activation=activation,
+                    key=mixer_key,
                 )
             )
 
@@ -58,14 +59,14 @@ def __call__(
         emb = filter_vmap(self.map_in)(emb)
         layer_in = (emb, start)
         h_out = []
-        for ff, recurrent_layer, h_i in zip(self.ff, self.layers, h):
+        for mixer, recurrent_layer, h_i in zip(self.mixers, self.layers, h):
             if key is None:
                 key, rkey = None, None
             else:
                 key, rkey = jax.random.split(key)
-            tmp, z = recurrent_layer(h_i, layer_in, key=rkey)
+            tmp, feat = recurrent_layer(h_i, layer_in, key=rkey)
             h_out.append(tmp)
-            z = filter_vmap(ff)(z)
+            z = filter_vmap(mixer)(feat)
             layer_in = (z, start)
         out = filter_vmap(self.map_out)(layer_in[0])
         return tuple(h_out), out

memax/equinox/semigroups/__init__.py

@@ -2,7 +2,7 @@
 They are generally much faster than standard RNNs.
 
 Each RNN type gets its own file.
-+ `memax.equinox.semigroups.attn` provides dot-product attention layer. 
++ `memax.equinox.semigroups.attn` provides dot-product attention layer.
 + `memax.equinox.semigroups.delta` provides the DeltaNet layer.
 + `memax.equinox.semigroups.deltap` provides the DeltaProduct layer.
 + `memax.equinox.semigroups.stack` provides framestacking (sliding-window) as an RNN.
@@ -11,8 +11,8 @@
 + `memax.equinox.semigroups.fwp` provides the Fast Weight Programmer layer.
 + `memax.equinox.semigroups.lru` provides the Linear Recurrent Unit layer.
 + `memax.equinox.semigroups.gdn` provides the Gated DeltaNet layer.
-+ `memax.equinox.semigroups.lrnn` provides a basic linear recurrence. 
-+ `memax.equinox.semigroups.mlp` provides an MLP (no memory) for completeness.
++ `memax.equinox.semigroups.lrnn` provides a basic linear recurrence.
++ `memax.equinox.semigroups.identity` provides a memory-free baseline cell (no cross-time state).
 + `memax.equinox.semigroups.s6` provides the Selective State Space Model (Mamba) layer.
 + `memax.equinox.semigroups.spherical` provides Rotational RNN layer (spherical projection).
-"""
+"""