python/elidedb/ctxprune.py

"""Post-training cellular turnover for the context tower — PPO + reverse
attention, ported from FDNN_BrainModel/models/pruning.py.

WHY PRUNE AT ALL
----------------
Because the only thing this index is for is answering fast. Every channel the
tower carries is multiplied over T timesteps, for every window, at every
rebuild. Capacity that is not earning its keep is latency the user pays for
forever. And with a few hundred labelled windows, an over-provisioned tower
is also memorising — so apoptosis buys accuracy and speed at the same time,
which is exactly the regime cellular turnover was designed for.

WHAT CHANGED FROM FDNN, AND WHY
-------------------------------
FDNN prunes per (layer, neuron): every layer has its own aliveness mask. That
is right for a plain stack. This tower is a RESIDUAL stack, and in a residual
stack channel j is a wire that runs the whole depth — masking it in layer 2
alone does not remove it, because the skip connection keeps carrying it. So
the prune unit here is the residual-stream CHANNEL, decided once and applied
to every block.

That is not a weakening of the idea, it is the correct form of it for this
topology (the same structured channel pruning people apply to ResNets), and
it has a property per-layer masking does not: a dead channel can be
physically deleted. `compact()` rebuilds the tower without it, so the mask
turns into real wall-clock speed instead of a multiply by zero. Pruning that
does not shrink the matmul is not pruning, it is decoration.

THE THREE MEASURES (unchanged in spirit from FDNN)
--------------------------------------------------
  utilization      Δ validation RETRIEVAL loss when the channel is silenced.
                   Not weight magnitude — the thing we actually care about.
  reverse attention softmax(-utilization): mass on the channels the network
                   attends AWAY from, i.e. the prune candidates.
  PPO              a clipped-PPO contextual bandit over per-channel features
                   that samples masks, scores them by
                   -(loss ratio) - λ·kept_fraction, and learns which channels
                   can actually go.
"""
from __future__ import annotations

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


# ===========================================================================
# 0. Channel mask plumbing
# ===========================================================================
def n_channels(model):
    return model.blocks[0].max_out_d


def get_channel_mask(model):
    return np.array(model.blocks[0].mask).copy()


def set_channel_mask(model, mask):
    """One decision, applied to every block — see module docstring."""
    m = np.asarray(mask, dtype=np.float32)
    for b in model.turnover_layers():
        b.set_active_mask(m)


# ===========================================================================
# 1. Utilization ("consumption")
# ===========================================================================
def measure_utilization(model, loss_of, verbose=False):
    """Per-channel ablation importance + the cheap descriptive stats.

    `loss_of(model)` must return the validation retrieval loss as a float.
    Ablation is the increase in that loss when a single channel is silenced
    everywhere — the most directly interpretable notion of how much the tower
    relies on it.
    """
    model.set_training(False)
    base = loss_of(model)
    mask = get_channel_mask(model)
    n = len(mask)

    abl = np.zeros(n, dtype=np.float32)
    for c in np.where(mask == 1.0)[0]:
        probe = mask.copy()
        probe[c] = 0.0
        set_channel_mask(model, probe)
        abl[c] = max(loss_of(model) - base, 0.0)
    set_channel_mask(model, mask)
    return {"mask": mask, "ablation": abl, "base_loss": base}


def descriptive_stats(model, X):
    """Activation magnitude per channel (averaged over blocks) and the
    downstream weight norm each channel feeds."""
    model.set_training(False)
    outs = []
    h = mx.array(X)
    for b in model.turnover_layers():
        outs.append(np.array(mx.mean(mx.abs(b.neuron_outputs(h)), axis=(0, 1))))
        h = b(h)
    act = np.mean(np.stack(outs), axis=0)

    n = n_channels(model)
    wf = np.array(model.gru_f.Wx)          # (3*hidden, n_hidden)
    wb = np.array(model.gru_b.Wx)
    hw = np.array(model.head.weight)       # (d_out, n_hidden + 2*gru_hidden)
    pw = np.array(model.pool.proj.weight)
    down = (np.linalg.norm(wf, axis=0) + np.linalg.norm(wb, axis=0)
            + np.linalg.norm(hw[:, :n], axis=0)
            + np.linalg.norm(pw[:, :n], axis=0))
    return act.astype(np.float32), down.astype(np.float32)


# ===========================================================================
# 2. Reverse attention
# ===========================================================================
def reverse_attention(importance, mask, temperature=1.0):
    """Forward attention concentrates on useful units; reverse attention
    inverts it so the LEAST useful surface as prune candidates. Dead channels
    score 0; scores over the live set sum to 1."""
    importance = np.asarray(importance, dtype=np.float64)
    alive = mask == 1.0
    r = np.zeros_like(importance)
    if alive.sum() == 0:
        return r.astype(np.float32)
    imp = importance[alive]
    imp = (imp - imp.mean()) / (imp.std() + 1e-8)
    logits = -imp / max(temperature, 1e-6)
    logits -= logits.max()
    e = np.exp(logits)
    r[alive] = e / (e.sum() + 1e-12)
    return r.astype(np.float32)


# ===========================================================================
# 3. PPO agent
# ===========================================================================
class PrunePolicy(nn.Module):
    def __init__(self, feat_dim, hidden=32):
        super().__init__()
        self.fc1 = nn.Linear(feat_dim, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.pi_head = nn.Linear(hidden, 1)
        self.v_head = nn.Linear(hidden, 1)

    def __call__(self, feats):
        h = nn.relu(self.fc1(feats))
        h = nn.relu(self.fc2(h))
        return self.pi_head(h)[:, 0], self.v_head(h)[:, 0]


def _bernoulli_logp(actions, logits):
    return actions * logits - mx.logaddexp(mx.zeros_like(logits), logits)


def _entropy(logits):
    p = mx.sigmoid(logits)
    return nn.softplus(logits) - p * logits


class PPOPruner:
    """Contextual bandit over prune masks of the residual stream.

    One step per episode: observe a fixed per-channel feature matrix from the
    frozen tower, emit a keep-probability per channel, sample a binary mask,
    receive one scalar reward. Maximising it means "prune as much as possible
    without inflating retrieval loss".
    """

    def __init__(self, model, X, loss_of, sparsity_coef=0.35,
                 loss_tolerance=1.0, hidden=32, lr=3e-3, clip_eps=0.2,
                 entropy_coef=0.01, value_coef=0.5, seed=0):
        self.model, self.loss_of = model, loss_of
        self.sparsity_coef = sparsity_coef
        self.loss_tolerance = loss_tolerance
        self.clip_eps, self.entropy_coef = clip_eps, entropy_coef
        self.value_coef = value_coef
        np.random.seed(seed)
        mx.random.seed(seed)

        self.base_mask = get_channel_mask(model)
        util = measure_utilization(model, loss_of)
        self.base_loss = util["base_loss"]
        self.ablation = util["ablation"]
        act, down = descriptive_stats(model, X)
        ra = reverse_attention(self.ablation, self.base_mask)
        omega = model.blocks[0].omegas_per_neuron

        self.candidates = np.where(self.base_mask == 1.0)[0]
        f = np.stack([self.ablation[self.candidates], act[self.candidates],
                      down[self.candidates], ra[self.candidates],
                      omega[self.candidates]], axis=1).astype(np.float32)
        mu, sd = f.mean(0, keepdims=True), f.std(0, keepdims=True) + 1e-6
        f[:, :4] = (f[:, :4] - mu[:, :4]) / sd[:, :4]
        self.feats = mx.array(f)
        self.reverse_attn = ra
        self.policy = PrunePolicy(f.shape[1], hidden=hidden)
        self.opt = optim.Adam(learning_rate=lr)
        self.history = {"iter": [], "reward": [], "kept_frac": [],
                        "loss_ratio": [], "entropy": []}

    def _apply(self, action):
        m = self.base_mask.copy()
        m[self.candidates] = action
        set_channel_mask(self.model, m)

    def _reward(self, action):
        self._apply(action)
        loss = self.loss_of(self.model)
        set_channel_mask(self.model, self.base_mask)
        kept = float(action.mean()) if len(action) else 0.0
        ratio = loss / (self.base_loss + 1e-12)
        return (-self.loss_tolerance * ratio - self.sparsity_coef * kept,
                kept, ratio)

    def train(self, n_iters=40, episodes_per_iter=12, ppo_epochs=4,
              verbose=True):
        for it in range(n_iters):
            logits, _ = self.policy(self.feats)
            probs = np.array(mx.sigmoid(logits))
            acts, rews, kfs, lrs = [], [], [], []
            for _ in range(episodes_per_iter):
                a = (np.random.uniform(size=probs.shape) < probs).astype(np.float32)
                r, kf, lr_ = self._reward(a)
                acts.append(a); rews.append(r); kfs.append(kf); lrs.append(lr_)
            acts = np.array(acts, np.float32)
            rews = np.array(rews, np.float32)
            adv = (rews - rews.mean()) / (rews.std() + 1e-6)

            acts_mx, adv_mx, rews_mx = mx.array(acts), mx.array(adv), mx.array(rews)
            old_logits, _ = self.policy(self.feats)
            old_logp = mx.array(np.array(
                mx.sum(_bernoulli_logp(acts_mx, old_logits[None, :]), axis=1)))
            for _ in range(ppo_epochs):
                _, grads = self._loss_and_grad(acts_mx, old_logp, adv_mx, rews_mx)
                self.opt.update(self.policy, grads)
                mx.eval(self.policy.parameters(), self.opt.state)

            ent = float(mx.mean(_entropy(logits)).item())
            self.history["iter"].append(it)
            self.history["reward"].append(float(rews.mean()))
            self.history["kept_frac"].append(float(np.mean(kfs)))
            self.history["loss_ratio"].append(float(np.mean(lrs)))
            self.history["entropy"].append(ent)
            if verbose and (it % 5 == 0 or it == n_iters - 1):
                print(f"  PPO {it:3d} | reward {rews.mean():+.3f} | kept "
                      f"{np.mean(kfs):.2f} | loss x{np.mean(lrs):.3f} | "
                      f"H {ent:.3f}", flush=True)
        return self.history

    def _loss_and_grad(self, actions, old_logp, adv, rewards):
        def loss_fn(policy):
            logits, values = policy(self.feats)
            v = mx.mean(values)
            logp = mx.sum(_bernoulli_logp(actions, logits[None, :]), axis=1)
            ratio = mx.exp(logp - old_logp)
            pl = -mx.mean(mx.minimum(ratio * adv,
                                     mx.clip(ratio, 1 - self.clip_eps,
                                             1 + self.clip_eps) * adv))
            vl = mx.mean(mx.square(v - rewards))
            return pl + self.value_coef * vl \
                - self.entropy_coef * mx.mean(_entropy(logits))
        return nn.value_and_grad(self.policy, loss_fn)(self.policy)

    def apply_greedy(self, keep_threshold=0.5):
        logits, _ = self.policy(self.feats)
        probs = np.array(mx.sigmoid(logits))
        m = self.base_mask.copy()
        m[self.candidates] = (probs >= keep_threshold).astype(np.float32)
        set_channel_mask(self.model, m)
        return m, probs


# ===========================================================================
# 4. Neurogenesis — revive dead channels with fresh, spectrally diverse ones
# ===========================================================================
def neurogenesis(model, n_revive, seed=0):
    """Rebirth. New channels start near-silent (tiny w2) so they cannot shock
    the forward pass, and are re-initialised with FINER's WIDE bias range —
    without that they spawn as low-frequency clones and the population slowly
    loses the spectral diversity that made the KAN basis worth having."""
    rng = np.random.default_rng(seed)
    mask = get_channel_mask(model)
    dead = np.where(mask == 0.0)[0]
    if len(dead) == 0 or n_revive <= 0:
        return 0
    revive = dead[:min(n_revive, len(dead))]

    for b in model.turnover_layers():
        w1 = np.array(b.w1)          # (kernel, in_d, out*k)
        b1 = np.array(b.b1)
        w2 = np.array(b.w2)
        ph = np.array(b.phases)
        gs = np.array(b.gabor_s)
        la = np.array(b.log_alpha)
        om_full = np.repeat(b.omegas_per_neuron, b.k)
        mean_om = float(b.omegas_per_neuron.mean())
        lim = float(np.sqrt(6.0 / (b.kernel * b.in_d)) / mean_om)
        for c in revive:
            s, e = c * b.k, (c + 1) * b.k
            w1[:, :, s:e] = rng.uniform(-lim, lim, (b.kernel, b.in_d, b.k))
            b1[s:e] = rng.uniform(-b.bias_range, b.bias_range, b.k)
            ph[s:e] = rng.uniform(0, 2 * np.pi, b.k)
            gs[s:e] = rng.uniform(0.3, 1.5, b.k)
            lw = np.log(np.clip(om_full[s:e], 1e-3, None))
            la[s:e] = rng.uniform(np.zeros_like(lw), lw)
            w2[c, :] = rng.standard_normal(b.k) * 1e-6
        b.w1 = mx.array(w1.astype(np.float32))
        b.b1 = mx.array(b1.astype(np.float32))
        b.w2 = mx.array(w2.astype(np.float32))
        b.phases = mx.array(ph.astype(np.float32))
        b.gabor_s = mx.array(gs.astype(np.float32))
        b.log_alpha = mx.array(la.astype(np.float32))
    mask[revive] = 1.0
    set_channel_mask(model, mask)
    return len(revive)


# ===========================================================================
# 5. Compaction — where the mask becomes actual speed
# ===========================================================================
def compact(model):
    """Physically delete dead channels: return a smaller, equivalent tower.

    Every dead channel is removed from each block's output slots, from every
    consumer's input slots (both GRUs, the attention pool, the head), and from
    the residual stream itself. Output is numerically equivalent to the masked
    model but the matmuls are genuinely smaller.
    """
    from .ctxtower import ContextTower
    keep = np.where(get_channel_mask(model) == 1.0)[0]
    n_new = len(keep)
    if n_new == 0:
        raise RuntimeError("every channel was pruned")
    if n_new == n_channels(model):
        return model, keep

    cfg = dict(model.cfg)
    cfg["n_hidden"] = int(n_new)
    new = ContextTower(**cfg)
    k = model.blocks[0].k
    sub = np.concatenate([np.arange(c * k, (c + 1) * k) for c in keep])

    for ob, nb in zip(model.blocks, new.blocks):
        w1 = np.array(ob.w1)
        if not ob.use_residual:                  # first block: in_d is d_in
            nb.w1 = mx.array(w1[:, :, sub])
        else:                                    # in_d is the stream: slice both
            nb.w1 = mx.array(w1[:, keep][:, :, sub])
        nb.b1 = mx.array(np.array(ob.b1)[sub])
        nb.phases = mx.array(np.array(ob.phases)[sub])
        nb.gabor_s = mx.array(np.array(ob.gabor_s)[sub])
        nb.log_alpha = mx.array(np.array(ob.log_alpha)[sub])
        nb.w2 = mx.array(np.array(ob.w2)[keep])
        # omegas is TRAINED, so it must be sliced from the live array — not
        # rebuilt from omegas_per_neuron, which is the frozen design band and
        # is stale the moment training starts. Rebuilding it silently reverted
        # every learned frequency and moved val loss by 0.015.
        nb.omegas = mx.array(np.array(ob.omegas)[sub])
        nb.omegas_per_neuron = ob.omegas_per_neuron[keep]   # design band only
        nb.basis_types = mx.array(np.array(ob.basis_types)[sub])
        nb.freeze(keys=["basis_types"], recurse=False)
        nb.set_active_mask(np.ones(n_new, np.float32))

    n_old = n_channels(model)
    g = model.cfg["gru_hidden"]
    for og, ng in ((model.gru_f, new.gru_f), (model.gru_b, new.gru_b)):
        ng.Wx = mx.array(np.array(og.Wx)[:, keep])
        ng.Wh = mx.array(np.array(og.Wh))
        ng.b = mx.array(np.array(og.b))
        ng.bhn = mx.array(np.array(og.bhn))
    # consumers of `cat` = [stream | gru_f | gru_b]: keep the stream slice,
    # carry the two GRU slices through untouched.
    tail = np.arange(n_old, n_old + 2 * g)
    cols = np.concatenate([keep, tail])
    new.pool.proj.weight = mx.array(np.array(model.pool.proj.weight)[:, cols])
    new.pool.proj.bias = mx.array(np.array(model.pool.proj.bias))
    new.pool.score.weight = mx.array(np.array(model.pool.score.weight))
    new.pool.score.bias = mx.array(np.array(model.pool.score.bias))
    new.head.weight = mx.array(np.array(model.head.weight)[:, cols])
    new.head.bias = mx.array(np.array(model.head.bias))
    new.log_gamma = mx.array(np.array(model.log_gamma))
    new.set_training(False)
    return new, keep


# ===========================================================================
# 6. Re-settle
# ===========================================================================
def finetune(model, loss_and_grad_fn, epochs=150, lr=5e-4,
             weight_decay=1e-3, loss_of=None, patience=40):
    """Short re-settling pass with aliveness frozen (see the mask gotcha).

    Validation-aware on purpose. Survivors re-settle on the TRAIN set, and
    with a couple of hundred windows an unguarded 400-epoch re-settle simply
    re-memorises: measured, it took a pruned tower from 0.311 val back up to
    0.357, undoing everything apoptosis had just gained. Keeping the best
    validation state means re-settling can only help, never hurt.
    """
    from .ctxtower import _clone_params, _load_params
    opt = optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
    frozen = get_channel_mask(model)
    best, best_w, since = float("inf"), None, 0
    if loss_of is not None:
        best, best_w = loss_of(model), _clone_params(model)
    for _ in range(epochs):
        model.set_training(True)
        loss, grads = loss_and_grad_fn(model)
        opt.update(model, grads)
        set_channel_mask(model, frozen)
        mx.eval(model.parameters(), opt.state)
        if loss_of is None:
            continue
        v = loss_of(model)
        if v < best - 1e-5:
            best, best_w, since = v, _clone_params(model), 0
        else:
            since += 1
            if since >= patience:
                break
    if best_w is not None:
        _load_params(model, best_w)
        set_channel_mask(model, frozen)
        mx.eval(model.parameters())
    model.set_training(False)
    return best


# ===========================================================================
# 7. The full cycle
# ===========================================================================
def run_pruning_cycle(model, X, loss_of, loss_and_grad_fn, ppo_iters=40,
                      episodes_per_iter=12, sparsity_coef=0.35,
                      keep_threshold=0.5, finetune_epochs=150,
                      rebirth_fraction=0.5, select_tolerance=0.02,
                      seed=0, verbose=True):
    """apoptosis → re-settle → neurogenesis → re-settle, on a frozen tower.

    The cycle is a SEARCH, so it returns the best point it visited, not the
    last one. Rebirth is a bet that the data can support more capacity; on a
    small corpus that bet often loses, and silently shipping the final state
    would hand back a tower worse than the one we started from.
    """
    from .ctxtower import _clone_params, _load_params
    rec = {"stages": []}
    saved = []          # (stage, loss, alive, params, mask) for every stage

    def snap(stage):
        u = measure_utilization(model, loss_of)
        alive = int(u["mask"].sum())
        if verbose:
            print(f"[{stage}] val loss {u['base_loss']:.4f} | "
                  f"alive {alive}/{len(u['mask'])}", flush=True)
        rec["stages"].append({"stage": stage, "loss": u["base_loss"],
                              "alive": alive,
                              "ablation": u["ablation"].tolist(),
                              "mask": u["mask"].tolist()})
        saved.append((stage, u["base_loss"], alive,
                      _clone_params(model), u["mask"].copy()))
        return u

    snap("before")
    if verbose:
        print("\n== PPO reverse-attention pruning search ==", flush=True)
    pruner = PPOPruner(model, X, loss_of, sparsity_coef=sparsity_coef, seed=seed)
    rec["ppo_history"] = pruner.train(n_iters=ppo_iters,
                                      episodes_per_iter=episodes_per_iter,
                                      verbose=verbose)
    _, probs = pruner.apply_greedy(keep_threshold)
    rec["keep_probs"] = probs.tolist()
    rec["reverse_attention"] = pruner.reverse_attn.tolist()
    snap("after_apoptosis")

    if finetune_epochs:
        if verbose:
            print("\n== re-settling survivors ==", flush=True)
        finetune(model, loss_and_grad_fn, epochs=finetune_epochs,
                 loss_of=loss_of)
        snap("after_finetune")

    if rebirth_fraction > 0:
        dead = int((get_channel_mask(model) == 0.0).sum())
        born = neurogenesis(model, int(round(rebirth_fraction * dead)), seed=seed)
        if verbose:
            print(f"\n== neurogenesis: {born} channels reborn ==", flush=True)
        if born:
            snap("after_rebirth")
            if finetune_epochs:
                finetune(model, loss_and_grad_fn,
                         epochs=finetune_epochs, loss_of=loss_of)
                snap("after_rebirth_finetune")

    # ---- operating-point selection ----------------------------------------
    # Not "lowest loss" — that rule can never prune, because the unpruned
    # tower is usually the most accurate one and the whole exercise then
    # returns its own input. The question a database actually asks is: what
    # is the CHEAPEST tower I can serve without giving up measurable quality?
    # So: take the best loss seen, allow `select_tolerance` relative slack,
    # and among everything inside that band keep the fewest channels.
    floor = min(s[1] for s in saved)
    budget = floor * (1.0 + select_tolerance)
    eligible = [s for s in saved if s[1] <= budget]
    stage, loss, alive, params, mask = min(eligible, key=lambda s: (s[2], s[1]))
    _load_params(model, params)
    set_channel_mask(model, mask)
    mx.eval(model.parameters())
    if verbose:
        print(f"\n== operating point: {stage} — {alive} channels, val loss "
              f"{loss:.4f} (best seen {floor:.4f}, "
              f"budget +{select_tolerance:.0%}) ==", flush=True)
    rec["selected"] = {"stage": stage, "loss": loss, "channels": alive,
                       "best_loss_seen": floor,
                       "tolerance": select_tolerance,
                       "candidates": [{"stage": s[0], "loss": s[1],
                                       "channels": s[2]} for s in saved]}
    return rec