"""FDNN-V: a video-native embedding encoder, distilled from SigLIP. WHY THIS MODEL EXISTS --------------------- The database must embed EVERY frame at write time. SigLIP cannot do that at scale because it is a 428M-parameter model that treats each frame as an unrelated photograph — measured 13-90 ms/frame on this machine. But video is not a pile of photographs: frame t is almost entirely explained by frame t-1 (Deep Feature Flow, arXiv 1611.07715, built a whole recognition system on that observation). A human watching a video does not re-parse the scene 30 times a second either; they maintain a scene model and update it with what changed. FDNN-V is that shape, made of the three FDNN rules: spatial glimpse (cheap, per frame) "what is in front of me right now" Gabor-initialised first conv — V1 simple cells ARE Gabor filters small conv pyramid — ventral stream | FDNN temporal core (rule 1) "what has been going on" recurrent state h_t; each channel is a KAN-style sum over k heterogeneous sub-functions of (glimpse, state): FINER variable-period oscillator — periodic motion (gait, wipers) Gabor temporal wavelet — bursts (a grasp, a brake light) poly chirp — acceleration (pulling away) omega bands partition the temporal spectrum: slow = scene identity, mid = object motion, fast = transitions | head -> SigLIP space (1152-d) so every existing index, text query, and centroid keeps working unchanged rules 2+3 (apoptosis -> fine-tune -> neurogenesis -> fine-tune, decided by PPO + reverse attention) run post-training in fdnnv_prune — and unlike the attempt to prune SigLIP itself, the fine-tune step EXISTS here, because distillation pairs are free: every frame of the corpus already has a teacher embedding in `frame_vectors`. THE STREAMING CONTRACT ---------------------- `step(frame, h) -> (embedding, h')` is causal and O(1) per frame: no lookahead, no window buffer. That is what makes embed-on-write real — the encoder can sit inside the ingest loop and emit an embedding as each frame arrives, like any other index maintenance. WHAT DISTILLATION CAN AND CANNOT GIVE ------------------------------------- The student lands in the teacher's embedding space, so text queries (embedded by the frozen SigLIP text tower) keep working. It can match the teacher ON THIS CORPUS's manifold; it is not a zero-shot model for arbitrary imagery. That is the correct trade for a database: specialise the index to the data it serves, keep the teacher for what it is — an offline labeller. """ from __future__ import annotations import json from pathlib import Path import mlx.core as mx import mlx.nn as nn import numpy as np EMBED_DIM = 1152 # SigLIP so400m space — compatibility is the point INPUT_HW = (144, 192) # decode width 192 -> 192x144 (h, w) # =========================================================================== # V1: Gabor-initialised first convolution # =========================================================================== def gabor_bank(n_filters, ksize, rng): """Oriented Gabor filters spanning orientation x frequency x phase. Not decoration: the first layer of every competent visual system — biological or learned — converges to oriented band-pass filters. Starting there instead of at noise removes the epochs a small model would spend rediscovering V1. """ k = np.zeros((n_filters, ksize, ksize, 3), dtype=np.float32) half = ksize // 2 ys, xs = np.mgrid[-half:half + 1, -half:half + 1] for i in range(n_filters): theta = np.pi * (i % 8) / 8.0 lam = ksize / (1.5 + (i // 8) % 3) psi = 0.0 if (i // 24) % 2 == 0 else np.pi / 2 sigma = 0.4 * lam xr = xs * np.cos(theta) + ys * np.sin(theta) yr = -xs * np.sin(theta) + ys * np.cos(theta) g = np.exp(-(xr**2 + 0.8 * yr**2) / (2 * sigma**2)) \ * np.cos(2 * np.pi * xr / lam + psi) g -= g.mean() g /= (np.abs(g).sum() + 1e-6) # colour-opponent weighting: some filters luminance, some R-G, B-Y cw = [(1, 1, 1), (1, -1, 0), (0.5, 0.5, -1)][i % 3] for c in range(3): k[i, :, :, c] = g * cw[c] k += rng.normal(0, 0.01, k.shape).astype(np.float32) return k # =========================================================================== # Ventral stream: small depthwise-separable pyramid # =========================================================================== class DWBlock(nn.Module): def __init__(self, c_in, c_out, stride): super().__init__() self.dw = nn.Conv2d(c_in, c_in, 3, stride=stride, padding=1, groups=c_in) self.pw = nn.Conv2d(c_in, c_out, 1) self.norm = nn.LayerNorm(c_out) self.res = (c_in == c_out and stride == 1) def __call__(self, x): y = self.norm(self.pw(nn.silu(self.dw(x)))) return x + y if self.res else y class ConvStem(nn.Module): """Per-frame spatial encoder -> one glimpse vector. Small on purpose: fine discrimination is amortised into the temporal state, and capacity here is paid for EVERY frame FOREVER.""" def __init__(self, width=48, out_dim=256, seed=0): super().__init__() rng = np.random.default_rng(seed) w = width self.v1 = nn.Conv2d(3, w, 7, stride=2, padding=3) self.v1.weight = mx.array(gabor_bank(w, 7, rng)) self.s1 = DWBlock(w, w * 2, 2) self.s2 = DWBlock(w * 2, w * 4, 2) self.s3 = DWBlock(w * 4, w * 8, 2) self.r1 = DWBlock(w * 8, w * 8, 1) self.r2 = DWBlock(w * 8, w * 8, 1) self.proj = nn.Linear(w * 16, out_dim) self.norm = nn.LayerNorm(out_dim) def __call__(self, x): # (B, H, W, 3) in [-1, 1] h = nn.silu(self.v1(x)) h = self.s1(h) h = self.s2(h) h = self.s3(h) h = self.r2(self.r1(h)) mean = mx.mean(h, axis=(1, 2)) peak = mx.max(h, axis=(1, 2)) # mean = layout, max = salient return self.norm(self.proj(mx.concatenate([mean, peak], axis=-1))) class ViTStem(nn.Module): """The alternative stem: patchify + tiny transformer. Pure matmuls, which Apple-silicon GEMM kernels love; raced against ConvStem by measurement, never by taste.""" def __init__(self, dim=192, depth=4, heads=3, out_dim=256, patch=16, hw=INPUT_HW): super().__init__() self.patch = patch self.nh, self.nw = hw[0] // patch, hw[1] // patch self.embed = nn.Linear(patch * patch * 3, dim) self.pos = mx.zeros((1, self.nh * self.nw, dim)) self.blocks = [_TinyBlock(dim, heads) for _ in range(depth)] self.proj = nn.Linear(dim * 2, out_dim) self.norm = nn.LayerNorm(out_dim) def __call__(self, x): # (B, H, W, 3) B, H, W, _ = x.shape p = self.patch t = x.reshape(B, self.nh, p, self.nw, p, 3).transpose(0, 1, 3, 2, 4, 5) t = t.reshape(B, self.nh * self.nw, p * p * 3) t = self.embed(t) + self.pos for blk in self.blocks: t = blk(t) pooled = mx.concatenate([mx.mean(t, axis=1), mx.max(t, axis=1)], axis=-1) return self.norm(self.proj(pooled)) class _TinyBlock(nn.Module): def __init__(self, dim, heads): super().__init__() self.n1 = nn.LayerNorm(dim) self.att = nn.MultiHeadAttention(dim, heads) self.n2 = nn.LayerNorm(dim) self.fc1 = nn.Linear(dim, dim * 2) self.fc2 = nn.Linear(dim * 2, dim) def __call__(self, x): y = self.n1(x) x = x + self.att(y, y, y) return x + self.fc2(nn.silu(self.fc1(self.n2(x)))) # =========================================================================== # Rule 1: the temporal core — every channel is a sub-network # =========================================================================== class FDNNTemporalCell(nn.Module): """Recurrent state whose channels are KAN-style sums over k heterogeneous temporal sub-functions, gated GRU-style so state persists by default. The bases read the JOINT signal (current glimpse, previous state), so a FINER channel can oscillate with repeated motion, a Gabor channel can fire on a burst of change, and a poly-phase channel can track acceleration — while the gate decides how much of the old scene model each step is allowed to overwrite. This is FDNN's HybridBiomimeticLayer with time as the signal axis instead of a coordinate. """ def __init__(self, in_dim=256, channels=256, k_width=4, omega_bands=(0.8, 2.5, 8.0), band_fractions=(0.34, 0.33, 0.33), seed=0): # omega bands are LOWER than the context tower's (2, 6, 18) on # purpose: that tower convolved over time feed-forward, this cell # FEEDS BACK. sin bases at omega 15 inside a 32-step recurrence give # chaotic gradients — measured: stage-2 loss climbed from 0.086 to # 0.112 and val fidelity fell 0.019. Same bases, calmer spectrum. super().__init__() rng = np.random.default_rng(seed) C, k = channels, k_width self.C, self.k = C, k omegas = [] for om, fr in zip(omega_bands, band_fractions): omegas.extend([om] * int(round(fr * C))) omegas = (omegas + [omega_bands[-1]] * C)[:C] self.omegas_per_neuron = np.array(omegas, dtype=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([0] * half + [1] * quarter + [3] * max(k - half - quarter, 0), np.int32)[:k] self.basis_types = mx.array(np.tile(per, C).astype(np.int32)) # A categorical selector, not a weight: unfrozen, the optimizer # promotes it to float and drifts it off its exact values, silently # rerouting every Gabor neuron to the poly branch (measured on the # context tower: 1.0 -> 0.9992 and `== 1` matched nothing). self.freeze(keys=["basis_types"], recurse=False) mean_om = float(self.omegas_per_neuron.mean()) lim_g = float(np.sqrt(6.0 / in_dim) / mean_om) lim_h = float(np.sqrt(6.0 / C) / mean_om) self.Wg = mx.array(rng.uniform(-lim_g, lim_g, (in_dim, C * k)).astype(np.float32)) self.Wh = mx.array(rng.uniform(-lim_h, lim_h, (C, 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)) # alpha starts log-uniform between 1 and omega. With sub-unit omegas # (the calmed recurrent bands) log-omega is negative, so the interval # must be ordered explicitly — uniform(0, negative) is an error. 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)) zlim = float(np.sqrt(6.0 / (in_dim + C))) self.Wz = mx.array(rng.uniform(-zlim, zlim, (in_dim, C)).astype(np.float32)) self.Uz = mx.array(rng.uniform(-zlim, zlim, (C, C)).astype(np.float32)) # Gate bias starts NEGATIVE: sigmoid(-1) ~ 0.27, so at init the state # persists — a scene model that forgets everything every frame is just # a per-frame model with extra steps. self.bz = mx.array(np.full((C,), -1.0, np.float32)) self.mask = mx.array(np.ones((C,), np.float32)) # aliveness (rule 2) # Aliveness is set by the pruning cycle, never by the optimizer — # unfrozen, AdamW weight-decays it off 1.0 and every channel quietly # shrinks (the measured context-tower failure mode). self.freeze(keys=["mask"], recurse=False) def set_active_mask(self, m): self.mask = mx.array(np.asarray(m, dtype=np.float32)) def _cand(self, g, h): pre = g @ self.Wg + h @ self.Wh + 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 == 0, finer, mx.where(self.basis_types == 1, gab, mx.where(self.basis_types == 2, sine, poly))) acts = acts.reshape(-1, self.C, self.k) return mx.sum(acts * self.w2, axis=-1) def neuron_outputs(self, g, h): """Per-neuron candidate BEFORE gate and mask — the pruning signal.""" return self._cand(g, h) def __call__(self, g, h): cand = self._cand(g, h) * self.mask z = mx.sigmoid(g @ self.Wz + h @ self.Uz + self.bz) * self.mask return (1.0 - z) * h + z * cand # =========================================================================== # The encoder # =========================================================================== class FDNNVideoEncoder(nn.Module): def __init__(self, stem="conv", stem_width=48, glimpse=256, channels=256, k_width=4, embed_dim=EMBED_DIM, vit_depth=4, seed=0): super().__init__() self.cfg = dict(stem=stem, stem_width=stem_width, glimpse=glimpse, channels=channels, k_width=k_width, embed_dim=embed_dim, vit_depth=vit_depth, seed=seed) if stem == "conv": self.stem = ConvStem(width=stem_width, out_dim=glimpse, seed=seed) else: self.stem = ViTStem(dim=stem_width * 4, depth=vit_depth, out_dim=glimpse) self.cell = FDNNTemporalCell(in_dim=glimpse, channels=channels, k_width=k_width, seed=seed) self.head_g = nn.Linear(glimpse, embed_dim) # Temporal head starts at zero: at init the model IS the per-frame # model (stage 1), and training can only add information from state. # Same identity-safe discipline as every other init in this repo. self.head_h = nn.Linear(channels, embed_dim) self.head_h.weight = mx.zeros(self.head_h.weight.shape) self.head_h.bias = mx.zeros((embed_dim,)) self.channels = channels # ---- streaming: this is the embed-on-write contract ------------------- def init_state(self, batch=1): return mx.zeros((batch, self.channels)) def step(self, frame, h): """One frame in, one embedding out, O(1) state carried. Causal.""" g = self.stem(frame) h = self.cell(g, h) e = self.head_g(g) + self.head_h(h) return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8), h # ---- batched sequences (training / bulk ingest) ----------------------- def __call__(self, seq, h0=None): """(B, T, H, W, 3) -> (B, T, D). Stem runs on all frames as one big batch (the GEMM-friendly part); only the tiny cell recurs.""" B, T = seq.shape[0], seq.shape[1] g = self.stem(seq.reshape(B * T, *seq.shape[2:])).reshape(B, T, -1) h = self.init_state(B) if h0 is None else h0 outs = [] for t in range(T): h = self.cell(g[:, t], h) outs.append(h) hs = mx.stack(outs, axis=1) e = self.head_g(g) + self.head_h(hs) return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8), h def embed_frames_np(self, frames_u8, batch=64, chunk=None): """uint8 (N, H, W, 3) of ONE stream, in time order -> (N, 1152). Stateful across batches — one continuous pass over the stream.""" h = self.init_state(1) out = [] for i in range(0, len(frames_u8), batch): x = mx.array(frames_u8[i:i + batch].astype(np.float32) / 127.5 - 1.0)[None] e, h = self(x, h0=h) out.append(np.array(e[0], dtype=np.float32)) return np.concatenate(out, axis=0) # =========================================================================== # Distillation loss: pointwise + affinity mimicking # =========================================================================== def distill_loss(student, teacher, affinity_w=0.25, mu=None, centered_w=1.0, anchors=None, anchor_w=50.0): """Distillation aimed at RETRIEVAL, not at raw closeness. Plain pointwise cosine is a trap on a homogeneous corpus: every teacher vector shares a huge common mode, so matching that alone buys ~0.9 cosine while scrambling the thin discriminative residual that ranking runs on. Measured: a student at fidelity 0.907 kept only 4.4% of the teacher's top-10 neighbours, while the teacher AGAINST ITSELF at a different input resolution — fidelity 0.921 — keeps 50.7%. Same closeness, 10x the retrieval agreement: the difference is WHERE the error lives. So three additional terms put the error where it does no harm: centered cosine on (v - mu): the mean-free residual is exactly what ranking compares, so it gets its own gradient. affinity within-batch similarity matching (TinyCLIP, arXiv 2309.12314): preserve the teacher's ordering structure. anchors similarity profile against real caption-text embeddings from this store: text queries live in those directions, and image-text sims occupy a band ~50x narrower than image-image sims — hence the weight. """ t = teacher * mx.rsqrt(mx.sum(teacher * teacher, axis=-1, keepdims=True) + 1e-8) loss = mx.mean(1.0 - mx.sum(student * t, axis=-1)) if mu is not None and centered_w > 0: sc = student - mu tc = t - mu sc = sc * mx.rsqrt(mx.sum(sc * sc, axis=-1, keepdims=True) + 1e-8) tc = tc * mx.rsqrt(mx.sum(tc * tc, axis=-1, keepdims=True) + 1e-8) loss = loss + centered_w * mx.mean(1.0 - mx.sum(sc * tc, axis=-1)) if affinity_w > 0: s2 = student.reshape(-1, student.shape[-1]) t2 = t.reshape(-1, t.shape[-1]) loss = loss + affinity_w * mx.mean(mx.square(s2 @ s2.T - t2 @ t2.T)) if anchors is not None and anchor_w > 0: sa = student.reshape(-1, student.shape[-1]) @ anchors.T ta = t.reshape(-1, t.shape[-1]) @ anchors.T loss = loss + anchor_w * mx.mean(mx.square(sa - ta)) return loss # =========================================================================== # persistence # =========================================================================== def save_encoder(model, meta, path): from mlx.utils import tree_flatten path = Path(path) path.mkdir(parents=True, exist_ok=True) np.savez(path / "weights.npz", **{k: np.array(v) for k, v in tree_flatten(model.parameters())}) (path / "encoder.json").write_text(json.dumps( {**meta, "cfg": model.cfg}, indent=2)) def fdnnv_dir() -> Path: """Where the FDNN-V encoder lives. Repo-level `models/fdnnv`, NOT inside any store. The encoder used to sit at `lake/bridge/models/fdnnv`, so clearing the stores deleted a trained model along with the data (2026-07-28). A model is not store data. The legacy path is still accepted for stores that predate the move.""" here = Path("models/fdnnv") if (here / "encoder.json").exists(): return here legacy = Path("lake/bridge/models/fdnnv") if (legacy / "encoder.json").exists(): return legacy return here def load_encoder(path): from mlx.utils import tree_unflatten path = Path(path) meta = json.loads((path / "encoder.json").read_text()) model = FDNNVideoEncoder(**meta["cfg"]) z = np.load(path / "weights.npz") model.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files])) model.cell.freeze(keys=["basis_types", "mask"], recurse=False) mx.eval(model.parameters()) return model, meta # =========================================================================== # The write path: chunked byte-range decode feeding the streaming encoder # =========================================================================== def embed_stream(store, model, rows, width=192, chunk=512, batch_cb=None): """Embed one stream's frames in time order, state carried across chunks. `chunk` bounds the decoder subprocess's rawvideo buffer (~1 GB at 512 frames of 640x480); the encoder state flows straight through, so the result is identical to one infinite pass. This loop is the write path: ingest can call it as frames land. Returns (ts int64 array, vectors float32 (N, D), decode_s, embed_s). """ import time as _time from .video import FrameSet ts_out, vecs = [], [] h = model.init_state(1) dec_s = emb_s = 0.0 for i in range(0, len(rows), chunk): t0 = _time.perf_counter() dec = FrameSet(store, "frames", rows.slice(i, chunk)).decode( width=width) dec_s += _time.perf_counter() - t0 if not dec: continue frames = np.stack([d[1] for d in dec]) # the stem takes EXACTLY (in_h, in_w); sources with a different # aspect ratio decode to other shapes (lab video came back square # and crashed the reshape). Stretch — the encoder was distilled on # stretched frames, so aspect distortion is in-distribution. ih, iw = model.cfg["in_hw"] if "in_hw" in model.cfg else (144, 192) if frames.shape[1] != ih or frames.shape[2] != iw: xr = np.linspace(0, frames.shape[2] - 1, iw).round().astype(int) yr = np.linspace(0, frames.shape[1] - 1, ih).round().astype(int) frames = frames[:, yr][:, :, xr] t0 = _time.perf_counter() x = mx.array(frames.astype(np.float32) / 127.5 - 1.0)[None] e, h = model(x, h0=h) e = np.array(e[0], dtype=np.float32) emb_s += _time.perf_counter() - t0 ts_out.extend(d[0] for d in dec) vecs.append(e) if batch_cb: batch_cb(len(ts_out)) if not vecs: return np.array([], np.int64), np.zeros((0, EMBED_DIM), np.float32), \ dec_s, emb_s return (np.array(ts_out, np.int64), np.concatenate(vecs), dec_s, emb_s)