python/elidedb/ctxtower.py

"""The context tower — a custom SigLIP video tower over frozen frame features.

WHAT IT IS
----------
SigLIP's image tower is a frozen spatial encoder. This module supplies the
*other* half of an R(2+1)D-style factorisation: a learned TEMPORAL tower that
consumes the sequence of per-frame SigLIP vectors in a window and emits one
vector aligned with SigLIP TEXT embeddings.

    frame vectors (T, 1152)          [frozen SigLIP image tower — the "2D"]
        |  PCA  (frozen, no params)
    (T, d)
        |  dilated temporal conv, taps learned      ← the CNN
        |  each output neuron = KAN sum over k
        |  heterogeneous bases (FINER/Gabor/poly)   ← FDNN thesis 1
    (T, n)
        |  bi-GRU                                    ← the RNN
    (T, 2h)
        |  attention pooling over time
    (2h,)
        |  head -> PCA^T of the caption space
    (1152,)  ~ a SigLIP text embedding

WHY A CONV *AND* A GRU
----------------------
They fail differently, which is the only good reason to have both. The dilated
conv is order-aware but translation-equivariant: it detects "a thing moved
left-to-right" wherever in the window it happens, with a fixed receptive
field. The GRU is order-aware and unbounded: it can carry "the car was
stationary the whole time, then braked" across the window, which no
fixed-width kernel expresses. CLIP4Clip (arXiv 2104.08860) measured exactly
this axis on frozen CLIP features — meanP vs seqLSTM vs seqTransf — and found
learned temporal aggregation helps once the target actually depends on order.
Mean pooling, what `embeddings` does today, is the degenerate case of both.

TSM (Lin et al., ICCV 2019) was the alternative for the conv slot: shifting
channels along time is free. It was not used because its shift is a fixed
±1 tap; here the whole point is that the taps are *frequency-selective*, and
a fixed shift cannot express that.

WHY THE BASES ARE THE FDNN BASES
--------------------------------
A window's feature trajectory is a signal, and these bases were built for
signals. Applied after a temporal convolution:
  - a Gabor sub-function is a temporal wavelet — a burst detector, localised
    in time (a door opening, a brake light);
  - a FINER sub-function is a variable-period oscillator — periodic motion
    (gait, wipers, a turning wheel);
  - a polynomial-phase sub-function is a chirp — monotonic acceleration
    (a vehicle pulling away, a zoom).
The frequency bands then partition the TEMPORAL spectrum instead of every
neuron competing for it: slow = scene identity, mid = object motion, fast =
transitions. Each neuron being a sum over k of these is FDNN thesis 1, and
here it buys genuine expressiveness rather than decoration.

IDENTITY-SAFE INIT
------------------
The head is initialised so the tower emits the mean caption vector for every
window. That is "I know nothing" — a legal, centred prediction — so training
can only add information, and a half-trained tower can never be worse than
the corpus prior. Same discipline as FDNN's identity-init residual.
"""
from __future__ import annotations

import json
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np


# ===========================================================================
# Temporal hybrid layer — FDNN's HybridBiomimeticLayer with a time axis
# ===========================================================================
class TemporalHybridLayer(nn.Module):
    """Dilated temporal conv whose output channels are KAN sub-network neurons.

    Shapes: (B, T, in_d) -> (B, T, max_out_d).

    The convolution and the sub-function bank are fused: the conv produces
    `max_out_d * k` pre-activations per timestep, the bases are applied
    elementwise, and the k sub-functions of each neuron are summed (KAN).
    So one neuron is not "a channel" — it is a little ensemble of temporal
    filters that disagree about what shape of motion to look for.

    `omega_bands` are TEMPORAL frequencies here. They are an order of
    magnitude smaller than FDNN's coordinate-network defaults on purpose: the
    input is a PCA of unit-norm embeddings, so |h| is O(1) rather than O(100),
    and reusing omega=200 would put every neuron in the chaotic regime where
    gradients are noise.
    """

    def __init__(self, in_d, max_out_d, initial_active=None, k_width=4,
                 kernel=3, dilation=1, omega_bands=(2.0, 6.0, 18.0),
                 band_fractions=(0.34, 0.33, 0.33), bias_range=2.0,
                 use_residual=True, dropout_p=0.0, seed=0):
        super().__init__()
        rng = np.random.default_rng(seed)
        self.in_d = in_d
        self.max_out_d = max_out_d
        self.k = k_width
        self.kernel = kernel
        self.dilation = dilation
        self.bias_range = bias_range
        self.use_residual = use_residual and (in_d == max_out_d)
        self.dropout_p = dropout_p
        self._training = True

        # ---- frequency-banded omega per neuron (temporal spectrum split) ----
        omegas = []
        for om, fr in zip(omega_bands, band_fractions):
            omegas.extend([om] * int(round(fr * max_out_d)))
        omegas = (omegas + [omega_bands[-1]] * max_out_d)[:max_out_d]
        self.omegas_per_neuron = np.array(omegas, dtype=np.float32)
        om_exp = np.repeat(self.omegas_per_neuron, k_width)
        self.omegas = mx.array(om_exp)

        # ---- heterogeneous basis assignment: 50% FINER, 25% Gabor, 25% poly
        half, quarter = max(k_width // 2, 1), max(k_width // 4, 1)
        per = np.array([0] * half + [1] * quarter
                       + [3] * max(k_width - half - quarter, 0),
                       dtype=np.int32)[:k_width]
        if per.size < k_width:
            per = np.concatenate([per, np.zeros(k_width - per.size, np.int32)])
        self.basis_types = mx.array(np.tile(per, max_out_d).astype(np.int32))

        # ---- the temporal convolution: (kernel, in_d, max_out_d*k) ----------
        # Fan-in is kernel*in_d, and the SIREN convention divides by omega so
        # the pre-activation lands in the basis's useful range at init.
        mean_om = float(self.omegas_per_neuron.mean())
        limit = float(np.sqrt(6.0 / (kernel * in_d)) / mean_om)
        self.w1 = mx.array(rng.uniform(
            -limit, limit, (kernel, in_d, max_out_d * k_width)).astype(np.float32))
        self.b1 = mx.array(rng.uniform(
            -bias_range, bias_range, (max_out_d * k_width,)).astype(np.float32))
        self.phases = mx.array(rng.uniform(
            0, 2 * np.pi, (max_out_d * k_width,)).astype(np.float32))
        self.gabor_s = mx.array(rng.uniform(
            0.3, 1.5, (max_out_d * k_width,)).astype(np.float32))
        log_om = np.log(np.clip(om_exp, 1e-3, None))
        self.log_alpha = mx.array(
            rng.uniform(0.0, log_om).astype(np.float32))

        w2s = float(np.sqrt(6.0 / (max_out_d * k_width)))
        self.w2 = mx.array(rng.uniform(
            -w2s, w2s, (max_out_d, k_width)).astype(np.float32))

        m = np.zeros((max_out_d,), dtype=np.float32)
        m[:(initial_active or max_out_d)] = 1.0
        self.mask = mx.array(m)

        # basis_types is a CATEGORICAL SELECTOR, not a weight. Every mx.array
        # attribute joins the parameter tree, so without this the optimizer
        # takes gradient steps on it: int32 gets promoted to float32 and the
        # values drift (measured: 1.0 -> 0.9992). Dispatch is
        # `where(basis_types == 1, gabor, ...)`, so a drifted selector matches
        # nothing and every Gabor neuron silently falls through to the poly
        # branch — the heterogeneous basis quietly stops being heterogeneous.
        # Freezing keeps it out of trainable_parameters() entirely.
        self.freeze(keys=["basis_types"], recurse=False)

    # -- temporal gather: centred taps so the window is read bidirectionally --
    def _conv(self, x):
        """x: (B, T, in_d) -> (B, T, max_out_d*k)."""
        T = x.shape[1]
        off0 = (self.kernel - 1) // 2
        acc = None
        for t_i in range(self.kernel):
            shift = (t_i - off0) * self.dilation
            if shift == 0:
                xs = x
            elif shift > 0:                       # look forward, edge-pad tail
                xs = mx.concatenate(
                    [x[:, shift:, :],
                     mx.repeat(x[:, -1:, :], shift, axis=1)], axis=1)
            else:                                  # look back, edge-pad head
                s = -shift
                xs = mx.concatenate(
                    [mx.repeat(x[:, :1, :], s, axis=1),
                     x[:, :T - s, :]], axis=1)
            term = xs @ self.w1[t_i]
            acc = term if acc is None else acc + term
        return acc + self.b1

    def _bases(self, h):
        omega_h = self.omegas * h
        h_sq = h * h
        alpha = mx.exp(self.log_alpha)
        finer = mx.sin(self.omegas * (mx.abs(h) + 1.0) * h + self.phases)
        gabor = mx.exp(-(self.gabor_s ** 2) * h_sq) * mx.sin(omega_h + self.phases)
        sine = mx.sin(omega_h + self.phases)
        poly = mx.sin(alpha * h_sq + omega_h + self.phases)
        return mx.where(self.basis_types == 0, finer,
                        mx.where(self.basis_types == 1, gabor,
                                 mx.where(self.basis_types == 2, sine, poly)))

    def neuron_outputs(self, x):
        """(B, T, max_out_d) per-neuron signal BEFORE mask and residual.
        Deterministic regardless of mode — every pruning metric reads this."""
        acts = self._bases(self._conv(x))
        acts = acts.reshape(acts.shape[0], acts.shape[1], self.max_out_d, self.k)
        return mx.sum(acts * self.w2, axis=-1)

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

    def __call__(self, x):
        acts = self._bases(self._conv(x))
        acts = acts.reshape(acts.shape[0], acts.shape[1], self.max_out_d, self.k)
        if self._training and self.dropout_p > 0.0:
            keep = 1.0 - self.dropout_p
            drop = (mx.random.uniform(0.0, 1.0, acts.shape) < keep).astype(acts.dtype)
            acts = acts * drop / keep
        out = mx.sum(acts * self.w2, axis=-1) * self.mask
        return out + x if self.use_residual else out


# ===========================================================================
# Attention pooling over time
# ===========================================================================
class AttnPool(nn.Module):
    """Learned soft-argmax over the window.

    Mean pooling says every instant matters equally, which is wrong for a
    driving clip where 1.8 s is empty road and 0.2 s is the pedestrian
    stepping off the kerb. Attention lets the tower spend its output on the
    part of the window that carries the event.
    """

    def __init__(self, d, hidden=32):
        super().__init__()
        self.proj = nn.Linear(d, hidden)
        self.score = nn.Linear(hidden, 1)

    def __call__(self, h):                       # (B, T, d) -> (B, d)
        a = self.score(mx.tanh(self.proj(h)))    # (B, T, 1)
        w = mx.softmax(a, axis=1)
        return mx.sum(h * w, axis=1), w[..., 0]


# ===========================================================================
# The tower
# ===========================================================================
class ContextTower(nn.Module):
    def __init__(self, d_in=48, n_hidden=48, k_width=4, gru_hidden=32,
                 d_out=32, kernel=3, dilations=(1, 2), dropout_p=0.1,
                 input_noise_std=0.02, omega_bands=(2.0, 6.0, 18.0), seed=0):
        super().__init__()
        self.cfg = dict(d_in=d_in, n_hidden=n_hidden, k_width=k_width,
                        gru_hidden=gru_hidden, d_out=d_out, kernel=kernel,
                        dilations=list(dilations), dropout_p=dropout_p,
                        input_noise_std=input_noise_std,
                        omega_bands=list(omega_bands), seed=seed)
        self.input_noise_std = input_noise_std
        self._training = True

        self.blocks = []
        for i, dil in enumerate(dilations):
            self.blocks.append(TemporalHybridLayer(
                in_d=d_in if i == 0 else n_hidden,
                max_out_d=n_hidden, k_width=k_width, kernel=kernel,
                dilation=dil, omega_bands=omega_bands,
                use_residual=(i > 0), dropout_p=dropout_p, seed=seed + i))
        self.gru_f = nn.GRU(n_hidden, gru_hidden)
        self.gru_b = nn.GRU(n_hidden, gru_hidden)
        self.pool = AttnPool(n_hidden + 2 * gru_hidden)
        self.head = nn.Linear(n_hidden + 2 * gru_hidden, d_out)
        # Head starts at exactly zero: the tower's first prediction is the
        # corpus prior (see module docstring). Bias is zero too — the prior
        # itself is added outside, in the frozen decode basis.
        self.head.weight = mx.zeros(self.head.weight.shape)
        self.head.bias = mx.zeros((d_out,))
        self.log_gamma = mx.array(np.float32(np.log(0.5)))

    def set_training(self, mode: bool):
        self._training = mode
        for b in self.blocks:
            b._training = mode
        return self

    def turnover_layers(self):
        return list(self.blocks)

    def trunk(self, x):
        """(B, T, d_in) -> (B, n_hidden + 2*gru_hidden) pooled features."""
        if self._training and self.input_noise_std > 0:
            x = x + mx.random.normal(x.shape) * self.input_noise_std
        h = x
        for b in self.blocks:
            h = b(h)
        f = self.gru_f(h)
        b_rev = self.gru_b(h[:, ::-1, :])[:, ::-1, :]
        cat = mx.concatenate([h, f, b_rev], axis=-1)
        pooled, attn = self.pool(cat)
        # NOTE: deliberately not stashed on `self`. An mx.array assigned to a
        # Module attribute joins the parameter tree, and the optimizer then
        # tries to take a step on it (KeyError: 'last_attn'). Use
        # `attention_weights()` when you want to inspect it.
        return pooled

    def attention_weights(self, x):
        """(B, T) softmax weights — which instants the tower spent itself on."""
        h = x
        for b in self.blocks:
            h = b(h)
        f = self.gru_f(h)
        b_rev = self.gru_b(h[:, ::-1, :])[:, ::-1, :]
        return self.pool(mx.concatenate([h, f, b_rev], axis=-1))[1]

    def __call__(self, x):
        """-> (B, d_out) latent coordinates in the caption PCA basis."""
        return self.head(self.trunk(x)) * mx.exp(self.log_gamma)


# ===========================================================================
# Codec: frozen PCA in, frozen PCA out. Zero learned parameters either side.
# ===========================================================================
class ContextCodec:
    """Frozen PCA on the INPUT side of the tower.

    We have a few hundred labelled windows. A learned 1152->d encoder would be
    ~150k parameters fitted from ~150 examples, which is not learning, it is
    memorising. PCA is the optimal linear compressor under reconstruction
    error and costs zero training examples, so all the sample budget goes to
    the temporal dynamics — the only part mean pooling cannot already do.

    The OUTPUT side is not a PCA of SigLIP text embeddings. That was measured
    and it lost: ranking a query against a predicted caption *embedding*
    scored +0.218 against the VLM judge, below plain appearance search at
    +0.276, because SigLIP is trained for image-text similarity and its text
    tower is not calibrated for text-to-text comparison. The output space is
    now `CaptionSpace` — LSA over the caption text itself — which measured
    +0.339. See docs/CONTEXT.md.
    """

    def __init__(self, P_in, mu_in):
        self.P_in, self.mu_in = P_in, mu_in       # (d_in, D), (D,)

    @staticmethod
    def fit(frame_vecs, d_in=48):
        from sklearn.decomposition import PCA
        d_in = min(d_in, *frame_vecs.shape)
        pi = PCA(n_components=d_in, random_state=0).fit(frame_vecs)
        return ContextCodec(pi.components_.astype(np.float32),
                            pi.mean_.astype(np.float32))

    def encode(self, seq):                        # (T, D) -> (T, d_in)
        return (seq - self.mu_in) @ self.P_in.T

    def save(self, path):
        np.savez(path, P_in=self.P_in, mu_in=self.mu_in)

    @staticmethod
    def load(path):
        z = np.load(path)
        return ContextCodec(z["P_in"], z["mu_in"])


# ===========================================================================
# SigLIP's own loss
# ===========================================================================
def siglip_loss(v, u, ignore=None, log_t=None, bias=None):
    """Pairwise sigmoid loss (Zhai et al., ICCV 2023, arXiv 2303.15343).

    Chosen over softmax-InfoNCE for one concrete reason: the paper's own
    ablation shows sigmoid wins below ~16k batch, and our batch is the whole
    labelled set — a couple of hundred. Softmax normalises over the batch, so
    at this size the partition function is estimated from almost nothing.
    Sigmoid treats every pair as an independent binary problem and never
    needs that global view.

    `ignore` masks pairs that must not be counted as negatives. Windows slide
    with 75% overlap, so window i+1 genuinely depicts the same moment as
    window i; calling it a negative would teach the tower to separate
    identical content. Those pairs are dropped, not down-weighted.
    """
    t = mx.exp(log_t)
    logits = t * (v @ u.T) + bias
    n = v.shape[0]
    y = 2.0 * mx.eye(n) - 1.0                     # +1 diagonal, -1 elsewhere
    z = y * logits
    # -log sigmoid(z) == softplus(-z), stable for both signs
    per_pair = nn.softplus(-z)
    if ignore is not None:
        keep = 1.0 - ignore
        return mx.sum(per_pair * keep) / mx.maximum(mx.sum(keep), 1.0)
    return mx.mean(per_pair)


def overlap_mask(windows):
    """1.0 where two windows are the same stream and overlap in time."""
    n = len(windows)
    m = np.zeros((n, n), dtype=np.float32)
    st = [w[0] for w in windows]
    t0 = np.array([w[1] for w in windows], dtype=np.int64)
    t1 = np.array([w[2] for w in windows], dtype=np.int64)
    for i in range(n):
        same = np.array([s == st[i] for s in st])
        ov = same & (t0 < t1[i]) & (t1 > t0[i])
        m[i] = ov.astype(np.float32)
    np.fill_diagonal(m, 0.0)                      # the positive stays a positive
    return m


# ===========================================================================
# Training
# ===========================================================================
def train_tower(seqs, targets, windows, val_idx, cfg=None, epochs=400,
                lr=3e-3, weight_decay=1e-3, align_w=0.3, verbose=True,
                seed=0):
    """Fit the tower. `seqs` are (T, d_in) encoded sequences, `targets` are
    (d_out,) caption coordinates in the output basis, both already codec-mapped.

    Returns (model, history). Validation is a TIME split supplied by the
    caller — never a random split, because 75%-overlapping windows would put
    near-duplicates on both sides and report a fantasy score.
    """
    import mlx.optimizers as optim
    cfg = cfg or {}
    mx.random.seed(seed)
    model = ContextTower(**cfg)

    T = max(s.shape[0] for s in seqs)
    X = np.zeros((len(seqs), T, seqs[0].shape[1]), dtype=np.float32)
    for i, s in enumerate(seqs):                  # edge-pad short windows
        X[i, :len(s)] = s
        if len(s) < T:
            X[i, len(s):] = s[-1]
    Y = np.asarray(targets, dtype=np.float32)

    val = np.zeros(len(seqs), dtype=bool)
    val[val_idx] = True
    tr = ~val
    Xtr, Ytr = mx.array(X[tr]), mx.array(Y[tr])
    Xva, Yva = mx.array(X[val]), mx.array(Y[val])
    ig_tr = mx.array(overlap_mask([w for w, m in zip(windows, tr) if m]))
    ig_va = mx.array(overlap_mask([w for w, m in zip(windows, val) if m]))

    # logit scale/bias are learned, initialised as in the SigLIP paper
    state = {"log_t": mx.array(np.float32(np.log(10.0))),
             "bias": mx.array(np.float32(-10.0))}

    def _norm(a):
        # eps INSIDE the sqrt, not added to the norm afterwards. The head is
        # zero-initialised on purpose, so the very first forward pass produces
        # an exactly-zero vector — and d||x||/dx = x/||x|| is 0/0 = NaN there.
        # Smoothing the radicand keeps the gradient finite at the origin.
        return a * mx.rsqrt(mx.sum(a * a, axis=-1, keepdims=True) + 1e-8)

    def loss_fn(m, x, y, ig):
        v = _norm(m(x))
        u = _norm(y)
        # alignment term: pull each prediction onto its own caption. The
        # contrastive term only fixes ORDER; this one fixes absolute position,
        # which is what makes the vectors usable against unseen query text.
        align = mx.mean(1.0 - mx.sum(v * u, axis=-1))
        return siglip_loss(v, u, ignore=ig, log_t=state["log_t"],
                           bias=state["bias"]) + align_w * align

    opt = optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
    lg = nn.value_and_grad(model, loss_fn)
    sopt = optim.Adam(learning_rate=lr)

    def state_loss(s):
        model.set_training(False)
        v = _norm(model(Xtr))
        return siglip_loss(v, _norm(Ytr), ignore=ig_tr,
                           log_t=s["log_t"], bias=s["bias"])

    # The aliveness masks are mx.array attributes, so they live in
    # model.parameters() and AdamW would happily weight-decay them off their
    # 0/1 values — silently rescaling every neuron and breaking the
    # `mask == 1.0` alive-check the pruner depends on. Snapshot and restore
    # after every step. (FDNN hit exactly this; it is not hypothetical.)
    frozen_masks = [np.array(l.mask).copy() for l in model.turnover_layers()]

    hist = {"train": [], "val": [], "val_r1": []}
    best, best_w, best_state = 1e9, None, None
    for ep in range(epochs):
        model.set_training(True)
        loss, grads = lg(model, Xtr, Ytr, ig_tr)
        opt.update(model, grads)
        for lyr, fm in zip(model.turnover_layers(), frozen_masks):
            lyr.mask = mx.array(fm)
        sl, sg = mx.value_and_grad(state_loss)(state)
        state = sopt.apply_gradients(sg, state)   # plain dict, not a Module
        mx.eval(model.parameters(), opt.state, state)

        model.set_training(False)
        vl = float(loss_fn(model, Xva, Yva, ig_va).item())
        r1 = retrieval_r1(np.array(_norm(model(Xva))), np.array(_norm(Yva)))
        hist["train"].append(float(loss.item()))
        hist["val"].append(vl)
        hist["val_r1"].append(r1)
        if vl < best:
            best, best_w = vl, _clone_params(model)
            best_state = {k: mx.array(np.array(v)) for k, v in state.items()}
        if verbose and (ep % 50 == 0 or ep == epochs - 1):
            print(f"  ep {ep:4d} train {float(loss.item()):.4f} "
                  f"val {vl:.4f} val_R@1 {r1:.3f}", flush=True)
    if best_w is not None:
        _load_params(model, best_w)
        state = best_state
    model.set_training(False)
    return model, {"history": hist, "best_val": best,
                   "log_t": float(state["log_t"].item()),
                   "bias": float(state["bias"].item())}


def retrieval_r1(v, u):
    """Fraction of windows whose own caption is its nearest caption.

    This is the metric that matters: not "is the vector close to the target"
    but "does the vector RANK the right target first" — the same question the
    query path asks.
    """
    if len(v) < 2:
        return float("nan")
    s = v @ u.T
    return float((s.argmax(axis=1) == np.arange(len(v))).mean())


def _clone_params(model):
    from mlx.utils import tree_flatten
    return {k: np.array(v) for k, v in tree_flatten(model.parameters())}


def _load_params(model, flat):
    from mlx.utils import tree_unflatten
    model.update(tree_unflatten([(k, mx.array(v)) for k, v in flat.items()]))


def save_tower(model, codec, meta, path):
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    np.savez(path / "tower.npz", **_clone_params(model))
    codec.save(path / "codec.npz")
    # model.cfg goes LAST and wins. `meta` is often a carried-forward copy of
    # a previous tower.json and still holds that tower's cfg; letting it
    # override would save compacted weights under the pre-compaction shape,
    # and the next load would reshape-crash.
    (path / "tower.json").write_text(json.dumps(
        {**meta, "cfg": model.cfg}, indent=2))


def load_tower(path):
    path = Path(path)
    meta = json.loads((path / "tower.json").read_text())
    model = ContextTower(**meta["cfg"])
    z = np.load(path / "tower.npz")
    _load_params(model, {k: z[k] for k in z.files})
    model.set_training(False)
    return model, ContextCodec.load(path / "codec.npz"), meta