| --- |
| license: other |
| tags: |
| - unified-embedding-field |
| - text-to-image |
| - scaling-laws |
| - research-checkpoints |
| library_name: pytorch |
| --- |
| |
| # UEF scaling curve — campaign A2 (d12 / d15 / d28) |
|
|
| Trunk-size scaling rungs for a **unified embedding field** trained jointly over |
| *frozen* understanding representations (SigLIP2 `so400m-patch14-224` image |
| states + `flan-t5-small` text states). Every rung is the same recipe at the same |
| data and batch — **only the trunk width/depth changes**. |
|
|
| These are **research checkpoints**, not a product release. Nothing here is |
| adjudicated: the numbers below are logged validation losses, not a verdict on |
| the shape of the scaling curve. |
|
|
| ## Rungs |
|
|
| | rung | hidden | depth_double | heads | trainable params | steps released | |
| |---|---|---|---|---|---| |
| | `d12` | 768 | 10 | 12 | 178,573,312 | 10k · 20k · 25k · 40k · 45k · 50k · best_val | |
| | `d15` | 960 | 13 | 15 | 334,866,688 | 10k · 20k · 25k · 40k · 45k · 50k | |
| | `d28` | 1792 | 26 | 28 | 2,121,404,416 | 10k · 15k · 20k · 25k | |
|
|
| All rungs share `text_preamble_depth: 2`, `head_dim: 64`, `patch_size: 14`, |
| `pca_channels: 128`, `time_cond: in_context`, `time_tokens: 4`. |
|
|
| ## Shared recipe (identical across rungs) |
|
|
| - global batch **4096**, lr **4e-4**, 25k steps then extended to 50k |
| - bf16 autocast, fp32 master weights + fp32 AdamW moments, fused AdamW |
| - dual EMA (slow + fast), both kept fp32 |
| - union pretrain pool 38.4M pairs, text window 128, `wds_shard_seed` 4242 |
| - 4 nodes x 4 H200, world size 16 |
|
|
| ## Validation loss (logged, final step of each segment) |
|
|
| | rung | step | long | short | jdb | img_t2i (long) | |
| |---|---|---|---|---|---| |
| | d12 | 25,000 | 0.6827 | 0.6681 | 0.6059 | 0.5852 | |
| | d12 | 50,000 | 0.6680 | 0.6539 | 0.5968 | 0.5729 | |
| | d15 | 25,000 | 0.6263 | 0.6125 | 0.5708 | 0.5339 | |
| | d15 | 50,000 | 0.6090 | 0.5951 | 0.5609 | 0.5198 | |
| | d28 | 25,000 | 0.5559 | 0.5443 | 0.5513 | 0.5446 | |
| |
| ## Provenance |
| |
| Each rung's 0→25k segment and its 25k→50k continuation are separate Slurm jobs; |
| the continuation auto-resumes from `checkpoint_025000.pt`. |
|
|
| | rung | segment | run id | slurm job | commit | |
| |---|---|---|---|---| |
| | d12 | 0→25k | `20260818-190255-a2-d12` | 40149805 | `d124497` | |
| | d12 | 25k→50k | `20260821-032904-a2ext-d12` | 40509242 | `fba6fa4` | |
| | d15 | 0→25k | `20260819-065052-a2-d15` | 40149806 | `1bd9777` | |
| | d15 | 25k→50k | `20260821-152650-a2ext-d15` | 40509287 | `fba6fa4` | |
| | d28 | 0→25k | `20260819-183555-a2-d28` | 40149811 | `1bd9777` | |
|
|
| Every file additionally carries its own `identity` block (experiment id, segment |
| id, parent segment id, config/dataset identity sha256, world size, global batch) |
| and an `export_provenance` block naming the exact source checkpoint it came from. |
|
|
| ## Contents of a checkpoint |
|
|
| Weights-only export — the optimizer state and RNG state have been stripped, so |
| these load for evaluation but are **not resumable**. |
|
|
| ```python |
| import torch |
| ck = torch.load("d15/checkpoint_050000.pt", map_location="cpu", weights_only=False) |
| |
| ck["model"] # raw trained weights (fp32) |
| ck["ema"] # slow EMA (fp32) |
| ck["ema_fast"] # fast EMA (fp32) |
| ck["config"] # full training config |
| ck["identity"] # run provenance |
| ck["representation_manifest"] # frozen repr specs + schedules |
| ck["step"], ck["best_val_loss"], ck["export_provenance"] |
| ``` |
|
|
| Tensors are **bitwise identical** to the training checkpoints they were cut |
| from; the export only drops keys, it does not cast or repack. |
|
|
| `d12/best_val.pt` is that run's lowest-validation checkpoint, step 47,000 — |
| a genuinely different point from its step-50,000 checkpoint. |
|
|
| **d15 has no `best_val.pt`.** Its best-validation step *was* 50,000, and the |
| file was verified bitwise identical to `d15/checkpoint_050000.pt`, so the |
| duplicate 4.02 GB of weights is not published. The validation numbers it |
| carried are preserved in `d15/best_val_metrics.json`. |
| |
| ## Caveats — read before using these in a comparison |
| |
| 1. **Single seed.** One run per rung. No seed repeats, so rung-to-rung gaps |
| carry no error bars. |
| 2. **d28 is incomplete.** Its 25k→50k continuation was still running when this |
| was published; only the 0→25k segment is here. `d28/checkpoint_015000.pt` is |
| a rolling checkpoint that the live job would otherwise have deleted. |
| 3. **Segmented continuation.** Resuming reseeds a fresh global data |
| permutation, so the ≥25k segment is not sample-order-aligned with a |
| hypothetical single 50k run. |
| 4. **JourneyDB pool repack.** The `jdb` shards are repacked copies with 14 bad |
| image members and their 14 text partners dropped (28 members total); |
| the jdb manifest count is 4,197,986. Composition-identity against the other |
| site's copy of the pool was not closed. |
| 5. **`d24` is absent** — that rung had not produced a usable checkpoint. |
| 6. Validation losses above are read from the run logs, are computed on small |
| val batches, and are the training-time metric only. No FID / GenEval / DPG |
| numbers are attached to these rungs. |
|
|
| ## Configs |
|
|
| `configs/f25k_scale_d{12,15,28}.yml` are the exact configs used, verbatim. |
|
|
