Add core src/ package for direct model instantiation

src/src/data/plant.py

@@ -0,0 +1,75 @@
+"""Plant CSV dataset for data/wa_plants/manifest.csv (observation-separated).
+Reuses RSNA train_mor pattern: manifest + ImageNet norm, 336px, Aug.
+"""
+from __future__ import annotations
+import csv
+from pathlib import Path
+import random
+from PIL import Image
+import warnings
+# faster decode: raise limit (100Mpx originals) + draft shrink for large JPEGs
+Image.MAX_IMAGE_PIXELS = 300_000_000
+warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning)
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms as T
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+class PlantDataset(Dataset):
+    def __init__(self, manifest: Path, split: str, img_size: int = 336, augment: bool = False):
+        self.split = split
+        self.img_size = img_size
+        rows = []
+        # species -> idx map built from manifest train split (500 classes)
+        with open(manifest, newline="", encoding="utf-8") as f:
+            for r in csv.DictReader(f):
+                if r["split"] == split and r["status"] in ("downloaded","skip_exists"):
+                    rows.append(r)
+        # stable class order sorted
+        species = sorted({r["species"] for r in rows})
+        self.species_to_idx = {s:i for i,s in enumerate(species)}
+        self.rows = rows
+        # class counts for balanced sampling
+        self.augment = augment
+        tfms = []
+        if augment:
+            tfms = [
+                T.RandomResizedCrop(img_size, scale=(0.7,1.0)),
+                T.RandomHorizontalFlip(),
+                T.ColorJitter(0.2,0.2,0.2,0.05),
+                T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
+            ]
+        else:
+            tfms = [T.Resize(int(img_size*1.14)), T.CenterCrop(img_size), T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)]
+        self.tf = T.Compose(tfms)
+
+    def __len__(self): return len(self.rows)
+    def __getitem__(self, i):
+        r = self.rows[i]
+        p = Path(r["path"])
+        # manifest stores absolute win path; if not found try relative to data/wa_plants
+        if not p.exists():
+            # try finding under train/val subfolders by gbifID
+            base = Path(__file__).resolve().parents[2] / "data" / "wa_plants"
+            for split in ("train","val"):
+                cand = base / split / r["species"].replace(" ","_").replace("/","_")[:120] / f"{r['gbifID']}.jpg"
+                if cand.exists():
+                    p = cand; break
+        try:
+            im = Image.open(p)
+            # fast draft for huge JPEGs (8× shrink before full decode) - no effect on small images
+            try:
+                if max(im.size) > 1024:
+                    # draft uses libjpeg shrink 1/2/4/8
+                    im.draft("RGB", (768, 768))
+            except:
+                pass
+            im = im.convert("RGB")
+        except Exception:
+            im = Image.new("RGB", (self.img_size, self.img_size))
+        # species_to_idx may have been remapped after init (train vs val) - lookup safely
+        y = self.species_to_idx.get(r["species"], 0)
+        x = self.tf(im)
+        return x, y, r["gbifID"]

src/src/data/taxonomy.py

@@ -0,0 +1,150 @@
+"""
+Taxonomic Hierarchy Engine for Western Australia Flora.
+Builds and maps Species -> Genus -> Family for hierarchical multi-task classification.
+"""
+import pandas as pd
+from pathlib import Path
+
+# Common WA Flora Genus -> Family mapping dictionary
+GENUS_TO_FAMILY = {
+    # Proteaceae
+    "Banksia": "Proteaceae", "Grevillea": "Proteaceae", "Hakea": "Proteaceae",
+    "Dryandra": "Proteaceae", "Adenanthos": "Proteaceae", "Isopogon": "Proteaceae",
+    "Petrophile": "Proteaceae", "Synaphea": "Proteaceae", "Conospermum": "Proteaceae",
+    "Lambertia": "Proteaceae", "Persoonia": "Proteaceae", "Xylomelum": "Proteaceae",
+    "Franklandia": "Proteaceae", "Stirlingia": "Proteaceae",
+    
+    # Myrtaceae
+    "Eucalyptus": "Myrtaceae", "Corymbia": "Myrtaceae", "Melaleuca": "Myrtaceae",
+    "Verticordia": "Myrtaceae", "Calothamnus": "Myrtaceae", "Darwinia": "Myrtaceae",
+    "Chamelaucium": "Myrtaceae", "Thryptomene": "Myrtaceae", "Beaufortia": "Myrtaceae",
+    "Baeckea": "Myrtaceae", "Regelia": "Myrtaceae", "Calytrix": "Myrtaceae",
+    "Eremaea": "Myrtaceae", "Kunzea": "Myrtaceae", "Micromyrtus": "Myrtaceae",
+    "Hypocalymma": "Myrtaceae", "Scholtzia": "Myrtaceae", "Phymatocarpus": "Myrtaceae",
+    "Leptospermum": "Myrtaceae", "Astartea": "Myrtaceae", "Taxandria": "Myrtaceae",
+
+    # Fabaceae (Legumes)
+    "Acacia": "Fabaceae", "Gastrolobium": "Fabaceae", "Daviesia": "Fabaceae",
+    "Jacksonia": "Fabaceae", "Bossiaea": "Fabaceae", "Hovea": "Fabaceae",
+    "Kennedia": "Fabaceae", "Hardenbergia": "Fabaceae", "Gompholobium": "Fabaceae",
+    "Mirbelia": "Fabaceae", "Chorizema": "Fabaceae", "Pultenaea": "Fabaceae",
+    "Isotropis": "Fabaceae", "Swainsona": "Fabaceae", "Senna": "Fabaceae",
+    "Indigofera": "Fabaceae", "Templetonia": "Fabaceae", "Latrobea": "Fabaceae",
+
+    # Orchidaceae (Orchids)
+    "Caladenia": "Orchidaceae", "Thelymitra": "Orchidaceae", "Diuris": "Orchidaceae",
+    "Pterostylis": "Orchidaceae", "Drakaea": "Orchidaceae", "Microtis": "Orchidaceae",
+    "Prasophyllum": "Orchidaceae", "Cyanicula": "Orchidaceae", "Ericksonella": "Orchidaceae",
+    "Leporella": "Orchidaceae", "Paracaleana": "Orchidaceae", "Pyrorchis": "Orchidaceae",
+    "Cryptostylis": "Orchidaceae", "Elythranthera": "Orchidaceae", "Spiculaea": "Orchidaceae",
+
+    # Haemodoraceae (Kangaroo Paws)
+    "Anigozanthos": "Haemodoraceae", "Macropidia": "Haemodoraceae", "Conostylis": "Haemodoraceae",
+    "Tribonanthes": "Haemodoraceae", "Haemodorum": "Haemodoraceae", "Phlebocarya": "Haemodoraceae",
+
+    # Ericaceae / Epacridaceae (Heaths)
+    "Leucopogon": "Ericaceae", "Astroloma": "Ericaceae", "Lysinema": "Ericaceae",
+    "Conostephium": "Ericaceae", "Andersonia": "Ericaceae", "Sphenotoma": "Ericaceae",
+    "Styphelia": "Ericaceae", "Cosmelia": "Ericaceae",
+
+    # Goodeniaceae
+    "Goodenia": "Goodeniaceae", "Dampiera": "Goodeniaceae", "Scaevola": "Goodeniaceae",
+    "Lechenaultia": "Goodeniaceae", "Velleia": "Goodeniaceae", "Anthotium": "Goodeniaceae",
+
+    # Stylidiaceae (Triggerplants)
+    "Stylidium": "Stylidiaceae", "Levenhookia": "Stylidiaceae",
+
+    # Droseraceae (Sundews / Carnivorous)
+    "Drosera": "Droseraceae",
+
+    # Asteraceae (Daisies)
+    "Rhodanthe": "Asteraceae", "Waitzia": "Asteraceae", "Schoenia": "Asteraceae",
+    "Podolepis": "Asteraceae", "Olearia": "Asteraceae", "Senecio": "Asteraceae",
+    "Brachyscome": "Asteraceae", "Lawrencella": "Asteraceae", "Xerochrysum": "Asteraceae",
+
+    # Xanthorrhoeaceae / Asphodelaceae
+    "Xanthorrhoea": "Asphodelaceae",
+
+    # Dasypogonaceae
+    "Kingia": "Dasypogonaceae", "Dasypogon": "Dasypogonaceae", "Calectasia": "Dasypogonaceae",
+
+    # Casuarinaceae (Sheoaks)
+    "Allocasuarina": "Casuarinaceae", "Casuarina": "Casuarinaceae",
+
+    # Pittosporaceae
+    "Marianthus": "Pittosporaceae", "Billardiera": "Pittosporaceae", "Pittosporum": "Pittosporaceae",
+    "Cheiranthera": "Pittosporaceae",
+
+    # Restionaceae (Rushes / Sedges)
+    "Desmocladus": "Restionaceae", "Alexgeorgea": "Restionaceae", "Loxocarya": "Restionaceae",
+    "Chordifex": "Restionaceae", "Lyginia": "Restionaceae",
+
+    # Chenopodiaceae / Amaranthaceae
+    "Atriplex": "Amaranthaceae", "Maireana": "Amaranthaceae", "Tecticornia": "Amaranthaceae",
+    "Ptilotus": "Amaranthaceae",
+
+    # Malvaceae
+    "Alyogyne": "Malvaceae", "Thomasia": "Malvaceae", "Guichenotia": "Malvaceae",
+    "Lasiopetalum": "Malvaceae", "Seringia": "Malvaceae", "Androcalva": "Malvaceae",
+
+    # Rutaceae (Boronias / Waxflowers)
+    "Boronia": "Rutaceae", "Philotheca": "Rutaceae", "Diplolaena": "Rutaceae",
+    "Correa": "Rutaceae", "Crowea": "Rutaceae", "Geleznowia": "Rutaceae",
+    "Asterolasia": "Rutaceae", "Eriostemon": "Rutaceae",
+
+    # Loranthaceae (Mistletoes / Christmas Tree)
+    "Nuytsia": "Loranthaceae", "Amyema": "Loranthaceae", "Lysiana": "Loranthaceae",
+}
+
+class TaxonomyHierarchy:
+    """Manages species -> genus -> family indices and mappings."""
+    def __init__(self, species_list: list[str]):
+        self.species_list = sorted(list(set(s for s in species_list if isinstance(s, str) and s.strip())))
+        
+        # 1. Species Index
+        self.species_to_idx = {s: i for i, s in enumerate(self.species_list)}
+        self.idx_to_species = {i: s for s, i in self.species_to_idx.items()}
+        self.n_species = len(self.species_list)
+
+        # 2. Genus Index
+        self.species_to_genus = {}
+        genera_set = set()
+        for s in self.species_list:
+            genus = s.split()[0] if s else "Unknown"
+            self.species_to_genus[s] = genus
+            genera_set.add(genus)
+        
+        self.genus_list = sorted(list(genera_set))
+        self.genus_to_idx = {g: i for i, g in enumerate(self.genus_list)}
+        self.idx_to_genus = {i: g for g, i in self.genus_to_idx.items()}
+        self.n_genera = len(self.genus_list)
+
+        # 3. Family Index
+        self.genus_to_family = {}
+        families_set = set()
+        for g in self.genus_list:
+            fam = GENUS_TO_FAMILY.get(g, f"{g}_Family")
+            self.genus_to_family[g] = fam
+            families_set.add(fam)
+
+        self.family_list = sorted(list(families_set))
+        self.family_to_idx = {f: i for i, f in enumerate(self.family_list)}
+        self.idx_to_family = {i: f for f, i in self.family_to_idx.items()}
+        self.n_families = len(self.family_list)
+
+        # Precompute lookup tables
+        self.species_idx_to_genus_idx = [
+            self.genus_to_idx[self.species_to_genus[self.idx_to_species[i]]]
+            for i in range(self.n_species)
+        ]
+        self.species_idx_to_family_idx = [
+            self.family_to_idx[self.genus_to_family[self.species_to_genus[self.idx_to_species[i]]]]
+            for i in range(self.n_species)
+        ]
+
+    def get_taxonomy_indices(self, species_name: str) -> tuple[int, int, int]:
+        """Returns (species_idx, genus_idx, family_idx)."""
+        s_idx = self.species_to_idx.get(species_name, -1)
+        if s_idx == -1:
+            return -1, -1, -1
+        return s_idx, self.species_idx_to_genus_idx[s_idx], self.species_idx_to_family_idx[s_idx]

src/src/models/moe_swiglu_v2.py

@@ -0,0 +1,136 @@
+"""
+Streamlined 16-Expert SwiGLU Mixture-of-Experts (MoE) Layer for Vision Transformers.
+Uses pure SwiGLU FFN experts, Gating Residuals (cross-layer routing memory),
+and Context-Aware Global Projections.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SwiGLU(nn.Module):
+    """3-Matrix SwiGLU Feed-Forward Network: SiLU(W1 * x) * (W3 * x) @ W2"""
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_ff, d_model, bias=False)
+        self.w3 = nn.Linear(d_model, d_ff, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class MoESwiGLUv2(nn.Module):
+    """
+    Pure SwiGLU MoE FFN layer.
+    - num_ffn: Number of SwiGLU experts (default 16)
+    - top_k: Active experts per token (default 2)
+    - gating_residual: Linear projection of previous-layer gate logits
+    - gate_ctx: Global CLS context projection into gate logits
+    """
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int = 512,
+        num_ffn: int = 16,
+        top_k: int = 2,
+        gating_residual: bool = True,
+        gate_ctx: bool = True,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.num_ffn = num_ffn
+        self.top_k = min(top_k, num_ffn)
+        self.gate_ctx = gate_ctx
+
+        # Router linear layer: [d_model -> num_ffn]
+        self.router = nn.Linear(d_model, num_ffn, bias=False)
+        self.experts = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(num_ffn)])
+
+        # Cross-layer gating residual projection
+        self.gating_residual = None
+        if gating_residual:
+            self.gating_residual = nn.Linear(num_ffn, num_ffn, bias=False)
+            nn.init.zeros_(self.gating_residual.weight)
+
+        # Context-aware global projection
+        self.ctx_proj = None
+        if gate_ctx:
+            self.ctx_proj = nn.Linear(d_model, num_ffn, bias=False)
+            nn.init.zeros_(self.ctx_proj.weight)
+
+    def forward(self, x: torch.Tensor, prev_gate: torch.Tensor = None, ctx: torch.Tensor = None):
+        """
+        x: [B, T, D] Token embeddings
+        prev_gate: [B, T, num_ffn] Gate logits from previous layer (optional)
+        ctx: [B, D] Global CLS context embedding (optional)
+        Returns:
+          output: [B, T, D]
+          aux_loss: Scalar load balancing loss
+          gate_logits: [B, T, num_ffn]
+        """
+        B, T, D = x.shape
+        x_2d = x.reshape(-1, D)
+        N = x_2d.size(0)
+
+        # 1. Compute Gate Logits
+        gate_logits = self.router(x_2d) # [N, num_ffn]
+
+        if self.ctx_proj is not None and ctx is not None:
+            # Broadcast global CLS context across all T tokens
+            gate_logits = gate_logits + self.ctx_proj(ctx).repeat_interleave(T, dim=0)
+
+        if self.gating_residual is not None and prev_gate is not None:
+            prev_flat = prev_gate.reshape(-1, self.num_ffn)
+            if prev_flat.size(0) == N:
+                gate_logits = gate_logits + self.gating_residual(prev_flat)
+
+        # 2. Softmax & Top-K Selection
+        gate_weights = F.softmax(gate_logits, dim=-1)
+        top_weights, top_indices = gate_weights.topk(self.top_k, dim=-1)
+        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)
+        route_weights = top_weights.to(dtype=x_2d.dtype)
+
+        # 3. Flatten tokens and sort by expert index for batch GEMM efficiency
+        flat_expert = top_indices.reshape(-1)
+        flat_token = (
+            torch.arange(N, device=x.device)
+            .unsqueeze(1)
+            .expand(-1, self.top_k)
+            .reshape(-1)
+        )
+        flat_weight = route_weights.reshape(-1, 1)
+
+        order = torch.argsort(flat_expert, stable=True)
+        sorted_tokens = flat_token[order]
+        sorted_weights = flat_weight[order]
+        sorted_x = x_2d.index_select(0, sorted_tokens)
+
+        counts = torch.bincount(flat_expert, minlength=self.num_ffn).to(torch.int64)
+        offsets = torch.cat(
+            [torch.zeros(1, dtype=torch.int64, device=x.device),
+             torch.cumsum(counts, 0)[:-1]]
+        )
+
+        output = torch.zeros_like(x_2d)
+
+        # 4. Dispatch each expert
+        for expert_idx in range(self.num_ffn):
+            lo = int(offsets[expert_idx])
+            hi = lo + int(counts[expert_idx])
+            if lo == hi:
+                continue
+            tids = sorted_tokens[lo:hi]
+            w = sorted_weights[lo:hi]
+            expert_out = self.experts[expert_idx](sorted_x[lo:hi]) * w
+            output.index_add_(0, tids, expert_out.to(dtype=output.dtype))
+
+        # 5. Load-Balancing Auxiliary Loss
+        fractions = counts.float() / counts.sum().clamp(min=1e-12).float()
+        avg_prob = gate_weights.mean(dim=0)
+        aux_loss = self.num_ffn * (fractions.to(dtype=gate_weights.dtype) * avg_prob).sum()
+
+        return (
+            output.view(B, T, D),
+            aux_loss,
+            gate_logits.view(B, T, self.num_ffn)
+        )

src/src/models/moe_vit.py

@@ -0,0 +1,452 @@
+"""MoE-ViT: MoE++ (zero/copy/constant experts) on a pretrained DINOv3 ViT stem.
+
+Replaces the MoR recursive core with a plain stacked-transformer core whose
+FFNs are MoE++ layers (arXiv:2410.07348):
+
+    dense MLP (d_ff=1536)  ->  4 FFN experts (d_ff=384) + 1 zero + 1 copy
+                                + 2 constant experts, top_k=2 per token
+
+Routing is per patch token, so easy image regions land on zero/copy/constant
+specialists (near-zero compute) while hard regions route to real FFN experts -
+adaptive compute at the finest granularity, with no recursion machinery.
+
+The frozen stem and the study-level pooling/head mirror MoR-ViT so the two
+models share the same ``features()`` / ``forward()`` contract and can be
+compared on identical benchmarks.
+"""
+
+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 SwiGLU(nn.Module):
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_ff, d_model, bias=False)
+        self.w3 = nn.Linear(d_model, d_ff, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class ConstantExpert(nn.Module):
+    """MoE++ constant expert: alpha1*x + alpha2*v (negligible params)."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.v = nn.Parameter(torch.zeros(d_model))
+        self.wc = nn.Linear(d_model, 2, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = F.softmax(self.wc(x), dim=-1)
+        return (alpha[..., 0:1] * x + alpha[..., 1:2] * self.v.to(dtype=x.dtype)).to(dtype=x.dtype)
+
+
+class MoEFFN(nn.Module):
+    """Sparse MoE FFN with MoE++ zero / copy / constant experts.
+
+    Index layout: [FFN | zero | copy | constant].  Routing is softmax top-k
+    with renormalisation; only tokens selected for a real FFN expert are
+    computed (``index_select``), so zero-compute routing is real skipping.
+    Returns the load-balance aux loss and pre-softmax gate logits.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int,
+        num_ffn: int = 4,
+        top_k: int = 2,
+        n_zero: int = 1,
+        n_copy: int = 1,
+        n_const: int = 2,
+        tau: float = 0.75,
+        gating_residual: bool = True,
+        gate_ctx: bool = False,
+    ):
+        super().__init__()
+        self.num_ffn = num_ffn
+        self.top_k = min(top_k, num_ffn + n_zero + n_copy + n_const)
+        self.n_zero = max(0, n_zero)
+        self.n_copy = max(0, n_copy)
+        self.n_const = max(0, n_const)
+        self.n_zc = self.n_zero + self.n_copy + self.n_const
+        self.total_experts = num_ffn + self.n_zc
+        self.last_counts: torch.Tensor | None = None
+        self.tau = float(tau)
+        self.gate_ctx = gate_ctx
+
+        self.copy_start = num_ffn + self.n_zero
+        self.const_start = self.copy_start + self.n_copy
+
+        self.router = nn.Linear(d_model, self.total_experts, bias=False)
+        self.experts = nn.ModuleList([SwiGLU(d_model, d_ff) for _ in range(num_ffn)])
+        self.const_experts = nn.ModuleList([ConstantExpert(d_model) for _ in range(self.n_const)])
+
+        self.gating_residual = None
+        if gating_residual:
+            self.gating_residual = nn.Linear(self.total_experts, self.total_experts, bias=False)
+            nn.init.zeros_(self.gating_residual.weight)
+
+        self.ctx_proj = None
+        if gate_ctx:
+            self.ctx_proj = nn.Linear(d_model, self.total_experts, bias=False)
+            nn.init.zeros_(self.ctx_proj.weight)
+
+        eta = torch.ones(self.total_experts)
+        if self.n_zc > 0:
+            eta[num_ffn:] = self.tau
+        self.register_buffer("eta", eta, persistent=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        prev_gate: torch.Tensor | None = None,
+        ctx: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Returns (output, aux_loss, gate_logits, gate_weights)."""
+        B, T, D = x.shape
+        x_2d = x.reshape(-1, D)
+        N = x_2d.size(0)
+
+        gate_logits = self.router(x_2d)
+        if self.ctx_proj is not None and ctx is not None:
+            gate_logits = gate_logits + self.ctx_proj(ctx).repeat_interleave(T, dim=0)
+        if self.gating_residual is not None and prev_gate is not None:
+            if prev_gate.dim() == 3:
+                prev_gate = prev_gate.reshape(-1, prev_gate.size(-1))
+            if prev_gate.size(0) == N and prev_gate.size(-1) == self.total_experts:
+                gate_logits = gate_logits + self.gating_residual(prev_gate)
+
+        gate_weights = F.softmax(gate_logits, dim=-1)
+        top_weights, top_indices = gate_weights.topk(self.top_k, dim=-1)
+        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True).clamp(min=1e-9)
+        route_weights = top_weights.to(dtype=x_2d.dtype)
+
+        flat_expert = top_indices.reshape(-1)
+        flat_token = (
+            torch.arange(N, device=x.device)
+            .unsqueeze(1)
+            .expand(-1, self.top_k)
+            .reshape(-1)
+        )
+        flat_weight = route_weights.reshape(-1, 1)
+
+        order = torch.argsort(flat_expert, stable=True)
+        sorted_tokens = flat_token[order]
+        sorted_weights = flat_weight[order]
+        sorted_x = x_2d.index_select(0, sorted_tokens)
+        counts = torch.bincount(
+            flat_expert, minlength=self.total_experts
+        ).to(torch.int64)
+        offsets = torch.cat(
+            [torch.zeros(1, dtype=torch.int64, device=x.device),
+             torch.cumsum(counts, 0)[:-1]]
+        )
+
+        output = torch.zeros_like(x_2d)
+
+        for expert_idx in range(self.num_ffn):
+            lo = int(offsets[expert_idx])
+            hi = lo + int(counts[expert_idx])
+            if lo == hi:
+                continue
+            tids = sorted_tokens[lo:hi]
+            w = sorted_weights[lo:hi]
+            output.index_add_(
+                0, tids,
+                (self.experts[expert_idx](sorted_x[lo:hi]) * w).to(dtype=output.dtype),
+            )
+
+        if self.n_zc > 0:
+            for i in range(self.n_copy):
+                expert_idx = self.copy_start + i
+                lo = int(offsets[expert_idx])
+                hi = lo + int(counts[expert_idx])
+                if lo == hi:
+                    continue
+                tids = sorted_tokens[lo:hi]
+                w = sorted_weights[lo:hi]
+                output.index_add_(
+                    0, tids, (sorted_x[lo:hi] * w).to(dtype=output.dtype)
+                )
+
+            for i in range(self.n_const):
+                expert_idx = self.const_start + i
+                lo = int(offsets[expert_idx])
+                hi = lo + int(counts[expert_idx])
+                if lo == hi:
+                    continue
+                tids = sorted_tokens[lo:hi]
+                w = sorted_weights[lo:hi]
+                output.index_add_(
+                    0, tids,
+                    (self.const_experts[i](sorted_x[lo:hi]) * w).to(dtype=output.dtype),
+                )
+
+        counts = torch.zeros(self.total_experts, device=x.device, dtype=gate_weights.dtype)
+        counts.scatter_add_(0, flat_expert, torch.ones_like(flat_expert, dtype=gate_weights.dtype))
+        self.last_counts = counts.detach()
+        fractions = counts / counts.sum().clamp(min=1e-12)
+        avg_prob = gate_weights.mean(dim=0)
+        if self.n_zc > 0:
+            aux_loss = self.total_experts * (self.eta.to(dtype=gate_weights.dtype) * fractions * avg_prob).sum()
+        else:
+            aux_loss = self.total_experts * (fractions * avg_prob).sum()
+
+        return (output.view(B, T, D), aux_loss, gate_logits,
+                gate_weights.view(B, T, self.total_experts))
+
+
+class MoEAttentionBlock(nn.Module):
+    """Pre-norm attention + MoE++ FFN.  Stores the FFN gate maps for the
+    spatial-smoothness (TV) loss."""
+
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int = 6,
+        d_ff: int = 384,
+        num_ffn: int = 4,
+        top_k: int = 2,
+        n_zero: int = 1,
+        n_copy: int = 1,
+        n_const: int = 2,
+        tau: float = 0.75,
+        gating_residual: bool = True,
+        gate_ctx: bool = False,
+    ):
+        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.moe = MoEFFN(
+            dim, d_ff, num_ffn=num_ffn, top_k=top_k,
+            n_zero=n_zero, n_copy=n_copy, n_const=n_const,
+            tau=tau, gating_residual=gating_residual, gate_ctx=gate_ctx,
+        )
+
+    def forward(
+        self, x: torch.Tensor, prev_gate: torch.Tensor | None = None,
+        ctx: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, 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)
+
+        y, aux, gate_logits, gate_weights = self.moe(
+            self.norm2(x), prev_gate=prev_gate, ctx=ctx)
+        x = x + y
+        return x, aux, gate_logits, gate_weights
+
+
+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 PerTargetPool(nn.Module):
+    """Per-finding attention over slices: each of n_classes findings owns a
+    query that weights the slices most relevant to it -> [n_classes, D]."""
+
+    def __init__(self, dim: int, n_classes: int):
+        super().__init__()
+        self.queries = nn.Parameter(torch.randn(n_classes, dim) * dim ** -0.5)
+        self.scale = dim ** -0.5
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [S, D] slice vectors -> [n_classes, S] attention weights
+        w = F.softmax(x @ self.queries.T * self.scale, dim=0)
+        return (x.unsqueeze(0) * w.T.unsqueeze(-1)).sum(1)
+
+
+class PerTargetHead(nn.Module):
+    """Shared MLP then per-target projection (12 findings, each its own weight)."""
+
+    def __init__(self, dim: int, n_classes: int, hidden: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU())
+        self.w = nn.Parameter(torch.randn(n_classes, hidden) * hidden ** -0.5)
+        self.b = nn.Parameter(torch.zeros(n_classes))
+
+    def forward(self, ctx: torch.Tensor) -> torch.Tensor:
+        # ctx: [n_classes, D] -> logits [n_classes]
+        h = self.mlp(ctx)
+        return (h * self.w).sum(-1) + self.b
+
+
+class MoEViT(nn.Module):
+    """Pretrained DINOv3 stem + stacked MoE++ 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,
+        core_blocks: int = 2,
+        n_heads: int = 6,
+        d_ff: int = 384,
+        num_ffn: int = 4,
+        top_k: int = 2,
+        n_zero: int = 1,
+        n_copy: int = 1,
+        n_const: int = 2,
+        tau: float = 0.75,
+        gating_residual: bool = True,
+        n_classes: int = 12,
+        gate_ctx: bool = False,
+        per_target: bool = False,
+    ):
+        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.num_ffn = num_ffn
+        self.top_k = top_k
+        self.gate_ctx = gate_ctx
+        self.per_target = per_target
+
+        self.stem = stem
+        stem_dim = self.stem.embed_dim
+        self.core_dim = 384
+        self.input_proj = nn.Linear(stem_dim, self.core_dim) if stem_dim != 384 else None
+        dim = self.core_dim
+        self.n_prefix = int(getattr(self.stem, "num_prefix_tokens", 1))
+
+        if freeze_stem:
+            for p in self.stem.parameters():
+                p.requires_grad_(False)
+
+        self._ffn_maps: list[torch.Tensor] = []
+
+        self.blocks = nn.ModuleList(
+            [
+                MoEAttentionBlock(
+                    dim, n_heads=n_heads, d_ff=d_ff, num_ffn=num_ffn, top_k=top_k,
+                    n_zero=n_zero, n_copy=n_copy, n_const=n_const,
+                    tau=tau, gating_residual=gating_residual, gate_ctx=gate_ctx,
+                )
+                for _ in range(core_blocks)
+            ]
+        )
+
+        self.exit_norm = nn.LayerNorm(dim)
+        if per_target:
+            self.pool = PerTargetPool(dim, n_classes)
+            self.head = PerTargetHead(dim, n_classes)
+        else:
+            self.pool = SliceAttentionPool(dim)
+            self.head = nn.Sequential(
+                nn.Linear(dim, dim * 2),
+                nn.GELU(),
+                nn.Linear(dim * 2, n_classes),
+            )
+
+    def moe_aux_loss(self) -> torch.Tensor:
+        total = None
+        for b in self.blocks:
+            if not hasattr(b, "_aux"):
+                continue
+            total = b._aux if total is None else total + b._aux
+        if total is None:
+            return torch.tensor(0.0, device=next(self.parameters()).device)
+        return total / len(self.blocks)
+
+    def routing_tv_loss(self) -> torch.Tensor:
+        """Total-variation penalty on the per-patch FFN-mass maps.
+
+        Knee findings occupy contiguous regions; scattered routing is a bug
+        signal.  The map is the probability mass routed to real FFN experts,
+        reshaped to the patch grid.
+        """
+        if not self._ffn_maps:
+            return torch.tensor(0.0, device=next(self.parameters()).device)
+        total = None
+        for p in self._ffn_maps:
+            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 torch.tensor(0.0, device=next(self.parameters()).device)
+        return total / len(self._ffn_maps)
+
+    def features(
+        self, x: torch.Tensor, 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, pool)
+
+    def features_from_tokens(
+        self, tokens: torch.Tensor, pool: str = "patchmean"
+    ) -> torch.Tensor:
+        """Features from cached stem tokens (feat cache path): [S, D]."""
+        tokens = tokens.float()  # core 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
+        self._ffn_maps = []
+        prev_gate = None
+        ctx = tokens.mean(dim=1) if self.gate_ctx else None  # [B, D] per-slice context
+        for b in self.blocks:
+            tokens, aux, gate_logits, gate_weights = b(
+                tokens, prev_gate=prev_gate, ctx=ctx)
+            b._aux = aux
+            self._ffn_maps.append(
+                gate_weights[:, self.n_prefix:, : self.num_ffn].sum(dim=-1)
+            )
+            prev_gate = gate_logits
+        tokens = self.exit_norm(tokens)
+        if pool == "cls":
+            return tokens[:, 0]
+        return tokens[:, self.n_prefix:].mean(1)
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, list[float]]:
+        """Forward one study: [S, 3, H, W] slices -> (logits [n_classes], aux_stats)."""
+        slice_feats = self.features(x)  # [S, D]
+        study = self.pool(slice_feats)  # [D] or [n_classes, D]
+        logits = self.head(study)
+        if logits.dim() == 2:
+            logits = logits.squeeze(-1)
+        return logits, [self.moe_aux_loss().item()]

src/src/models/mor_cnn.py

@@ -0,0 +1,239 @@
+"""MoR-CNN: Mixture-of-Recursions applied to a convolutional vision model.
+
+Ports the LLM MoR idea (llm-pipeline) to 2D medical imaging.  The mapping:
+
+    entry block        -> pretrained CNN stem (unique weights)
+    shared recursive   -> ``RecursiveConvBlock`` stack reused at every recursion
+    core
+    depth router       -> per-slice ``DepthRouter`` (expert choice over slices)
+    recursion emb      -> learned per-recursion embedding added before the core
+    slice freeze       -> a slice that stops routing keeps its current state
+
+At inference the router spends full recursion depth only on the "hard" slices
+(those most likely to be abnormal), so FLOPs scale with content — the lever the
+competition's Efficiency Track scores.
+"""
+
+from __future__ import annotations
+
+import math
+
+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 DepthScore(nn.Module):
+    """Per-slice continuation score: pooled feature -> scalar logit.
+
+    Mirrors ``_DepthScore`` in llm-pipeline's MoR router, but over a global
+    pooled per-slice feature instead of a per-token hidden state.
+    """
+
+    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 DepthRouter(nn.Module):
+    """Expert-choice depth router over slices.
+
+    One ``DepthScore`` head per recursion.  Each recursion also adds a learned
+    recursion embedding so the router can condition on how deep a slice
+    already is (the vision analogue of ``rec_emb`` in the LLM MoR).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        hidden: int,
+        n_recursions: int,
+        capacities: list[float],
+        init_bias: float = 0.0,
+        warmup_steps: int = 0,
+    ):
+        super().__init__()
+        self.n_recursions = n_recursions
+        self.capacities = capacities
+        self.warmup_steps = warmup_steps
+        self.heads = nn.ModuleList(
+            [DepthScore(dim, hidden, init_bias) for _ in range(n_recursions)]
+        )
+        self.rec_emb = nn.Parameter(torch.zeros(n_recursions, dim))
+
+    def capacity(self, r: int, step: int) -> float:
+        """Slice 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) -> torch.Tensor:
+        # x: [S, dim] pooled per-slice features
+        return self.heads[r](x + self.rec_emb[r])
+
+
+class RecursiveConvBlock(nn.Module):
+    """A shared inverted-residual conv block reused at every recursion.
+
+    Depthwise-separable (MobileNet-v2 style) so the recursive core stays cheap
+    while the stem does the heavy feature extraction.
+    """
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.pw1 = nn.Conv2d(dim, dim * 2, 1)
+        self.dw = nn.Conv2d(dim * 2, dim * 2, 3, padding=1, groups=dim * 2)
+        self.pw2 = nn.Conv2d(dim * 2, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [N, dim, h, w]
+        identity = x
+        x = self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        x = self.pw2(F.gelu(self.dw(F.gelu(self.pw1(x)))))
+        return identity + 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:
+        # x: [S, dim]
+        w = F.softmax(x @ self.query * self.scale, dim=0)
+        return (w.unsqueeze(1) * x).sum(0)
+
+
+class ChannelLayerNorm(nn.Module):
+    """LayerNorm over the channel dim of a [N, C, H, W] feature map."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+
+class MoRCNN(nn.Module):
+    """Pretrained CNN stem + MoR recursive core + per-slice depth router."""
+
+    def __init__(
+        self,
+        *,
+        stem_name: str = "convnext_tiny",
+        pretrained: bool = True,
+        n_recursions: int = 3,
+        capacities: tuple[float, ...] = (1.0, 2.0 / 3.0, 1.0 / 3.0),
+        router_hidden: int = 128,
+        core_blocks: int = 2,
+        n_classes: int = 12,
+        router_warmup_steps: int = 0,
+        router_init_bias: float = 0.0,
+    ):
+        super().__init__()
+        if not HAS_TIMM:
+            raise ImportError("timm is required for a pretrained stem")
+        self.n_recursions = n_recursions
+        self.capacities = list(capacities)
+        assert len(self.capacities) == n_recursions
+
+        self.stem = timm.create_model(
+            stem_name, pretrained=pretrained, features_only=True, num_classes=0
+        )
+        out_dim = self.stem.feature_info.channels()[-1]
+
+        self.core = nn.Sequential(*[RecursiveConvBlock(out_dim) for _ in range(core_blocks)])
+        self.router = DepthRouter(
+            out_dim,
+            router_hidden,
+            n_recursions,
+            self.capacities,
+            init_bias=router_init_bias,
+            warmup_steps=router_warmup_steps,
+        )
+        self.rec_emb = nn.Parameter(torch.zeros(n_recursions, out_dim))
+
+        self.exit = nn.Sequential(
+            ChannelLayerNorm(out_dim),
+            nn.Conv2d(out_dim, out_dim, 1),
+            nn.GELU(),
+        )
+        self.pool = SliceAttentionPool(out_dim)
+
+        self.head = nn.Sequential(
+            nn.Linear(out_dim, out_dim * 2),
+            nn.GELU(),
+            nn.Linear(out_dim * 2, n_classes),
+        )
+
+    def _run_core(self, feat: torch.Tensor, r: int) -> torch.Tensor:
+        emb = self.rec_emb[r].view(1, -1, 1, 1)
+        return self.core(feat + emb)
+
+    def _global_pool(self, feat: torch.Tensor) -> torch.Tensor:
+        return feat.mean(dim=(2, 3))
+
+    def forward(
+        self, x: torch.Tensor, step: int = 0
+    ) -> tuple[torch.Tensor, list[float]]:
+        """Forward one study.
+
+        Args:
+            x: [S, 3, H, W] sampled slices of a single study.
+            step: current optimizer step, used for router warmup.
+
+        Returns:
+            (logits [n_classes], route_stats) where route_stats is the active
+            slice fraction at each recursion.
+        """
+        S = x.size(0)
+        feat = self.stem(x)[-1]  # [S, C, h, w]
+        C = feat.size(1)
+        pooled = self._global_pool(feat)
+
+        stats: list[float] = []
+        for r in range(self.n_recursions):
+            probs = torch.sigmoid(self.router(pooled, r))
+            cap = self.router.capacity(r, step)
+            if r == self.n_recursions - 1:
+                active = torch.ones(S, dtype=torch.bool, device=x.device)
+            else:
+                k = max(1, int(round(S * cap)))
+                active = torch.zeros(S, dtype=torch.bool, device=x.device)
+                active[torch.topk(probs, k).indices] = True
+            stats.append(active.float().mean().item())
+
+            if active.all():
+                feat = self._run_core(feat, r)
+            else:
+                idx = active.nonzero(as_tuple=False).squeeze(1)
+                updated = self._run_core(feat[idx], r)
+                feat = feat.clone()
+                feat[idx] = updated
+            pooled = self._global_pool(feat)
+
+        slice_feats = self.exit(feat)  # [S, C, h, w]
+        slice_feats = self._global_pool(slice_feats)  # [S, C]
+        study = self.pool(slice_feats)  # [C]
+        logits = self.head(study)  # [n_classes]
+        return logits, stats

src/src/models/mor_vit.py

@@ -0,0 +1,386 @@
+"""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, []

src/src/models/mricore_stem.py

@@ -0,0 +1,108 @@
+"""MRI-CORE stem: MRI-pretrained ViT-B feature extractor.
+
+Loads the MRI-CORE checkpoint (DINOv2-style ViT-B, SAM-init, pretrained on
+6.1M MRI slices) and exposes a timm-compatible surface so the existing
+MoR/MoE cores can use it as their frozen stem:
+
+    stem.embed_dim          -> 768
+    stem.num_prefix_tokens  -> 1  (CLS only, no registers)
+    stem.pos_embed          -> [1, 197, 768]
+    stem.forward_features(x)-> [S, 197, 768]  (frozen, eval mode)
+
+Key layout in the checkpoint:  teacher.backbone.{cls_token, pos_embed,
+patch_embed.proj, blocks.<g>.<i>, norm}  with blocks nested 4 x 3 = 12.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+_CKPT = "MRI_CORE_vitb.pth"
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int = 768, n_heads: int = 12, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
+        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:
+        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class MRICoreStem(nn.Module):
+    """Frozen MRI-CORE ViT-B feature extractor (embed_dim=768, 197 tokens)."""
+
+    def __init__(self, ckpt_path: str | None = None, freeze: bool = True):
+        super().__init__()
+        if ckpt_path is None:
+            for base in (
+                r"G:\RSNA-Knee\cache",
+                os.path.join(os.path.dirname(__file__), "..", "..", "weights"),
+            ):
+                cand = os.path.join(base, _CKPT)
+                if os.path.isfile(cand):
+                    ckpt_path = cand
+                    break
+        if ckpt_path is None or not os.path.isfile(ckpt_path):
+            raise FileNotFoundError(f"MRI-CORE checkpoint not found ({_CKPT})")
+        self.embed_dim = 768
+        self.num_prefix_tokens = 1
+        self.patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16, bias=True)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
+        self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
+        self.mask_token = nn.Parameter(torch.zeros(1, 768))
+        self.blocks = nn.ModuleList(
+            [nn.ModuleList([Block() for _ in range(3)]) for _ in range(4)]
+        )
+        self.norm = nn.LayerNorm(768)
+        st = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+        inner = st["teacher"]
+        if isinstance(inner, dict) and "state_dict" in inner:
+            inner = inner["state_dict"]
+        own = self.state_dict()
+        prefix = "backbone."
+        missing = []
+        for k in list(inner.keys()):
+            if not k.startswith(prefix):
+                inner.pop(k)
+                continue
+            name = k[len(prefix):]
+            if name not in own:
+                continue
+            if tuple(inner[k].shape) == tuple(own[name].shape):
+                own[name] = inner[k]
+            else:
+                missing.append((name, tuple(inner[k].shape), tuple(own[name].shape)))
+        self.load_state_dict(own)
+        if missing:
+            print(f"mricore: {len(missing)} shape mismatches skipped "
+                  f"(e.g. {missing[0]})", flush=True)
+        if freeze:
+            for p in self.parameters():
+                p.requires_grad_(False)
+        self.eval()
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        """[S, 3, H, W] -> [S, 197, 768] (CLS + 196 patch tokens)."""
+        x = self.patch_embed(x).flatten(2).transpose(1, 2)  # [S, 196, 768]
+        x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1)
+        x = x + self.pos_embed
+        for group in self.blocks:
+            for blk in group:
+                x = blk(x)
+        return self.norm(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_features(x)

src/src/models/plant_vit.py

@@ -0,0 +1,137 @@
+"""PlantViT - thin adapter over RSNA MoR/MoE for single-image plant ID.
+Reuses moe_vit.MoEViT and mor_vit.MoRViT frozen DINOv3 stem + routing.
+Single image = single study with S=1 slice, so pool reduces to patchmean.
+Adds 512-D embedder head per architecture.md §4.
+
+Usage:
+  from src.models.plant_vit import PlantViT
+  model = PlantViT(stem_name="vit_base_patch16_dinov3", n_classes=500, use_moe=True, embed_dim=512)
+  logits, embed, aux = model(x)  # x [B,3,336,336] -> logits [B,500], embed [B,512] L2
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .moe_vit import MoEViT
+from .mor_vit import MoRViT
+
+
+class EmbedHead(nn.Module):
+    """Plant-specific projection: 768/384 -> 512 L2 per architecture.md §4."""
+    def __init__(self, in_dim: int, embed_dim: int = 512):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(in_dim)
+        self.proj = nn.Linear(in_dim, embed_dim, bias=False)
+        self.ln2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x [B, D]
+        x = self.ln1(x)
+        x = self.proj(x)
+        x = self.ln2(x)
+        return F.normalize(x, dim=-1)
+
+
+class PlantViT(nn.Module):
+    """Wrapper that exposes (logits, embed, aux_losses) with interchangeable stem.
+
+    Args:
+        stem_name: timm name, e.g. vit_small_patch16_dinov3, vit_base_patch16_dinov3, vit_large_patch16_dinov3, convnext_tiny
+        n_classes: 500 for WA top500
+        embed_dim: 512 (test 256/128 later)
+        use_moe: True -> MoE++ (MoEViT), False -> MoR (MoRViT), None -> dense
+        freeze_stem: keep DINOv3 frozen for CE baseline
+        All MoE/MoR kwargs (num_ffn, top_k, d_ff, n_recursions, capacities etc.) passed through.
+    """
+    def __init__(
+        self,
+        *,
+        stem_name: str = "vit_base_patch16_dinov3",
+        n_classes: int = 500,
+        embed_dim: int = 512,
+        use_moe: bool | None = True,
+        freeze_stem: bool = True,
+        # MoE knobs
+        num_ffn: int = 4,
+        top_k: int = 2,
+        d_ff: int = 384,
+        # MoR knobs
+        n_recursions: int = 3,
+        capacities: tuple[float, ...] = (1.0, 0.667, 0.333),
+        core_blocks: int = 2,
+        router_hidden: int = 128,
+        pretrained: bool = True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.use_moe = use_moe
+        # build core - keep RSNA defaults, just swap n_classes and stem
+        if use_moe is True:
+            self.core = MoEViT(
+                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
+                core_blocks=core_blocks, d_ff=d_ff, num_ffn=num_ffn, top_k=top_k,
+                n_classes=n_classes, **{k: v for k, v in kwargs.items() if k in ["gate_ctx","per_target","tau","n_zero","n_copy","n_const","gating_residual"]},
+            )
+            feat_dim = self.core.core_dim  # 384
+        elif use_moe is False:
+            self.core = MoRViT(
+                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
+                n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
+                router_hidden=router_hidden, n_classes=n_classes,
+            )
+            feat_dim = self.core.core_dim
+        else:
+            # dense baseline - MoR dense=True
+            self.core = MoRViT(
+                stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
+                n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
+                router_hidden=router_hidden, n_classes=n_classes, dense=True,
+            )
+            feat_dim = self.core.core_dim
+
+        # embedder on pooled slice feat (before head pool, use same dim)
+        # For MoE/MoR, slice feat dim = core.core_dim (384) ; for base ViT-B 768 -> still 384 via input_proj
+        # If stem is vit_base (768) we still route through input_proj -> 384, so embed in 384
+        self.embed = EmbedHead(feat_dim, embed_dim)
+        self.n_classes = n_classes
+        self.embed_dim = embed_dim
+        self.stem_name = stem_name
+
+    def forward(self, x: torch.Tensor, step: int = 0):
+        """x [B,3,H,W] single image per plant -> logits [B,C], embed [B,512], aux dict"""
+        # RSNA core expects [S,3,H,W] where S is slices per study; for plant S=B single slice per sample
+        # Use core.features single-batch path
+        feats = self.core.features(x)  # [B, D]  (MoEViT/MoRViT handle bfloat16 internally)
+        logits = self.core.head(self.core.pool(feats.unsqueeze(1).repeat(1,1,1)) if False else feats)  # bypass buggy repeat - just pool slice feats
+        # Actually MoEViT/MoRViT head expects pooled study vector [D]; we have [B,D] already pooled per image
+        # Re-derive logits via head directly for single-slice case:
+        # MoEViT/MoRViT: features -> [S,D] -> pool -> [D] -> head -> [C]
+        # For B independent images, features is [B,D] already per-image slice feat, so head per row:
+        # Use core.head on each row individually via batched linear - head is Sequential Linear/GELU/Linear
+        # For per_target head, need [B, n_classes, D] - not used for plant (per_target=False by default)
+        # Simpler: call core.head on feats directly (feats is [B,D] -> head expects [D] but Linear broadcasts on last dim)
+        # Do batched:
+        if isinstance(self.core.head, nn.Sequential):
+            logits = self.core.head(feats)  # [B,C] broadcasts
+        else:
+            # PerTarget not used
+            logits = self.core.head(feats)
+
+        embed = self.embed(feats)
+        aux = {}
+        # surface aux losses if present
+        try:
+            aux["moe_aux"] = float(self.core.moe_aux_loss().item()) if hasattr(self.core, "moe_aux_loss") else 0.0
+        except:
+            aux["moe_aux"] = 0.0
+        try:
+            aux["tv"] = float(self.core.routing_tv_loss().item()) if hasattr(self.core, "routing_tv_loss") else float(self.core.routing_smoothness_loss().item() if hasattr(self.core, "routing_smoothness_loss") else 0.0)
+        except:
+            aux["tv"] = 0.0
+        return logits, embed, aux
+
+    def features(self, x: torch.Tensor) -> torch.Tensor:
+        return self.core.features(x)

src/src/models/plant_vit_v2.py

@@ -0,0 +1,169 @@
+"""
+PlantViT v2 Architecture (PlantDetect v2.0):
+1. DINOv3 ViT-Base Backbone @ 448x448 High-Res (784 patch tokens).
+2. Multi-Layer Feature Pyramid Fusion (Block 8 Micro-Texture + Block 11 Macro-Semantics).
+3. Pure 16 SwiGLU Mixture-of-Experts (MoESwiGLUv2 with Gating Residuals & Context Projection).
+4. Sub-Center ArcFace Metric Head (K=3 Phenology Centers per Class).
+5. Hierarchical Multi-Task Taxonomic Classification Heads (Species + Genus + Family).
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+
+from .moe_swiglu_v2 import MoESwiGLUv2
+from .subcenter_arcface import SubCenterArcFace
+
+class MultiLayerFeatureFusion(nn.Module):
+    def __init__(self, dim: int = 768, out_dim: int = 512):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+
+        self.norm_mid = nn.LayerNorm(dim)
+        self.norm_deep = nn.LayerNorm(dim)
+
+        self.fusion_proj = nn.Sequential(
+            nn.Linear(dim * 2, out_dim),
+            nn.GELU(),
+            nn.LayerNorm(out_dim),
+            nn.Linear(out_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+
+    def forward(self, mid_tokens: torch.Tensor, deep_tokens: torch.Tensor) -> torch.Tensor:
+        # Pool spatial tokens over sequence dimension
+        mid_pooled = self.norm_mid(mid_tokens.mean(dim=1))
+        deep_pooled = self.norm_deep(deep_tokens.mean(dim=1))
+
+        fused = torch.cat([mid_pooled, deep_pooled], dim=-1)
+        emb = self.fusion_proj(fused)
+        return F.normalize(emb, dim=-1)
+
+class PlantViTv2(nn.Module):
+    def __init__(
+        self,
+        stem_name: str = "vit_base_patch16_dinov3.lvd1689m",
+        img_size: int = 448,
+        n_species: int = 1000,
+        n_genera: int = 358,
+        n_families: int = 80,
+        embed_dim: int = 512,
+        use_moe: bool = True,
+        num_ffn: int = 16,
+        d_ff: int = 512,
+        top_k: int = 2,
+        k_subcenters: int = 3,
+        pretrained: bool = True,
+        freeze_stem: bool = False
+    ):
+        super().__init__()
+        self.stem_name = stem_name
+        self.img_size = img_size
+        self.n_species = n_species
+        self.n_genera = n_genera
+        self.n_families = n_families
+        self.embed_dim = embed_dim
+        self.use_moe = use_moe
+
+        # 1. DINOv3 Backbone
+        stem_model_name = stem_name if stem_name.startswith("vit_") else f"{stem_name}.lvd1689m"
+        self.stem = timm.create_model(
+            stem_model_name,
+            pretrained=pretrained,
+            img_size=img_size,
+            num_classes=0,
+            dynamic_img_size=True
+        )
+        self.dim = self.stem.embed_dim
+
+        if freeze_stem:
+            for p in self.stem.parameters():
+                p.requires_grad = False
+
+        # 2. Pure 16-SwiGLU MoE Layer
+        if self.use_moe:
+            self.moe_layer = MoESwiGLUv2(
+                d_model=self.dim,
+                d_ff=d_ff,
+                num_ffn=num_ffn,
+                top_k=top_k,
+                gating_residual=True,
+                gate_ctx=True
+            )
+            self.moe_norm = nn.LayerNorm(self.dim)
+        else:
+            self.moe_layer = None
+
+        # 3. Multi-Layer Feature Fusion (Block 8 + Block 11)
+        self.fusion = MultiLayerFeatureFusion(dim=self.dim, out_dim=embed_dim)
+
+        # 4. Multi-Task Taxonomic Heads
+        self.head_species = nn.Linear(embed_dim, n_species, bias=False)
+        self.head_genus = nn.Linear(embed_dim, n_genera, bias=False)
+        self.head_family = nn.Linear(embed_dim, n_families, bias=False)
+
+        # 5. Sub-Center ArcFace Metric Head (K=3 Sub-Centers)
+        self.subcenter_arcface = SubCenterArcFace(
+            embed_dim=embed_dim,
+            n_classes=n_species,
+            k_subcenters=k_subcenters,
+            m=0.5,
+            s=30.0
+        )
+
+    def _to_tokens(self, feat: torch.Tensor) -> torch.Tensor:
+        """Converts [B, C, H, W] or [B, T, C] to standard [B, T, C]."""
+        if feat.dim() == 4:
+            return feat.flatten(2).transpose(1, 2)
+        return feat
+
+    def forward_features(self, x: torch.Tensor):
+        intermediates, final_feat = self.stem.forward_intermediates(x, indices=[7, 11])
+        
+        if isinstance(intermediates, (list, tuple)):
+            mid_raw = intermediates[0]
+            deep_raw = intermediates[1] if len(intermediates) > 1 else final_feat
+        else:
+            mid_raw = intermediates
+            deep_raw = final_feat
+
+        if isinstance(mid_raw, (list, tuple)): mid_raw = mid_raw[0]
+        if isinstance(deep_raw, (list, tuple)): deep_raw = deep_raw[0]
+
+        mid_tokens = self._to_tokens(mid_raw)
+        deep_tokens = self._to_tokens(deep_raw)
+
+        aux_loss = torch.tensor(0.0, device=x.device)
+        gate_logits = None
+        if self.use_moe and self.moe_layer is not None:
+            ctx = deep_tokens[:, 0]
+            moe_out, aux_loss, gate_logits = self.moe_layer(self.moe_norm(deep_tokens), ctx=ctx)
+            deep_tokens = deep_tokens + moe_out
+
+        return mid_tokens, deep_tokens, aux_loss, gate_logits
+
+    def forward(self, x: torch.Tensor, y_species: torch.Tensor = None):
+        mid_tokens, deep_tokens, moe_aux, _ = self.forward_features(x)
+
+        # 1. Multi-Layer Pyramid Fusion (Block 8 Texture + Block 11 Macro)
+        emb = self.fusion(mid_tokens, deep_tokens)
+
+        # 2. Multi-Task Linear Logits
+        logits_species = self.head_species(emb)
+        logits_genus = self.head_genus(emb)
+        logits_family = self.head_family(emb)
+
+        # 3. Sub-Center ArcFace Margin Logits
+        arcface_logits = self.subcenter_arcface(emb, y_species)
+
+        aux_dict = {"moe_aux": moe_aux}
+
+        return (
+            logits_species,
+            logits_genus,
+            logits_family,
+            arcface_logits,
+            emb,
+            aux_dict
+        )

src/src/models/subcenter_arcface.py

@@ -0,0 +1,69 @@
+"""
+Sub-Center ArcFace with K Sub-Centers per Class.
+Designed for fine-grained botanical vision to capture multi-modal phenology
+(e.g., Flower vs. Leaf/Foliage vs. Fruit/Seed Pods) without intra-class noise.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SubCenterArcFace(nn.Module):
+    """
+    Sub-Center ArcFace for 512-D L2 Embeddings.
+    Maintains K sub-centers per class (default K=3).
+    """
+    def __init__(self, embed_dim: int, n_classes: int, k_subcenters: int = 3, m: float = 0.5, s: float = 30.0, easy_margin: bool = False):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.n_classes = n_classes
+        self.k_subcenters = k_subcenters
+        self.m = m
+        self.s = s
+        self.easy_margin = easy_margin
+
+        # Weight matrix: [n_classes * k_subcenters, embed_dim]
+        self.W = nn.Parameter(torch.randn(n_classes * k_subcenters, embed_dim) * 0.01)
+        nn.init.xavier_uniform_(self.W)
+
+        # Precomputed angular constants
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+        self.th = math.cos(math.pi - m)
+        self.mm = math.sin(math.pi - m) * m
+
+    def forward(self, emb: torch.Tensor, y: torch.Tensor = None) -> torch.Tensor:
+        """
+        emb: [B, D] L2 normalized embedding
+        y: [B] Ground-truth class labels (0 .. n_classes-1)
+        Returns: [B, n_classes] Logits scaled by s
+        """
+        # Normalize weights and embeddings
+        W_norm = F.normalize(self.W, dim=1) # [C * K, D]
+        emb_norm = F.normalize(emb, dim=1)   # [B, D]
+
+        # Compute cosine similarity across all sub-centers: [B, C * K]
+        cos_all = F.linear(emb_norm, W_norm)
+        B = emb.size(0)
+
+        # Reshape to [B, C, K] and select the maximum cosine per class
+        cos_sub = cos_all.view(B, self.n_classes, self.k_subcenters)
+        cos_max, _ = cos_sub.max(dim=-1) # [B, C]
+        cos = cos_max.clamp(-1.0 + 1e-7, 1.0 - 1e-7)
+
+        # Inference mode: return scaled max cosine logits
+        if y is None:
+            return cos * self.s
+
+        # Training mode: Apply additive angular margin to ground-truth class
+        sin = torch.sqrt((1.0 - cos ** 2).clamp(0.0, 1.0))
+        cos_m = cos * self.cos_m - sin * self.sin_m
+
+        if self.easy_margin:
+            cos_m = torch.where(cos > 0, cos_m, cos)
+        else:
+            cos_m = torch.where(cos > self.th, cos_m, cos - self.mm)
+
+        one_hot = F.one_hot(y, num_classes=self.n_classes).bool()
+        logits = torch.where(one_hot, cos_m, cos)
+        return logits * self.s