python/elidedb/fdnnstudent.py

"""FDNN students shaped like their teachers.

A student is not a generic head bolted onto whatever vectors happen to
be lying around. That was tried and it failed for a reason worth
keeping: an attention-pool + ReLU MLP over FDNN-V frame vectors reached
0.9168 cosine to the PE teacher while a CONSTANT prediction of the
corpus mean scored 0.8629, and nearest-neighbour agreement was 0.005.
The head had learned the mean. A least-squares oracle on the same input
reached nn_top1 0.0155, so no head over those vectors could have worked
- the information was not in them.

So the student takes the TEACHER'S OWN ARCHITECTURE as its skeleton,
reads pixels like the teacher does, and shrinks. What changes is the
neuron, per the three FDNN rules:

  RULE 1  each neuron is a sub-network. Every transformer block's MLP -
          the 2/3 of a ViT's parameters that is elementwise GELU -
          becomes a KAN-style sum over k heterogeneous bases of the
          same pre-activation: FINER (variable-period oscillator),
          Gabor (wavelet, fires on a burst), sine, poly-phase (chirp,
          fires on acceleration). omega bands partition the spectrum so
          different channels answer to different rates of change.
          A GELU unit can only answer "how much"; these answer "how
          much, how fast, and is it accelerating" - which is what a
          video model's MLP is being asked for in the first place.

  RULE 2  neurogenesis and apoptosis after training, each followed by a
          re-settle fine-tune. `mask` is the aliveness vector the cycle
          writes; it is frozen against the optimizer, because AdamW
          weight-decays an unfrozen mask off 1.0 and every channel
          quietly shrinks (measured on the context tower).

  RULE 3  PPO + reverse attention decide who lives: utilization is the
          drop in held-out fidelity when a channel is silenced.

WHY THE SKELETON MATTERS. V-JEPA2 is 24 blocks of width 1024 over 8,192
spatiotemporal tokens (256px, 64 frames, patch 16, tubelet 2). Most of
that cost is token count, not depth. A student that keeps the tubelet
embedding and the attentive pooler - the parts that decide WHAT is
compared - while cutting resolution, frames, width and depth, is doing
the teacher's computation at the teacher's shape. A student with a
different shape is a different model that happens to be trained on the
teacher's outputs, which is what failed.
"""
from __future__ import annotations

import numpy as np

import mlx.core as mx
import mlx.nn as nn

# Basis identifiers, matching FDNNTemporalCell so the two implementations
# cannot drift apart: 0 FINER, 1 Gabor, 2 sine, 3 poly-phase.
FINER, GABOR, SINE, POLY = 0, 1, 2, 3


class FDNNFeedForward(nn.Module):
    """Rule 1, applied to a transformer block's MLP.

    Standard: y = W2 @ gelu(W1 @ x). Here each of C output channels is a
    KAN sum over k sub-functions of the SAME pre-activation, so one
    channel can be periodic, another can fire only on a burst, another
    only on acceleration - heterogeneous by construction rather than by
    hoping a homogeneous nonlinearity specialises.

    omega bands are spatial-temporal here, not purely temporal: tokens
    are tubelets, so a band is a rate of change across the token grid.
    """

    def __init__(self, dim, channels, k_width=4,
                 omega_bands=(0.8, 2.5, 8.0),
                 band_fractions=(0.34, 0.33, 0.33), seed=0):
        super().__init__()
        rng = np.random.default_rng(seed)
        C, k = channels, k_width
        self.C, self.k = C, k

        om = []
        for b, fr in zip(omega_bands, band_fractions):
            om.extend([b] * int(round(fr * C)))
        om = (om + [omega_bands[-1]] * C)[:C]
        self.omegas_per_neuron = np.array(om, np.float32)
        self.omegas = mx.array(np.repeat(self.omegas_per_neuron, k))

        half, quarter = max(k // 2, 1), max(k // 4, 1)
        per = np.array([FINER] * half + [GABOR] * quarter
                       + [POLY] * max(k - half - quarter, 0), np.int32)[:k]
        self.basis_types = mx.array(np.tile(per, C).astype(np.int32))
        # categorical, not a weight: unfrozen the optimizer drifts it off
        # its exact values and silently reroutes every Gabor to poly
        self.freeze(keys=["basis_types"], recurse=False)

        mean_om = float(self.omegas_per_neuron.mean())
        lim = float(np.sqrt(6.0 / dim) / mean_om)
        self.W1 = mx.array(rng.uniform(-lim, lim, (dim, C * k)).astype(np.float32))
        self.b1 = mx.array(rng.uniform(-2.0, 2.0, (C * k,)).astype(np.float32))
        self.phases = mx.array(rng.uniform(0, 2 * np.pi, (C * k,)).astype(np.float32))
        self.gabor_s = mx.array(rng.uniform(0.3, 1.5, (C * k,)).astype(np.float32))
        log_om = np.log(np.clip(np.repeat(self.omegas_per_neuron, k), 1e-3, None))
        self.log_alpha = mx.array(rng.uniform(np.minimum(0.0, log_om),
                                              np.maximum(0.0, log_om) + 1e-6
                                              ).astype(np.float32))
        w2s = float(np.sqrt(6.0 / (C * k)))
        self.w2 = mx.array(rng.uniform(-w2s, w2s, (C, k)).astype(np.float32))
        pl = float(np.sqrt(6.0 / C))
        self.Wp = mx.array(rng.uniform(-pl, pl, (C, dim)).astype(np.float32))

        self.mask = mx.array(np.ones((C,), np.float32))     # rule 2
        self.freeze(keys=["mask"], recurse=False)

    def channels_out(self, x):
        """Per-channel activation before the mask - rule 3's signal."""
        pre = x @ self.W1 + self.b1
        om_h = self.omegas * pre
        sq = pre * pre
        alpha = mx.exp(self.log_alpha)
        finer = mx.sin(self.omegas * (mx.abs(pre) + 1.0) * pre + self.phases)
        gab = mx.exp(-(self.gabor_s ** 2) * sq) * mx.sin(om_h + self.phases)
        sine = mx.sin(om_h + self.phases)
        poly = mx.sin(alpha * sq + om_h + self.phases)
        acts = mx.where(self.basis_types == FINER, finer,
                        mx.where(self.basis_types == GABOR, gab,
                                 mx.where(self.basis_types == SINE, sine, poly)))
        acts = acts.reshape(*pre.shape[:-1], self.C, self.k)
        return mx.sum(acts * self.w2, axis=-1)

    def __call__(self, x):
        return (self.channels_out(x) * self.mask) @ self.Wp

    def set_active_mask(self, m):
        self.mask = mx.array(np.asarray(m, np.float32))


class Block(nn.Module):
    """A ViT block with the teacher's shape and an FDNN neuron inside."""

    def __init__(self, dim, heads, channels, k_width=4, seed=0):
        super().__init__()
        self.n1 = nn.LayerNorm(dim)
        self.attn = nn.MultiHeadAttention(dim, heads)
        self.n2 = nn.LayerNorm(dim)
        self.ff = FDNNFeedForward(dim, channels, k_width, seed=seed)

    def __call__(self, x):
        h = self.n1(x)
        x = x + self.attn(h, h, h)
        return x + self.ff(self.n2(x))


class TubeletEmbed(nn.Module):
    """The teacher's input stage: non-overlapping spatiotemporal patches.

    Kept because it is what decides what a token IS. Resolution, frame
    count, width and depth are all shrunk; the tokenisation is not
    changed, so the student is comparing the same kind of thing.
    """

    def __init__(self, dim, patch=16, tubelet=2, in_ch=3):
        super().__init__()
        self.patch, self.tubelet = patch, tubelet
        self.proj = nn.Linear(in_ch * patch * patch * tubelet, dim)

    def __call__(self, v):                      # (B, T, H, W, 3) in [-1,1]
        B, T, H, W, C = v.shape
        p, t = self.patch, self.tubelet
        gt, gh, gw = T // t, H // p, W // p
        v = v[:, :gt * t, :gh * p, :gw * p]
        v = v.reshape(B, gt, t, gh, p, gw, p, C)
        v = v.transpose(0, 1, 3, 5, 2, 4, 6, 7).reshape(B, gt * gh * gw, -1)
        return self.proj(v), (gt, gh, gw)


class AttentivePool(nn.Module):
    """The teacher pools with a learned query, so the student does too.
    Mean pooling buries the few tokens that carry the event under the
    many that carry the unchanged room."""

    def __init__(self, dim, out_dim):
        super().__init__()
        self.q = mx.array((np.random.default_rng(0).normal(size=(dim,))
                           / np.sqrt(dim)).astype(np.float32))
        self.norm = nn.LayerNorm(dim)
        self.proj = nn.Linear(dim, out_dim)

    def __call__(self, x):
        a = mx.softmax(x @ self.q, axis=1)[..., None]
        y = self.proj(self.norm(mx.sum(a * x, axis=1)))
        return y / (mx.linalg.norm(y, axis=-1, keepdims=True) + 1e-8)


class VideoStudent(nn.Module):
    """V-JEPA2's skeleton at 1/N the size, with FDNN neurons.

    teacher  256px, 64 frames, patch 16, tubelet 2 -> 8192 tokens,
             24 blocks, width 1024, 16 heads
    student  configurable; the default cuts tokens 16x and width 5x,
             which is where a ViT's cost actually lives.
    """

    def __init__(self, out_dim=1024, dim=192, depth=4, heads=3,
                 channels=256, k_width=4, patch=16, tubelet=2,
                 size=128, frames=16, seed=0):
        super().__init__()
        self.size, self.frames = size, frames
        self.embed = TubeletEmbed(dim, patch, tubelet)
        n_tok = (frames // tubelet) * (size // patch) ** 2
        self.pos = mx.array((np.random.default_rng(seed).normal(
            size=(1, n_tok, dim)) * 0.02).astype(np.float32))
        self.blocks = [Block(dim, heads, channels, k_width, seed=seed + i)
                       for i in range(depth)]
        self.pool = AttentivePool(dim, out_dim)

    def __call__(self, v):
        x, _ = self.embed(v)
        x = x + self.pos[:, :x.shape[1]]
        for b in self.blocks:
            x = b(x)
        return self.pool(x)

    def utilization(self, v):
        """Rule 3's raw signal: per-channel mean |activation| per block."""
        x, _ = self.embed(v)
        x = x + self.pos[:, :x.shape[1]]
        out = []
        for b in self.blocks:
            h = b.n1(x)
            x = x + b.attn(h, h, h)
            c = b.ff.channels_out(b.n2(x))
            out.append(np.abs(np.array(c)).mean((0, 1)))
            x = x + b.ff(b.n2(x))
        return out