README.md

metadata

license: apache-2.0
base_model: google/gemma-4-31B-it
tags:
  - cross-modal
  - retrieval
  - srt
  - frozen-backbone

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.

method	i2t R@1	i2t R@5	i2t R@10	t2i R@1
centered cosine (zero training)	0.288	0.523	0.648	0.173
this head (117k COCO pairs, InfoNCE)	0.661	0.911	0.967	0.506

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

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:

you have	use	cross-runtime behavior
nothing	v3	i2t 0.636 / t2i 0.469 with zero calibration
~200 unpaired states from your runtime (42KB mean recal)	v3 + recal	i2t 0.658, within 0.2 points of same-runtime
generic caution, no target runtime known	v2	degrades most gracefully across arbitrary runtimes
exact v1 reproduction	v1 + recal	the paper's reference numbers

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).