A Night at the LLM Zoo: Data Conditioned Flow Matching for Generating Pretrained LLM Weights
Organizations: Khoury College of Computer Science, Northeastern University, Silicon Valley · MLAI Lab, KAIST · Khoury College of Computer Science, Northeastern University, Boston · Slater Lab, University of Oxford ·
Abstract
A challenge to Weight Space Learning as it expands from task specific and vision models to Large Language Models (LLMs) is the absence of a well curated and analyzed LLM population. To address this, we train such a model zoo for application to Weight Space Generation by systematically varying the data mixture composition between web text, books, mathematics, coding, and French across populations of 51M, 124M, and 355M parameter GPT-2 Models, each tier with count N = 100. We analyze these populations by cosine similarity, functional diversity, Leave-One-Mixture-Out modeling, and variance decomposition. We show that the zoo contains diverse models whose variation is structured by training mixture composition. We embed the training mixture composition to condition a Flow Matching model of complete pretrained LLM weights in Gram-PCA compressed weight space. By demonstrating zero-shot generalization to held out mixtures, measuring Continuous Ranked Probability Score skill, as well as finding minimal changes in downstream test-set perplexity and general capability benchmark scores, we show that our generator generalizes on the simplex of data mixtures to produce models that perform similarly to those in the zoo ex nihilo. Finally, we produce all source code and artifacts to support the continued expansion of Weight Space Learning and Generation into the LLM domain.
Figures & tables
| Mini (51.5M) | Small (124.4M) | Medium (354.8M) | |||||
|---|---|---|---|---|---|---|---|
| pool | pairs | raw | drift rem. | raw | drift rem. | raw | drift rem. |
| within-anchor (same ) | 30 | ||||||
| between-anchor (different ) | 160 | ||||||
| singleton–singleton | 3160 | ||||||
| anchor–singleton | 1600 | ||||||
| all pairs | 4950 | ||||||
| Mini | Small | Medium | |||||||
|---|---|---|---|---|---|---|---|---|---|
| region | of | of | of | ||||||
| whole model | |||||||||
| embeddings | |||||||||
| blocks only | |||||||||
| pool | Mini | Small | Medium |
|---|---|---|---|
| within-anchor (same ) | |||
| between-anchor (different ) | |||
| singleton–singleton | |||
| anchor–singleton | |||
| all pairs | |||
| between/within anchor |
| Mini (51M) | Small (124M) | Medium (355M) | ||||
| source of weights | % PPL | ceil | % PPL | ceil | % PPL | ceil |
| replicate ceiling (seed noise) | 1.01 | 1.0 | 1.09 | 1.0 | 0.97 | 1.0 |
| rank truncation floor | 1.80 | 1.8 | 1.93 | 1.8 | 2.46 | 2.5 |
| flow (ours) | 4.05 | 4.0 | 4.17 | 3.8 | 5.11 | 5.3 |
| retrieval (1-NN, matched blinding) | 5.75 | 5.7 | 5.79 | 5.3 | 6.01 | 6.2 |
| quadratic surface on | 6.54 | 6.5 | 7.27 | 6.7 | 7.88 | 8.1 |
| stratum | real | flow | flow real ( ) | flow rank ( ) | |
|---|---|---|---|---|---|
| Mini (51M) | interior | ( ) | ( ) | ||
| vertices | ( ) | ( ) | |||
| Small (124M) | interior | ( ) | ( ) | ||
| vertices | ( ) | ( ) | |||
| Medium (355M) | interior | ( ) | ( ) | ||
| vertices | ( ) | ( ) |
| Pool | Mini | Small | Medium |
|---|---|---|---|
| Flow vs. its target | 0.1322 | 0.1331 | 0.1452 |
| 1-NN retrieval vs. target | 0.1680 | 0.1596 | 0.1651 |
| Flow vs. other | 0.2056 | 0.1954 | 0.1967 |
| Rank-50 floor | 0.0931 | 0.0891 | 0.0935 |
| Flow cohort spread | 0.0825 | 0.1018 | 0.1257 |
| CRPS, % PPL | retrieval flow, % PPL | marginal flow, CRPS | |||||||
|---|---|---|---|---|---|---|---|---|---|
| flow | marginal | skill ( ) | gap | won | gap | won | |||
| Mini (51M, n=32) | 3.10 | 8.34 | +0.628 | 61/85 | 84/85 | ||||
| Small (124M, n=32) | 3.26 | 8.38 | +0.611 | 61/85 | 83/85 | ||||
| Medium (355M, n=16) | 3.51 | 8.83 | +0.603 | 43/85 | 82/85 | ||||
| round-trip rel. | rel. vs. 200-step | mean pairwise dist. | |||||||
|---|---|---|---|---|---|---|---|---|---|
| steps | Mini | Small | Med. | Mini | Small | Med. | Mini | Small | Med. |
| 1 | 1.0077 | 0.7129 | 0.5826 | 0.4158 | 0.3786 | 0.3735 | 2.306 | 3.209 | 4.558 |
| 5 | 0.2413 | 0.1552 | 0.1208 | 0.0920 | 0.0715 | 0.0534 | 1.725 | 2.732 | 4.087 |
| 50 | 0.0281 | 0.0166 | 0.0122 | 0.0062 | 0.0050 | 0.0037 | 1.922 | 2.889 | 4.139 |
| 200 | 0.0070 | 0.0042 | 0.0030 | 0.0000 | 0.0000 | 0.0000 | 1.942 | 2.904 | 4.145 |
Appendix figures & tables29 assets
Supplementary material from the paper’s appendix.
Appendix
| Mini | Small | Medium | |
| embedding share of | |||
| tokens per block parameter |
| Optimizer | AdamW ( Loshchilov and Hutter, 2019 ) |
|---|---|
| Weight decay | , applied uniformly |
| Gradient clipping | global norm |
| Precision | bf16 mixed |
| Peak learning rate | |
| Warmup | linear, first of steps |
| Scale | total | trunk | branch | |
|---|---|---|---|---|
| Mini | ||||
| Mini | ||||
| Small | ||||
| Medium |
| Domain | Dataset | Config | Column |
|---|---|---|---|
| web | HuggingFaceFW/fineweb-edu ( Penedo et al., 2024 ) | sample-10BT | text |
| code † | codeparrot/codeparrot-clean ( CodeParrot, ) | — | content |
| math | open-web-math ( Paster et al., 2024 ) | — | text |
| books | manu/project_gutenberg ( manu, ) | — | text |
| multilingual ‡ | HuggingFaceFW/fineweb-2 ( Penedo et al., 2025 ) | fra_Latn | text |
| between/within distance ratio | |||
|---|---|---|---|
| displacement from trunk |
| Mini | Small | Medium | ||||
|---|---|---|---|---|---|---|
| within | between | within | between | within | between | |
| worst interior barrier | ||||||
| mean dip below chord | ||||||
| midpoint best endpoint | ||||||
| midpoint beats both endpoints | ||||||
| Rank , population statistics | |
|---|---|
| Rank , generative experiments | (Appendix B.2 ) |
| One-dimensional parameters | retained ( of ) |
| Code normalization | one global scalar to unit training RMS |
| Synthetic weight-space noise | none |
| epochs | |||
|---|---|---|---|
| — |
| Mini | Small | Medium | |
|---|---|---|---|
| flow round-trip | |||
| truncation floor | |||
| generator’s own excess |
| Block | share | within/ens. | of all within | ||
|---|---|---|---|---|---|
| c0–c3 | |||||
| c4 | |||||
| c5–c9 | |||||
| c10–c24 | |||||
| c25+ |
| Mini (8 blk) | Small (12 blk) | Medium (24 blk) | ||||
| region, | eff. rank | var. share | eff. rank | var. share | eff. rank | var. share |
| whole stack | ||||||
| blocks only | ||||||
| embeddings only | ||||||
| same eigenvalues, embeddings only, read at lower : | ||||||
| Mini | Small | Medium | |||||
|---|---|---|---|---|---|---|---|
| region | design | interior | vertex | interior | vertex | interior | vertex |
| whole model | |||||||
| quadratic | |||||||
| blocks only | |||||||
| Mini | Small | Medium | |||||
|---|---|---|---|---|---|---|---|
| region | design | interior | vertex | interior | vertex | interior | vertex |
| embeddings | |||||||
| quadratic | |||||||
| function space | |||||||
| Mini (51M) | Small (124M) | Medium (355M) | |||||||
| stratum | flow | retr. | gap | flow | retr. | gap | flow | retr. | gap |
| lower tercile ( ) | 2.60 | 3.63 | +1.03 | 2.21 | 3.39 | +1.18 | 3.20 | 3.72 | +0.52 |
| middle tercile ( ) | 2.38 | 3.94 | +1.56 | 2.48 | 3.61 | +1.13 | 3.24 | 3.66 | +0.41 |
| upper tercile ( ) | 7.40 | 9.96 | +2.55 | 8.10 | 10.71 | +2.61 | 9.18 | 10.99 | +1.81 |
| interior ( ) | 2.66 | 4.43 | +1.77 | 2.57 | 4.26 | +1.68 | 3.39 | 4.37 | +0.98 |
| one-hot vertices ( ) | 26.27 | 26.83 | +0.56 | 29.79 | 30.38 | +0.59 | 32.65 | 32.23 | -0.42 |
| CRPS, % PPL | retrieval flow, % PPL | marginal flow, CRPS | |||||||
|---|---|---|---|---|---|---|---|---|---|
| flow | marginal | skill | gap | won | gap | won | |||
| Mini (51M) | 3.10 | 8.34 | +0.628 | 61/85 | 84/85 | ||||
| Small (124M) | 3.26 | 8.38 | +0.611 | 61/85 | 83/85 | ||||
| Medium (355M) | 3.51 | 8.83 | +0.603 | 43/85 | 82/85 | ||||
| Type | rectified flow ( Liu et al., 2023a ; Lipman et al., 2023 ) |
|---|---|
| Width | |
| Parameters | k |
| Integration | explicit Euler steps |
| Conditioning | -layer SiLU MLP, concatenated ( Perez et al., 2018 ) |
| Conditioning dropout | ( Ho and Salimans, 2021 ) |
| Guidance scale at evaluation |
| Mini | Small | |||
|---|---|---|---|---|
| pooled | pooled | |||
| (committed) | ||||
| round-trip error | vs. 200-step sample | dispersion | ||||
|---|---|---|---|---|---|---|
| steps | rel. | rel. | cosine | sample RMS | mean pair dist. | |
| 1 | 0.5902 | 0.590 | 0.2772 | 0.9617 | 1.050 | 4.432 |
| 2 | 0.3116 | 0.623 | 0.1346 | 0.9911 | 1.062 | 4.046 |
| 5 | 0.1272 | 0.636 | 0.0464 | 0.9989 | 1.058 | 3.982 |
| 10 | 0.0633 | 0.633 | 0.0217 | 0.9998 | 1.055 | 3.996 |
| 25 | 0.0252 | 0.630 | 0.0079 | 1.0000 | 1.054 | 4.011 |
| steps | PPL % | CRPS % | CRPS skill | |
|---|---|---|---|---|
| 1 | 4.365 | 3.134 | 0.6259 | 0.839 |
| 10 | 4.165 | 3.252 | 0.6119 | 0.673 |
| 50 | 4.173 | 3.231 | 0.6143 | 0.668 |
| Mini | Small | Medium | |
| lower bound vs. matched twin (nats) | |||
| identification, top- (chance ) | |||
| identification, top- | |||
| median rank of true (of ) | |||
| tie fraction | |||
| C2ST, joint |
| Mini (51M) | Small (124M) | Medium (355M) | ||||
| domain | real zoo | flow | real zoo | flow | real zoo | flow |
| web | ||||||
| code | ||||||
| math | ||||||
| books | ||||||
| multilingual | ||||||
| pool | Mini | Small | Medium |
|---|---|---|---|
| flow vs. its target | 0.3036 | 0.2991 | 0.2939 |
| 1-NN retrieval vs. target | 0.3136 | 0.2995 | 0.2887 |
| flow vs. other | 0.2744 | 0.2503 | 0.2427 |
| replicate floor (same ) | 0.1657 | 0.1791 | 0.1892 |
| rank-50 floor | 0.0192 | 0.0133 | 0.0120 |
| flow cohort spread | 0.1302 | 0.1343 | 0.1585 |
| scale | gen tr | test tr | ||
|---|---|---|---|---|
| Mini | ||||
| Medium | ||||
| stratum | real | flow | ||
|---|---|---|---|---|
| Mini (51M) | interior | |||
| vertices | ||||
| Small (124M) | interior | |||
| vertices | ||||
| Medium (355M) | interior | |||
| vertices |