Remove src/src/models/mor_vit.py

src/src/models/mor_vit.py

@@ -1,386 +0,0 @@
-"""MoR-ViT: Mixture-of-Recursions on a pretrained DINOv3 ViT stem.
-
-Maps the llm-pipeline MoR recipe to a vision transformer:
-
-    entry block       -> pretrained DINOv3 stem (unique weights)
-    shared core       -> ``RecursiveAttentionBlock`` stack reused at every recursion
-    depth router      -> per-TOKEN ``TokenRouter`` (expert choice over patch tokens)
-    recursion emb     -> learned per-recursion embedding added before the core
-    token freeze      -> a patch token that stops routing keeps its current state
-
-Routing is per patch token (adaptive compute over the image), not per slice:
-every study keeps the full recursion budget, but only the "hard" patch tokens
-are recursed deeply.  Slice features are mean-pooled over patch tokens, then
-weighted by a ``SliceAttentionPool`` into a study vector.
-
-``features()`` exposes the frozen feature extractor used by linear probing
-without running the classification head.
-"""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-try:
-    import timm
-
-    HAS_TIMM = True
-except ImportError:  # pragma: no cover - optional dependency
-    HAS_TIMM = False
-
-
-class TokenScore(nn.Module):
-    """Per-token continuation score: token embedding -> scalar logit.
-
-    ``slice_gain`` (zero-initialised) couples the slice-level abnormality
-    score into the patch decision; the forward is exactly legacy at init.
-    """
-
-    def __init__(self, dim: int, hidden: int, init_bias: float = 0.0):
-        super().__init__()
-        self.in_proj = nn.Linear(dim, hidden)
-        self.out_proj = nn.Linear(hidden, 1)
-        self.slice_gain = nn.Parameter(torch.zeros(1))
-        with torch.no_grad():
-            self.out_proj.bias.fill_(init_bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.out_proj(F.silu(self.in_proj(x))).squeeze(-1)
-
-
-class SliceScore(nn.Module):
-    """Per-slice continuation score: pooled patch feature -> scalar logit.
-
-    Gives the patch router a slice-level view, so routing stays consistent
-    across the slices of a study (a finding shows on several adjacent slices).
-    """
-
-    def __init__(self, dim: int, hidden: int, init_bias: float = 0.0):
-        super().__init__()
-        self.in_proj = nn.Linear(dim, hidden)
-        self.out_proj = nn.Linear(hidden, 1)
-        with torch.no_grad():
-            self.out_proj.bias.fill_(init_bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.out_proj(F.silu(self.in_proj(x))).squeeze(-1)
-
-
-class TokenRouter(nn.Module):
-    """Expert-choice depth router over patch tokens.
-
-    One ``TokenScore`` head per recursion, plus a learned recursion embedding
-    so the router can condition on how deep a token already is.  The router
-    input can be spatially conditioned (stem patch position embedding) and
-    slice-conditioned (``SliceScore`` coupled through a zero-init gain).
-    Router biases initialise from the per-recursion capacity so the right
-    fraction of tokens continues from step 0.
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        hidden: int,
-        n_recursions: int,
-        capacities: list[float],
-        init_bias: float = 0.0,
-        warmup_steps: int = 0,
-        init_from_capacity: bool = True,
-    ):
-        super().__init__()
-        self.n_recursions = n_recursions
-        self.capacities = list(capacities)
-        self.warmup_steps = warmup_steps
-        self.heads = nn.ModuleList(
-            [TokenScore(dim, hidden, init_bias) for _ in range(n_recursions)]
-        )
-        self.rec_emb = nn.Parameter(torch.zeros(n_recursions, dim))
-        self.slice_score = SliceScore(dim, hidden)
-        if init_from_capacity:
-            for r, head in enumerate(self.heads):
-                cap = max(1e-3, min(1.0 - 1e-3, self.capacities[min(r, len(self.capacities) - 1)]))
-                with torch.no_grad():
-                    head.out_proj.bias.fill_(float(torch.logit(torch.tensor(cap))))
-
-    def capacity(self, r: int, step: int) -> float:
-        """Token fraction kept at recursion ``r``, ramping from 1.0 during warmup."""
-        target = self.capacities[r]
-        if self.warmup_steps > 0 and step < self.warmup_steps:
-            t = step / self.warmup_steps
-            return 1.0 - (1.0 - target) * t
-        return target
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        r: int,
-        patch_pos: torch.Tensor | None = None,
-        slice_feat: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        h = x + self.rec_emb[r]
-        if patch_pos is not None:
-            h = h + patch_pos
-        score = self.heads[r](h)
-        if slice_feat is not None:
-            score = score + self.heads[r].slice_gain * self.slice_score(slice_feat).unsqueeze(1)
-        return score
-
-
-class RecursiveAttentionBlock(nn.Module):
-    """Shared pre-norm attention block reused at every recursion.
-
-    Operates on the active token group only (the caller gathers/scatters), so
-    attention is plain unmasked full attention - no -inf masks, no NaN paths.
-    """
-
-    def __init__(self, dim: int, n_heads: int = 6, mlp_ratio: float = 4.0, dropout: float = 0.0):
-        super().__init__()
-        self.n_heads = n_heads
-        self.head_dim = dim // n_heads
-        self.q = nn.Linear(dim, dim)
-        self.k = nn.Linear(dim, dim)
-        self.v = nn.Linear(dim, dim)
-        self.proj = nn.Linear(dim, dim)
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-        self.mlp = nn.Sequential(
-            nn.Linear(dim, int(dim * mlp_ratio)),
-            nn.GELU(),
-            nn.Linear(int(dim * mlp_ratio), dim),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        S, T, D = x.shape
-        H = self.n_heads
-        xn = self.norm1(x)
-        q = self.q(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
-        k = self.k(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
-        v = self.v(xn).reshape(S, T, H, self.head_dim).transpose(1, 2)
-        attn = (q @ k.transpose(-1, -2) / (self.head_dim ** 0.5)).softmax(dim=-1)
-        out = (attn @ v).transpose(1, 2).reshape(S, T, D)
-        x = x + self.proj(out)
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-
-class SliceAttentionPool(nn.Module):
-    """Weight slices by learned relevance before aggregating into a study vector."""
-
-    def __init__(self, dim: int):
-        super().__init__()
-        self.query = nn.Parameter(torch.randn(dim))
-        self.scale = dim ** -0.5
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        w = F.softmax(x @ self.query * self.scale, dim=0)
-        return (w.unsqueeze(1) * x).sum(0)
-
-
-class MoRViT(nn.Module):
-    """Pretrained DINOv3 stem + MoR recursive core with per-token routing."""
-
-    def __init__(
-        self,
-        *,
-        stem_name: str = "vit_small_patch16_dinov3",
-        stem: nn.Module | None = None,
-        pretrained: bool = True,
-        freeze_stem: bool = True,
-        use_mor: bool = True,
-        dense: bool = False,
-        n_recursions: int = 3,
-        capacities: tuple[float, ...] = (1.0, 2.0 / 3.0, 1.0 / 3.0),
-        core_blocks: int = 2,
-        router_hidden: int = 128,
-        n_heads: int = 6,
-        n_classes: int = 12,
-        router_warmup_steps: int = 0,
-        router_init_bias: float = 0.0,
-    ):
-        super().__init__()
-        if stem is None:
-            if not HAS_TIMM:
-                raise ImportError("timm is required for a pretrained stem")
-            stem = timm.create_model(stem_name, pretrained=pretrained, num_classes=0)
-        self.n_recursions = n_recursions
-        self.capacities = list(capacities)
-        assert len(self.capacities) == n_recursions
-        self.use_mor = use_mor
-        self.dense = dense
-
-        self.stem = stem
-        stem_dim = self.stem.embed_dim
-        self.core_dim = stem_dim
-        if stem_dim != 384:
-            self.input_proj = nn.Linear(stem_dim, 384)
-            self.core_dim = 384
-        else:
-            self.input_proj = None
-        dim = self.core_dim
-        self.n_prefix = int(getattr(self.stem, "num_prefix_tokens", 1))
-
-        pe = getattr(self.stem, "pos_embed", None)
-        if isinstance(pe, torch.Tensor) and pe.dim() == 3 and pe.size(1) > self.n_prefix:
-            patch_pos = pe[:, self.n_prefix:].detach().clone()  # [1, P, stem_dim]
-            if self.input_proj is not None:
-                with torch.no_grad():
-                    patch_pos = self.input_proj(patch_pos)
-            self.register_buffer("patch_pos", patch_pos, persistent=False)
-        else:
-            self.patch_pos = None
-
-        if freeze_stem:
-            for p in self.stem.parameters():
-                p.requires_grad_(False)
-
-        self._soft_route_maps: list[torch.Tensor] = []
-
-        self.core = nn.ModuleList(
-            [RecursiveAttentionBlock(dim, n_heads=n_heads) for _ in range(core_blocks)]
-        )
-        self.router = (
-            TokenRouter(
-                dim,
-                router_hidden,
-                n_recursions,
-                self.capacities,
-                init_bias=router_init_bias,
-                warmup_steps=router_warmup_steps,
-            )
-            if not dense
-            else None
-        )
-        self.rec_emb = nn.Parameter(torch.zeros(n_recursions, dim))
-
-        self.exit_norm = nn.LayerNorm(dim)
-        self.pool = SliceAttentionPool(dim)
-        self.head = nn.Sequential(
-            nn.Linear(dim, dim * 2),
-            nn.GELU(),
-            nn.Linear(dim * 2, n_classes),
-        )
-
-    def _run_core(self, tokens: torch.Tensor, r: int) -> torch.Tensor:
-        x = tokens + self.rec_emb[r]
-        for block in self.core:
-            x = block(x)
-        return x
-
-    def _moR(self, tokens: torch.Tensor, step: int = 0) -> tuple[torch.Tensor, list[float]]:
-        """Run the recursion loop over a batch of token sequences.
-
-        Dense mode runs the shared core over every token at every recursion
-        (the same-budget control for adaptive routing).  Adaptive mode gathers
-        the top-k patch tokens (by router score) plus the prefix tokens,
-        processes them through the core, and scatters them back; non-selected
-        patch tokens keep their previous state.
-
-        The returned tokens are a soft mixture of the per-recursion states
-        weighted by the router's continuation probabilities, so the router
-        receives real gradient (hard top-k alone is non-differentiable).  The
-        soft routing maps are stored for the optional smoothness loss.
-        """
-        if self.dense:
-            for r in range(self.n_recursions):
-                tokens = self._run_core(tokens, r)
-            return tokens, [1.0] * self.n_recursions
-
-        S, T, D = tokens.shape
-        n_pref = self.n_prefix
-        P = T - n_pref
-        stats: list[float] = []
-        self._soft_route_maps = []
-
-        num = tokens.clone()
-        den = torch.ones(S, T, device=tokens.device)
-        w_pat = torch.ones(S, P, device=tokens.device)
-        pref_ones = torch.ones(S, n_pref, device=tokens.device)
-        pe = (
-            self.patch_pos
-            if (self.patch_pos is not None and self.patch_pos.shape[1] == P)
-            else None
-        )
-
-        for r in range(self.n_recursions - 1):
-            cap = self.router.capacity(r, step)
-            pat = tokens[:, n_pref:]
-            pat_scores = self.router(pat, r, patch_pos=pe, slice_feat=pat.mean(1))
-            p = torch.sigmoid(pat_scores)
-            self._soft_route_maps.append(p)
-            k = max(1, min(P, int(round(P * cap))))
-            sel = pat_scores.topk(k, dim=1).indices  # [S, k]
-            pref = torch.arange(n_pref, device=tokens.device).expand(S, n_pref)
-            idx = torch.cat([pref, sel + n_pref], dim=1)  # [S, M]
-            idx3 = idx.unsqueeze(-1).expand(S, idx.size(1), D)
-            stats.append((n_pref + k) / T)
-            gathered = tokens.gather(1, idx3)
-            updated = self._run_core(gathered, r)
-            tokens = tokens.scatter(1, idx3, updated)
-
-            w_pat = w_pat * p
-            w = torch.cat([pref_ones, w_pat], dim=1)
-            num = num + tokens * w.unsqueeze(-1)
-            den = den + w
-
-        tokens = self._run_core(tokens, self.n_recursions - 1)
-        stats.append(1.0)
-        num = num + tokens * w.unsqueeze(-1)
-        den = den + w
-        mixed = num / den.unsqueeze(-1)
-        return mixed, stats
-
-    def routing_smoothness_loss(self) -> torch.Tensor | float:
-        """Total-variation penalty on the soft routing maps (spatial coherence).
-
-        Knee findings occupy contiguous regions; a scattered routing map is a
-        bug signal.  Add ``weight * this`` to the training loss.
-        """
-        if not self._soft_route_maps:
-            return 0.0
-        total = None
-        for p in self._soft_route_maps:  # [S, P]
-            P = p.size(1)
-            h = int(round(P ** 0.5))
-            if h * h != P:
-                continue
-            g = p.reshape(p.size(0), h, h)
-            tv = (g[:, 1:, :] - g[:, :-1, :]).abs().mean() + (
-                g[:, :, 1:] - g[:, :, :-1]
-            ).abs().mean()
-            total = tv if total is None else total + tv
-        if total is None:
-            return 0.0
-        return total / len(self._soft_route_maps)
-
-    def features(
-        self, x: torch.Tensor, step: int = 0, pool: str = "patchmean"
-    ) -> torch.Tensor:
-        """Per-slice features (no head): [S, D] from [S, 3, H, W] slices."""
-        tokens = self.stem.forward_features(x)  # [S, T, D]
-        return self.features_from_tokens(tokens, step, pool)
-
-    def features_from_tokens(
-        self, tokens: torch.Tensor, step: int = 0, pool: str = "patchmean"
-    ) -> torch.Tensor:
-        """Features from cached stem tokens (feat cache path): [S, D]."""
-        tokens = tokens.float()  # _moR internals are f32; bf16 cache upcasts exactly
-        if self.input_proj is not None:
-            with torch.autocast("cuda", enabled=False):
-                tokens = self.input_proj(tokens).float()  # keep f32; autocast would make it bf16
-        if self.use_mor or self.dense:
-            tokens, _ = self._moR(tokens, step)
-        tokens = self.exit_norm(tokens)
-        if pool == "cls":
-            return tokens[:, 0]
-        return tokens[:, self.n_prefix :].mean(1)
-
-    def forward(
-        self, x: torch.Tensor, step: int = 0
-    ) -> tuple[torch.Tensor, list[float]]:
-        """Forward one study: [S, 3, H, W] slices -> (logits [n_classes], route_stats)."""
-        slice_feats = self.features(x, step)  # [S, D]
-        study = self.pool(slice_feats)  # [D]
-        logits = self.head(study)
-        return logits, []