Ced-Collab/uef-wave2-paradigm

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/loss is 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	✅ released (50k)
a2n	continuous, flow	flow, mixture [0, .5, .5]	✅ released (50k)
i2t	— (no image-target rows)	flow, mixture [0, 0, 1]	✅ released (50k)

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.

disc_disc — paradigm cell (a)

Fully discrete: the image is a grid of VQ code indices predicted by masked diffusion, and the text is masked-diffusion CE over t5 ids. This is the arm whose 10k checkpoint was pre-registered as a binding tripwire.

Trunk and representations

trunk	d19: hidden_size 1216, depth_double 17, 19 heads (same as text_token_ar)
image codec (frozen)	amused/amused-256, subfolder vqvae — 256 px → 16×16 codes, codebook 8192
image objective	masked-diffusion CE over code indices, image_code_embed_std 0.9861, t_shift 1.0
text encoder (frozen)	google/flan-t5-small, txt_len 128
text objective	masked-diffusion CE (token_discrete), text_t_shift 1.0
task mixture	[0.556, 0.222, 0.222], stratified — identical to text_token_ar

Validation trajectory

Chance for an 8192-way code prediction is ln 8192 = 9.01 nats and 0.012 % accuracy. Both image metrics move off chance immediately and improve monotonically through 50k:

step	val/loss	val/loss_img_ce	val/img_acc_masked
10,000	7.7826	7.1040	3.57 %
20,000	7.5479	6.9167	4.28 %
30,000	7.4410	6.8321	4.59 %
40,000	7.3888	6.7961	4.68 %
50,000	7.3493	6.7658	4.89 %

Final text side: val/loss_txt_ce 2.3382.

On the tripwire. The stop condition as written — image CE flat at 9.0 for more than 1k steps — was never met: CE was already 8.8 within the first hundred steps and reached 7.1 by 10k. The image side is learning. Whether it is learning enough is a separate question this repo does not answer: masked-code accuracy is 4.9 % at 50k, which is ~400× chance but low in absolute terms, and the validation curve is still descending at the final step. The own-space binding screen is the instrument that adjudicates that, and it is not attached here.

Because the two arms optimise different image objectives — a flow regression in SigLIP2 latent space versus a cross-entropy over VQ indices — val/loss 7.35 here and 0.71 for text_token_ar are not comparable numbers.

a2n — task-mixture arm [0, .5, .5]

Same continuous-flow paradigm as the M0 recipe, but the joint (product) rows are removed: the mixture is half text-to-image and half image-to-text, with no rows that carry both targets at once. It isolates what the joint rows contribute, against the [0.556, 0.222, 0.222] mixture used by the two paradigm arms above.

trunk	d19: hidden_size 1216, depth_double 17, 19 heads
image side	continuous SigLIP2 latents, flow matching, 224 px / patch 14, t_shift 8.485
text side	flow matching in T5 latent space (text_objective_mode: separate), text_t_shift 1.0
task mixture	[0.0, 0.5, 0.5] (product / t2i / i2t), stratified

Validation trajectory

step	val/loss	val/loss_img_t2i	val/loss_txt_i2t
10,000	0.6156	0.5486	0.2504
20,000	0.5795	0.5172	0.2310
30,000	0.5640	0.5038	0.2220
40,000	0.5547	0.4953	0.2190
50,000	0.5496	0.4909	0.2169

Monotone through 50k, best validation at the final step.

This arm is loss-comparable with i2t and with the mixture arms of the scaling curve — same objective, same latent spaces, same units — but it reports no loss_img_product term, since it trains no product rows. It is not comparable with text_token_ar or disc_disc, whose text and image objectives differ.

i2t — task-mixture arm [0, 0, 1]

The same continuous-flow paradigm, reduced to image-to-text only: every row carries a text target conditioned on the image, and there are no text-to-image or product rows. Against a2n it isolates the i2t direction on its own; against the mixture arms it is the pure-i2t endpoint of the task-mixture axis.

trunk	d19: hidden_size 1216, depth_double 17, 19 heads
image side	frozen SigLIP2 conditioning only, 224 px / patch 14 (no image-target rows)
text side	flow matching in T5 latent space (text_objective_mode: separate), text_t_shift 1.0
task mixture	[0.0, 0.0, 1.0] (product / t2i / i2t), stratified

Validation trajectory

step	val/loss	val/loss_txt_i2t
10,000	0.0599	0.2237
20,000	0.0555	0.2065
30,000	0.0533	0.1981
40,000	0.0524	0.1943
50,000	0.0520	0.1928

Monotone through 50k. Unlike the other arms, best validation is step 48000 (val/loss 0.05185), not the final step — 49k and 50k were marginally higher — so the best weights are shipped separately as best_val.pt. The arm reports no loss_img_t2i term, since it trains no image-target rows.

This arm is loss-comparable with a2n and with the mixture arms of the scaling curve — same objective, same latent spaces, same units. It is not comparable with text_token_ar or disc_disc.

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

disc_disc/checkpoint_010000.pt         weights-only export (the tripwire release)
disc_disc/checkpoint_025000.pt
disc_disc/checkpoint_050000.pt
disc_disc/best_val_metrics.json        best-val block (step 50000)
disc_disc/precision_contract.json

a2n/checkpoint_010000.pt               weights-only export
a2n/checkpoint_025000.pt
a2n/checkpoint_050000.pt
a2n/best_val_metrics.json              best-val block (step 50000)
a2n/precision_contract.json

i2t/checkpoint_010000.pt               weights-only export
i2t/checkpoint_025000.pt
i2t/checkpoint_050000.pt
i2t/best_val.pt                        weights-only best-val (step 48000, != 50k)
i2t/best_val_metrics.json              best-val block (step 48000)
i2t/precision_contract.json

configs/f50k_text_token_ar_2x8_kmp.yml the exact configs used
configs/f50k_disc_disc_2x8_kmp.yml
configs/tmx_a2n_b4096_50k_kmp.yml
configs/tmx_i2t_b4096_50k_kmp.yml
logs/text_token_ar_rank0.log           full rank-0 training logs
logs/disc_disc_rank0.log
logs/a2n_rank0.log
logs/i2t_rank0.log
logs/*_manifest.yaml                   run id, job id, nodelist, commit

Reproducing disc_disc needs the amused/amused-256 vqvae weights (diffusion_pytorch_model.safetensors, sha256 1241a5c8…); they are not redistributed here.

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

	text_token_ar
---	---
run id	20260906-012208-w2-text_token_ar
scheduler job	kempner slurm 44743117
wall time to 50k	14 h 13 m, 0 nan

All arms ran on 4 nodes × 4 H200 = world 16, partition kempner_h200, from code tag kempner-wave2-v1 (3d1c5f3) plus the site delta, commit 48d65d7.

Both arms were certified before launch by a 100-step env-cert twin on the identical 2×8-equivalent production shape.

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.