Escarda-VE: JEPA+HRM vision encoder + Byrne-VE comparison

README.md

@@ -0,0 +1,38 @@
+---
+license: apache-2.0
+tags:
+- vision-encoder
+- jepa
+- hrm
+- dinov2-distillation
+---
+
+# Escarda-VE
+
+**Escarda-VE** is the vision encoder of the **Escarda** model family (~39.6M params,
+448px / patch16 / 784 tokens). It is the **Byrne-VE** architecture **plus a JEPA head
+added alongside the HRM refinement block** — Escarda's defining family trait.
+
+- ViT-style: RMSNorm, 2D-axial RoPE, QK-Norm, SwiGLU, HRM refinement.
+- **JEPA head** (auxiliary): from a patch token it predicts the raster-order neighbour
+  patch's own representation (stop-gradient target, 1−cosine loss). **Zero inference cost.**
+- Trained by distillation from a frozen **DINOv2-base** teacher (50k steps, CLS + patch),
+  then 20k steps of DINO-style **self-distillation** (EMA teacher, cosine prototypes, no collapse).
+
+## Escarda-VE vs Byrne-VE (DINOv2 teacher-alignment, n=1024 held-out)
+
+| | Byrne-VE | Escarda-VE |
+|---|---|---|
+| Params | 39.34M | 39.60M (+JEPA head) |
+| CLS cosine | **0.776** | 0.771 |
+| PATCH cosine | **0.600** | 0.584 |
+| JEPA self-consistency | — | 0.040 |
+
+**Read:** Escarda trades ~1–3% teacher-alignment (the JEPA objective diverts a little
+capacity away from pure supervised mimicry) in exchange for a self-supervised spatial
+neighbour-prediction signal that costs nothing at inference. Same size class as Byrne-VE.
+
+## Related
+- [Byrne-VE](https://huggingface.co/Quazim0t0/Byrne-VE) — HRM-only sibling
+- [Escarda-Docling-126M](https://huggingface.co/Quazim0t0/Escarda-Docling-126M) — document VLM built on this encoder
+- [Escarda-86M-Base](https://huggingface.co/Quazim0t0/Escarda-86M-Base)

escarda_ve.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:056e8ba04e6f4727a340eccf1950a86a0424759f07a94af8dd256654590db113
+size 158447229

requirements.txt

@@ -0,0 +1,4 @@
+torch
+numpy
+pillow
+datasets

vision/__init__.py

@@ -0,0 +1,3 @@
+from .model import VisionEncoder, VisionConfig
+
+__all__ = ["VisionEncoder", "VisionConfig"]

vision/dataset.py

@@ -0,0 +1,89 @@
+"""Image inputs for distillation. PIL-only (no torchvision dependency).
+
+Two sources, same preprocessing (resize shorter side to image_size+32, center-crop,
+ImageNet-normalize — what DINOv2 expects):
+  - ImageFolder:    a local folder of images, searched recursively.
+  - HFStreamImages: an HF streaming dataset (e.g. timm/imagenet-1k-wds) — shards are
+                    pulled on the fly, nothing is stored on disk.
+"""
+import os
+import glob
+import random
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, IterableDataset
+
+_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+_EXTS = (".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff")
+
+
+def preprocess(img, size=224):
+    """PIL.Image -> normalized [3,H,W] float tensor."""
+    img = img.convert("RGB")
+    w, h = img.size
+    s = size + 32
+    nw, nh = (s, round(h * s / w)) if w < h else (round(w * s / h), s)
+    img = img.resize((nw, nh), Image.BICUBIC)
+    left, top = (nw - size) // 2, (nh - size) // 2
+    img = img.crop((left, top, left + size, top + size))
+    arr = (np.asarray(img, dtype=np.float32) / 255.0 - _MEAN) / _STD   # [H,W,3]
+    return torch.from_numpy(arr).permute(2, 0, 1)                       # [3,H,W]
+
+
+def load_image(path, size=224):
+    return preprocess(Image.open(path), size)
+
+
+class ImageFolder(Dataset):
+    def __init__(self, root, image_size=224):
+        self.size = image_size
+        self.files = sorted(p for p in glob.glob(os.path.join(root, "**", "*"), recursive=True)
+                            if p.lower().endswith(_EXTS))
+        if not self.files:
+            raise FileNotFoundError(f"no images found under {root}")
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, i):
+        try:
+            return load_image(self.files[i], self.size)
+        except Exception:                              # skip a corrupt file
+            return self.__getitem__(random.randrange(len(self.files)))
+
+
+def _find_image(example, key=None):
+    """Pull the PIL image out of an HF example dict (auto-detect if key not given)."""
+    if key is not None:
+        return example[key]
+    for v in example.values():
+        if isinstance(v, Image.Image):
+            return v
+    raise KeyError("no PIL image column found; pass --image-key")
+
+
+class HFStreamImages(IterableDataset):
+    """Streams images from a Hugging Face dataset (no full download).
+
+    e.g. name='timm/imagenet-1k-wds', split='train'. Reshuffles shard/sample order
+    each epoch with a bounded buffer so DataLoader can keep shuffle=False.
+    """
+    def __init__(self, name, split="train", image_size=224, config=None,
+                 image_key=None, shuffle_buffer=2000, seed=0):
+        from datasets import load_dataset
+        self.size = image_size
+        self.image_key = image_key
+        ds = load_dataset(name, config, split=split, streaming=True)
+        if shuffle_buffer:
+            ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
+        self.ds = ds
+
+    def __iter__(self):
+        for ex in self.ds:
+            try:
+                yield preprocess(_find_image(ex, self.image_key), self.size)
+            except Exception:                          # skip a corrupt/odd sample
+                continue

vision/eval_probe.py

@@ -0,0 +1,81 @@
+"""Downstream probe: weighted k-NN classification accuracy of a frozen encoder.
+
+A label-free distillation/self-distillation objective tells you how close two
+feature spaces are, but not whether the features are *useful*. This measures that
+directly: freeze the encoder, embed a labelled set, and classify by weighted k-NN
+in feature space (no training, no learned head). Use it to compare checkpoints —
+e.g. the teacher-distilled `vision_last.pt` vs the self-distilled
+`vision_self_last.pt` — under identical preprocessing.
+
+  python -m vision.eval_probe --ckpt checkpoints/vision_last.pt
+  python -m vision.eval_probe --ckpt checkpoints/vision_self_last.pt
+
+CIFAR-10 (10 classes, ungated, ~170 MB) is the default probe set: small and fast,
+and since both checkpoints get the exact same pipeline the *relative* number is a
+fair read even though CIFAR is out-of-domain for an ImageNet encoder.
+"""
+import argparse
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .infer import load_encoder
+from .dataset import preprocess
+
+
+@torch.no_grad()
+def embed(model, cfg, images, device, bs=128):
+    feats = []
+    for i in range(0, len(images), bs):
+        batch = torch.stack([preprocess(im, cfg.image_size) for im in images[i:i + bs]]).to(device)
+        feats.append(F.normalize(model(batch)["pooled"], dim=-1).cpu())
+    return torch.cat(feats)
+
+
+def knn_accuracy(bank, bank_y, query, query_y, k=20, T=0.07, n_classes=10):
+    sim = query @ bank.T                                  # cosine (features L2-normed)
+    topk_sim, topk_idx = sim.topk(min(k, bank.shape[0]), dim=1)
+    topk_y = bank_y[topk_idx]                             # (Q, k)
+    w = (topk_sim / T).exp()                              # distance weighting
+    votes = torch.zeros(query.shape[0], n_classes)
+    votes.scatter_add_(1, topk_y, w)
+    pred = votes.argmax(1)
+    return (pred == query_y).float().mean().item()
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--dataset", default="uoft-cs/cifar10")
+    ap.add_argument("--n-bank", type=int, default=5000)
+    ap.add_argument("--n-query", type=int, default=1000)
+    ap.add_argument("--k", type=int, default=20)
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    from datasets import load_dataset
+    ds = load_dataset(args.dataset)
+    train, test = ds["train"], ds["test"]
+    img_key = "img" if "img" in train.column_names else "image"
+    lab_key = "label" if "label" in train.column_names else "fine_label"
+    n_classes = train.features[lab_key].num_classes
+
+    train = train.shuffle(seed=args.seed).select(range(args.n_bank))
+    test = test.shuffle(seed=args.seed).select(range(args.n_query))
+    model, cfg = load_encoder(args.ckpt, device)
+
+    bank = embed(model, cfg, train[img_key], device)
+    bank_y = torch.tensor(train[lab_key])
+    query = embed(model, cfg, test[img_key], device)
+    query_y = torch.tensor(test[lab_key])
+    acc = knn_accuracy(bank, bank_y, query, query_y, k=args.k, n_classes=n_classes)
+
+    chance = 1.0 / n_classes
+    print(f"{args.ckpt}: k-NN(k={args.k}) top-1 = {acc*100:.2f}%  "
+          f"({args.n_bank} bank / {args.n_query} query, {n_classes}-way, chance {chance*100:.0f}%)")
+
+
+if __name__ == "__main__":
+    main()

vision/infer.py

@@ -0,0 +1,39 @@
+"""Encode an image with the distilled VisionEncoder.
+
+  python -m vision.infer --ckpt checkpoints/vision_last.pt --image photo.jpg
+"""
+import argparse
+
+import torch
+
+from .model import VisionEncoder, VisionConfig
+from .dataset import load_image
+
+
+def load_encoder(ckpt="checkpoints/vision_last.pt", device="cpu"):
+    ck = torch.load(ckpt, map_location=device)
+    cfg = VisionConfig(**ck["cfg"]) if "cfg" in ck else VisionConfig()
+    m = VisionEncoder(cfg).to(device)
+    m.load_state_dict(ck["student"])
+    m.eval()
+    return m, cfg
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="checkpoints/vision_last.pt")
+    ap.add_argument("--image", required=True)
+    ap.add_argument("--out", default="feature.pt")
+    args = ap.parse_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    m, cfg = load_encoder(args.ckpt, device)
+    x = load_image(args.image, cfg.image_size).unsqueeze(0).to(device)
+    with torch.no_grad():
+        out = m(x)
+    print(f"pooled feature {tuple(out['pooled'].shape)} | {out['tokens'].shape[1]} patch tokens")
+    torch.save(out["pooled"][0].cpu(), args.out)
+    print(f"saved {args.out}")
+
+
+if __name__ == "__main__":
+    main()

vision/model.py

@@ -0,0 +1,392 @@
+"""Tiny Vision Encoder (~40M params) — from scratch in PyTorch.
+
+A ViT-style image encoder that shares the SpikeWhale/Inflect family DNA used by
+the rest of this repo (see reference/model_v2.py):
+
+  - RMSNorm everywhere (pre-norm blocks + QK-Norm)
+  - 2D axial RoPE on the patch grid (replaces learned/sinusoidal position embeds)
+  - QK-Norm attention (Gemma2/OLMo2-style logit stability)
+  - SwiGLU FFN (silu(gate) * up -> down)
+
+It is non-causal (full bidirectional attention) and emits both a pooled CLS
+embedding and the full sequence of patch tokens, so it can feed a classifier
+head, a contrastive projection, or a multimodal decoder.
+
+Default config (`VisionConfig()`) targets ~40M parameters; run this file
+directly to print the exact count.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class VisionConfig:
+    image_size: int = 224
+    patch_size: int = 16
+    in_channels: int = 3
+    hidden_size: int = 512
+    num_layers: int = 12
+    num_heads: int = 8
+    ffn_mult: float = 8 / 3          # SwiGLU effective expansion
+    ffn_multiple_of: int = 64        # round ffn dim up to this
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    use_qk_norm: bool = True
+    use_cls_token: bool = True
+    drop_path: float = 0.0           # stochastic depth rate (linearly scaled)
+    use_hrm_refine: bool = True      # family HRM refinement before output norm
+    hrm_refine_dim: int = 256
+    hrm_refine_steps: int = 2
+    # Escarda family trait: JEPA auxiliary head ALONGSIDE HRM (both active). Predicts the
+    # trunk's own neighbouring patch representation (detached target) in embedding space.
+    use_jepa: bool = False           # Byrne = False; Escarda = True
+    jepa_pred_dim: int = 256
+    jepa_horizon: int = 1
+
+    @property
+    def grid_size(self) -> int:
+        assert self.image_size % self.patch_size == 0, "image_size must divide by patch_size"
+        return self.image_size // self.patch_size
+
+    @property
+    def num_patches(self) -> int:
+        return self.grid_size ** 2
+
+    @property
+    def head_dim(self) -> int:
+        assert self.hidden_size % self.num_heads == 0, "hidden_size must divide by num_heads"
+        return self.hidden_size // self.num_heads
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x * self.weight).to(dtype)
+
+
+def build_axial_rope(grid_size: int, head_dim: int, theta: float, device, dtype):
+    """2D axial RoPE cos/sin for a (grid_size x grid_size) patch grid.
+
+    Half of head_dim encodes the row coordinate, the other half the column.
+    Returns cos, sin each of shape (num_patches, head_dim).
+    """
+    assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D axial RoPE"
+    quarter = head_dim // 4
+    inv_freq = 1.0 / (theta ** (torch.arange(0, quarter, device=device).float() / quarter))
+
+    coords = torch.arange(grid_size, device=device).float()
+    rows = coords.view(-1, 1).expand(grid_size, grid_size).reshape(-1)  # (P,)
+    cols = coords.view(1, -1).expand(grid_size, grid_size).reshape(-1)  # (P,)
+
+    row_ang = torch.outer(rows, inv_freq)  # (P, quarter)
+    col_ang = torch.outer(cols, inv_freq)  # (P, quarter)
+    ang = torch.cat([row_ang, col_ang], dim=-1)  # (P, head_dim//2)
+    emb = torch.cat([ang, ang], dim=-1)          # (P, head_dim)
+    return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    # x: (B, H, P, D), cos/sin: (P, D)
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    return x * cos + rotate_half(x) * sin
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.num_heads = cfg.num_heads
+        self.head_dim = cfg.head_dim
+        self.use_qk_norm = cfg.use_qk_norm
+        self.qkv = nn.Linear(cfg.hidden_size, 3 * cfg.hidden_size, bias=False)
+        self.proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
+        if cfg.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                rope_offset: int) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).view(B, N, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q = q.transpose(1, 2)  # (B, H, N, D)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # RoPE only on patch tokens; the leading CLS token (offset) is left untouched.
+        if rope_offset:
+            q_cls, q_pat = q[:, :, :rope_offset], q[:, :, rope_offset:]
+            k_cls, k_pat = k[:, :, :rope_offset], k[:, :, rope_offset:]
+            q_pat = apply_rope(q_pat, cos, sin)
+            k_pat = apply_rope(k_pat, cos, sin)
+            q = torch.cat([q_cls, q_pat], dim=2)
+            k = torch.cat([k_cls, k_pat], dim=2)
+        else:
+            q = apply_rope(q, cos, sin)
+            k = apply_rope(k, cos, sin)
+
+        out = F.scaled_dot_product_attention(q, k, v)  # full bidirectional attention
+        out = out.transpose(1, 2).reshape(B, N, C)
+        return self.proj(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        hidden = int(cfg.ffn_mult * cfg.hidden_size)
+        hidden = cfg.ffn_multiple_of * math.ceil(hidden / cfg.ffn_multiple_of)
+        self.gate = nn.Linear(cfg.hidden_size, hidden, bias=False)
+        self.up = nn.Linear(cfg.hidden_size, hidden, bias=False)
+        self.down = nn.Linear(hidden, cfg.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class DropPath(nn.Module):
+    """Per-sample stochastic depth on the residual branch."""
+
+    def __init__(self, p: float):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.p == 0.0 or not self.training:
+            return x
+        keep = 1 - self.p
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = torch.empty(shape, dtype=x.dtype, device=x.device).bernoulli_(keep)
+        return x * mask / keep
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: VisionConfig, drop_path: float):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.ffn = SwiGLU(cfg)
+        self.drop_path = DropPath(drop_path)
+
+    def forward(self, x, cos, sin, rope_offset):
+        x = x + self.drop_path(self.attn(self.attn_norm(x), cos, sin, rope_offset))
+        x = x + self.drop_path(self.ffn(self.ffn_norm(x)))
+        return x
+
+
+class HRMRefinementBlock(nn.Module):
+    """Iterative zero-init-gated refinement (ported from reference/model_v2.py).
+
+    Polishes the final token sequence; each step re-injects the pre-refinement
+    anchor. Gates init at zero (tanh(0)=0) so it is a safe no-op at start.
+    """
+
+    def __init__(self, hidden_size: int, refine_dim: int, steps: int, eps: float = 1e-6):
+        super().__init__()
+        self.steps = steps
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False)
+        self.up = nn.Linear(refine_dim, hidden_size, bias=False)
+        self.gate = nn.Parameter(torch.zeros(steps))
+        nn.init.normal_(self.down.weight, std=0.02)
+        # up was zero-init too -> with gate=0 AND up=0 neither gets gradient (dead block).
+        # Zero ONLY the scalar gate; a small non-zero up lets the gate receive gradient and
+        # wake up, while the output is still ~0 at init (tanh(0)=0) so it's a safe no-op.
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        anchor = x
+        h = x
+        for t in range(self.steps):
+            inp = torch.cat([self.norm(h), anchor], dim=-1)
+            update = self.up(F.silu(self.down(inp)))
+            h = h + torch.tanh(self.gate[t]) * update
+        return h
+
+
+class JEPAPredictorBlock(nn.Module):
+    """Escarda JEPA head (ported from the Escarda LM's JEPAPredictorBlock, shaped like
+    HRMRefinementBlock). Runs ALONGSIDE the HRM block (not replacing it). From patch
+    token t it predicts the trunk's own token t+k in *representation* space; the target
+    is stop-gradient (detached) — the JEPA asymmetry that prevents collapse. Auxiliary
+    only: used in the training loss, never in the forward output, so inference cost is 0.
+
+    gate starts at zero (identity at init, zero-risk); up is normal-init (not zero) so the
+    gate receives gradient — same double-zero-saddle fix as the HRM block.
+    """
+
+    def __init__(self, hidden_size: int, pred_dim: int, horizon: int, eps: float = 1e-6):
+        super().__init__()
+        self.horizon = horizon
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size, pred_dim, bias=False)
+        self.up = nn.Linear(pred_dim, hidden_size, bias=False)
+        self.gate = nn.Parameter(torch.zeros(horizon))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, h: torch.Tensor, k: int) -> torch.Tensor:
+        update = self.up(F.silu(self.down(self.norm(h))))
+        return h + torch.tanh(self.gate[k - 1]) * update
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.proj = nn.Conv2d(cfg.in_channels, cfg.hidden_size,
+                              kernel_size=cfg.patch_size, stride=cfg.patch_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)                 # (B, C, gh, gw)
+        return x.flatten(2).transpose(1, 2)  # (B, P, C)
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, cfg: VisionConfig | None = None):
+        super().__init__()
+        self.cfg = cfg or VisionConfig()
+        cfg = self.cfg
+
+        self.patch_embed = PatchEmbed(cfg)
+        self.rope_offset = 1 if cfg.use_cls_token else 0
+        if cfg.use_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.hidden_size))
+
+        dpr = [cfg.drop_path * i / max(cfg.num_layers - 1, 1) for i in range(cfg.num_layers)]
+        self.blocks = nn.ModuleList([Block(cfg, dpr[i]) for i in range(cfg.num_layers)])
+        self.hrm_refine = (
+            HRMRefinementBlock(cfg.hidden_size, cfg.hrm_refine_dim,
+                               cfg.hrm_refine_steps, cfg.rms_norm_eps)
+            if cfg.use_hrm_refine else None
+        )
+        self.jepa = (                                    # Escarda: JEPA alongside HRM
+            JEPAPredictorBlock(cfg.hidden_size, cfg.jepa_pred_dim,
+                               cfg.jepa_horizon, cfg.rms_norm_eps)
+            if getattr(cfg, "use_jepa", False) else None
+        )
+        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+        cos, sin = build_axial_rope(cfg.grid_size, cfg.head_dim, cfg.rope_theta,
+                                    device="cpu", dtype=torch.float32)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+        self.apply(self._init_weights)
+        if cfg.use_cls_token:
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        # Re-apply HRM init that the global _init_weights pass would clobber. Gate-only
+        # zero-init (up is non-zero) so the block is a safe no-op at start AND can train.
+        if self.hrm_refine is not None:
+            nn.init.normal_(self.hrm_refine.down.weight, std=0.02)
+            nn.init.normal_(self.hrm_refine.up.weight, std=0.02)
+
+    @staticmethod
+    def _init_weights(m: nn.Module):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, images: torch.Tensor) -> dict:
+        """images: (B, C, H, W) in [-1, 1] or normalized. Returns pooled + tokens."""
+        cfg = self.cfg
+        x = self.patch_embed(images)  # (B, P, C)
+        if cfg.use_cls_token:
+            cls = self.cls_token.expand(x.shape[0], -1, -1)
+            x = torch.cat([cls, x], dim=1)
+
+        cos = self.rope_cos.to(x.dtype)
+        sin = self.rope_sin.to(x.dtype)
+        for blk in self.blocks:
+            x = blk(x, cos, sin, self.rope_offset)
+        if self.hrm_refine is not None:
+            x = self.hrm_refine(x)
+        x = self.norm(x)
+
+        if cfg.use_cls_token:
+            pooled = x[:, 0]
+            tokens = x[:, 1:]
+        else:
+            pooled = x.mean(dim=1)
+            tokens = x
+        return {"pooled": pooled, "tokens": tokens, "last_hidden_state": x}
+
+    def jepa_aux_loss(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Escarda JEPA auxiliary loss on the patch tokens: from token t predict token
+        t+k (raster-order neighbour) in representation space, target detached, 1-cosine
+        averaged over offsets. Returns 0 if JEPA is disabled."""
+        if self.jepa is None:
+            return tokens.new_zeros(())
+        total = tokens.new_zeros(())
+        n = 0
+        for k in range(1, self.jepa.horizon + 1):
+            pred = self.jepa(tokens[:, :-k, :], k)
+            target = tokens[:, k:, :].detach()
+            cos = F.cosine_similarity(pred.float(), target.float(), dim=-1)
+            total = total + (1.0 - cos).mean()
+            n += 1
+        return total / max(n, 1)
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+def _format_millions(n: int) -> str:
+    return f"{n / 1e6:.2f}M"
+
+
+if __name__ == "__main__":
+    cfg = VisionConfig()
+    model = VisionEncoder(cfg)
+    total = model.num_parameters()
+
+    # Per-group breakdown
+    groups = {"patch_embed": 0, "blocks": 0, "hrm_refine": 0, "norm/cls": 0}
+    for name, p in model.named_parameters():
+        if name.startswith("patch_embed"):
+            groups["patch_embed"] += p.numel()
+        elif name.startswith("blocks"):
+            groups["blocks"] += p.numel()
+        elif name.startswith("hrm_refine"):
+            groups["hrm_refine"] += p.numel()
+        else:
+            groups["norm/cls"] += p.numel()
+
+    print(f"VisionEncoder config: dim={cfg.hidden_size} depth={cfg.num_layers} "
+          f"heads={cfg.num_heads} patch={cfg.patch_size} img={cfg.image_size} "
+          f"({cfg.num_patches} patches)")
+    for g, v in groups.items():
+        print(f"  {g:14s}: {_format_millions(v)}")
+    print(f"  {'TOTAL':14s}: {_format_millions(total)}  ({total:,} params)")
+
+    x = torch.randn(2, cfg.in_channels, cfg.image_size, cfg.image_size)
+    out = model(x)
+    print(f"\nforward OK: pooled={tuple(out['pooled'].shape)} "
+          f"tokens={tuple(out['tokens'].shape)}")