UEF wave 2 — paradigm baselines and task-mixture arms (50k)
A unified embedding field trained jointly over frozen understanding representations (SigLIP2 image states + flan-T5 text states). Wave 2 varies the objective paradigm and the task mixture at a fixed trunk, fixed data pool and a fixed recipe, so that each arm differs from the others in one declared way.
All arms are sample-paired (seed 4242, wds_shard_seed 4242) with each other and
with the 50k arms of
Ced-Collab/uef-scaling-a2.
These are research checkpoints, not a product release, and nothing here is adjudicated.
⚠️ Losses across arms are NOT comparable
The arms in this repo do not optimise the same quantity. The AR-text arm's text loss is an autoregressive cross-entropy over flan-t5-small token ids (nats per token); the flow arms' text loss is a flow-matching regression in T5 latent space. They have different units and different floors. The reported
val/lossis a sum of an image flow loss and a text loss, so it inherits that incomparability.A paradigm comparison needs decode-side or downstream metrics (GenEval, DPG, VQA, caption). None of those are attached here.
Arms
| arm | image side | text side | status |
|---|---|---|---|
text_token_ar |
continuous SigLIP2 latents, flow matching | autoregressive CE over flan-t5-small ids | ✅ released (50k) |
disc_disc |
amused-256 VQ codes, masked diffusion CE | masked-diffusion CE over t5 ids | ⏳ training |
a2n |
continuous, flow | flow, mixture [0, .5, .5] |
⏳ queued |
i2t |
— (no image-target rows) | flow, mixture [0, 0, 1] |
⏳ queued |
Arms are added to this repo as they finish; the table above is the source of truth for what is actually present.
text_token_ar — paradigm cell (b)
Continuous image flow + autoregressive text. The single declared delta against the
flow-text control is text_objective_mode: token_ar.
Trunk and representations
| trunk | d19: hidden_size 1216, depth_double 17, 19 heads |
| image encoder (frozen) | google/siglip2-so400m-patch14-224, latent dim 1152, 224 px |
| text encoder (frozen) | google/flan-t5-small, hidden 512, txt_len 128 |
| image objective | flow matching, image_prediction: x, loss space v, t_shift 8.485 |
| text objective | autoregressive CE over t5 ids, tied embeddings, text_t_shift 1.0 |
| task mixture | [0.556, 0.222, 0.222] (product / t2i / i2t), stratified |
Recipe
Batch 4096 (micro 128 × 16 ranks × grad-accum 2), 50,000 steps, lr 4e-4 with 5,000 warmup steps, weight decay 0.01, dual EMA (0.99995 slow / 0.9996 fast), bf16 autocast with fp32 master weights and fp32 EMA. Data pool: the 38.4M-sample union of BLIP3o long/short captions and JourneyDB. COCO val/test never enters training.
Validation trajectory
| step | val/loss |
val/loss_txt_ce |
val/loss_img_t2i |
|---|---|---|---|
| 10,000 | 0.8108 | 0.9473 | 0.5700 |
| 20,000 | 0.7575 | 0.8531 | 0.5399 |
| 30,000 | 0.7341 | 0.8116 | 0.5263 |
| 40,000 | 0.7227 | 0.7907 | 0.5203 |
| 50,000 | 0.7141 | 0.7743 | 0.5164 |
Monotone through 50k with no reversal; best validation is the final step. Full
per-pool numbers (long / short / jdb) are in text_token_ar/best_val_metrics.json.
Files
text_token_ar/checkpoint_010000.pt weights-only export
text_token_ar/checkpoint_025000.pt
text_token_ar/checkpoint_050000.pt
text_token_ar/best_val_metrics.json best-val block (step 50000)
text_token_ar/precision_contract.json recorded precision contract
configs/f50k_text_token_ar_2x8_kmp.yml the exact config used
logs/text_token_ar_rank0.log full rank-0 training log
logs/text_token_ar_manifest.yaml run id, job id, nodelist, commit
Checkpoints are weights-only: model, ema, ema_fast, config, identity and
representation manifest are kept; optimizer moments and RNG state are dropped. They
load for evaluation but are not resumable. Dtypes are preserved verbatim (fp32).
Provenance
| run id | 20260906-012208-w2-text_token_ar |
| scheduler job | kempner slurm 44743117, partition kempner_h200 |
| hardware | 4 nodes × 4 H200 = world 16 |
| wall time | 14 h 13 m to 50k, 0 nan |
| code | tag kempner-wave2-v1 (3d1c5f3) + site delta, commit 48d65d7 |
| config | configs/kempner_wave2/kempner/f50k_text_token_ar_2x8_kmp.yml |
The site delta over the hand-off tag changes deployment keys only — data manifest
paths for this cluster's pool copy, a JourneyDB sample count corrected for the
2026-08-18 repack, the W&B entity, and checkpoint retention. Task mixture, seeds,
t_shift, trunk shape and step count are untouched.
Loading
import torch
ck = torch.load("text_token_ar/checkpoint_050000.pt", map_location="cpu", weights_only=False)
ck["step"] # 50000
ck["model"] # trunk weights
ck["ema"] # slow EMA (the one to evaluate)
ck["ema_fast"] # fast EMA
ck["config"] # full training config
Evaluate ema unless you have a reason to prefer the raw model weights.