srt-sunstone-linear-head

SRT-Sunstone Linear Head

Train once, read everywhere. This head was trained once on datacenter bf16 states and reads the same structure unchanged across a 10× host-scale reduction, 4-bit quantization (−0.01 R@1), and a change of silicon (CUDA → Apple-Silicon MLX: 97 % head-space text agreement against a 99.96 % same-runtime ceiling, and on-device image→text retrieval within 3 R@1 points of the datacenter reference). One artifact, every deployment tier from edge to fleet.

A single trained linear projection pair that turns frozen gemma-4-31B-it into an image↔text retrieval engine. No backbone weights are touched; the head reads layer-47 hidden states from the forward pass you are already running.

Protocol: 1,000 held-out COCO val2017 images vs their 5,000 captions (chance R@1 ≈ 0.001); training pairs from train2017 (cross-partition). An MLP trained on the same data never beats this linear map at any training size, and the gap widens with data: the cross-modal correspondence in this substrate is linear (paper §11.6.4, space-bacon/SRT).

Contents

Two dtypes of the same head. sunstone_linear_head_v3_drift.pt is fp32 at 44.1 MB; sunstone_linear_head_v3_drift_bf16.pt is 22.0 MB and is what the "22 MB head" figure in the papers refers to. The export is measured, not assumed: on 5,000 L47 pairs the two agree to cosine 0.999997 (minimum 0.999992) and text-to-image R@1 moves from 0.3016 to 0.3018 with median rank 4 unchanged. Receipt in bf16_head_equivalence.json, produced by scripts/bf16_head_equivalence.py. Scope: that pool is 5,000 images, well short of the shipped 123,287 gallery, so those are agreement figures and not deployment numbers.

sunstone_linear_head.pt (weights_only=True-safe):

• img, txt — state dicts for two nn.Linear(5376, 1024) heads
• mu_img, mu_txt — per-modality centering means (mandatory; raw cosine is dominated by the anisotropy attractor)
• meta — training config + eval numbers

sunstone_linear_head_v2_jitter30.pt — same format, same protocol, trained with mean-jitter augmentation (each batch perturbed by a random constant displacement per modality, norm ~U(0, 30), covering the ~22-unit mean displacement a runtime change actually produces).

sunstone_linear_head_v3_drift.pt — same format, trained with measured drift-family augmentation: each training sample is perturbed by a real per-vector cross-runtime residual (MLX-Q4 minus CUDA-bf16, 3,000 measured pairs, scaled ~U(0, 1.5)). On held-out cross-runtime data this internalizes the runtime drift: no-recal i2t R@1 0.636 (v1 needs the recalibration to reach 0.634), t2i 0.469 vs v1's 0.424 ceiling, and clean same-runtime performance improves (i2t 0.679 vs 0.661). The residual augmentation acts as a regularizer.

Which to use. The three heads form a ladder of target-runtime knowledge:

The per-runtime mean recalibration (average ~200 states from your own runtime, use instead of the shipped mu_*) is still recommended with every head; with v1 it is load-bearing (skipping it costs ~24 i2t R@1 points, a failure mode invisible to same-transform agreement metrics and only visible on end-task retrieval). The drift analysis behind all of this is in artifacts/nla/q4/*_2026080[56].json in the repo, developed through a public five-round review exchange on the release thread.

Usage

import torch, torch.nn.functional as F

ck = torch.load("sunstone_linear_head.pt", weights_only=True)
head_i = torch.nn.Linear(5376, 1024); head_i.load_state_dict(ck["img"])
head_t = torch.nn.Linear(5376, 1024); head_t.load_state_dict(ck["txt"])

# h_img: mean L47 state over image soft-token positions
# h_txt: BOS-prefixed last-token L47 state
zi = F.normalize(head_i(h_img - ck["mu_img"]), dim=-1)
zt = F.normalize(head_t(h_txt - ck["mu_txt"]), dim=-1)
score = zi @ zt.T

Calibration notes (read before deploying)

1. Center per-modality, always.
2. Anchor domain match beats anchor size: recalibrate mu_img on ~150+ images from your deployment's own query domain.
3. Prepend BOS on every text encode (gemma-4 is BOS-sensitive).
4. Retrieval readout needs a single prefill pass, not generation.

Training recipe, deployment tiers, and reproduction scripts: docs/CROSSMODAL_LINEAR_HEAD.md in space-bacon/SRT. Pair encodings for retraining: RiverRider/srt-nla-gemma4-artifacts (procrustes/). Trained on COCO captions (CC-BY 4.0).