A Night at the LLM Zoo: Data Conditioned Flow Matching for Generating Pretrained LLM Weights

Scott Biggs; Bedionita Soro; Saumya Gupta; Bruno Andreis; Sung Ju Hwang

cs.LGNov 4, 2026

A Night at the LLM Zoo: Data Conditioned Flow Matching for Generating Pretrained LLM Weights

Authors: Scott Biggs, Bedionita Soro, Saumya Gupta, Bruno Andreis, Sung Ju Hwang

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 $\pi$ )	30	$0.9153$	$0.8920$	$0.9024$	$0.8717$	$0.8505$	$0.8000$
between-anchor (different $\pi$ )	160	$0.2564$	$-0.0697$	$0.2743$	$-0.0468$	$0.2423$	$-0.0160$
singleton–singleton	3160	$0.6047$	$0.1721$	$0.6172$	$0.1193$	$0.6060$	$0.0844$
anchor–singleton	1600	$0.3534$	$-0.1381$	$0.3568$	$-0.1276$	$0.3132$	$-0.1167$
all pairs	4950	$0.5141$	$0.0684$	$0.5236$	$0.0387$	$0.5011$	$0.0205$

Table 1: Branch delta cosine similarity by pool, raw and after removing the drift inherited from the shared trunk equation 2 . The pools are disjoint. The random-direction nulls are

\leq\mathcal{O}(1\text{E-4})

across scales.

	Mini			Small			Medium
region	$\eta^{2}$	$\eta^{2}_{\text{corr}}$	of $D$	$\eta^{2}$	$\eta^{2}_{\text{corr}}$	of $D$	$\eta^{2}$	$\eta^{2}_{\text{corr}}$	of $D$
whole model	$0.9416$	$0.9224$	$100\%$	$0.9220$	$0.8965$	$100\%$	$0.8661$	$0.8230$	$100\%$
embeddings	$0.9455$	$0.9276$	$51.0\%$	$0.9451$	$0.9271$	$31.7\%$	$0.9377$	$0.9172$	$14.8\%$
blocks only	$0.9112$	$0.8822$	$49.0\%$	$0.8754$	$0.8353$	$68.3\%$	$0.8186$	$0.7610$	$85.2\%$

Table 2: One-way between/within decomposition over the twenty anchors (

d=5

mixtures

\times

s=4

replicate branches,

\mathrm{df}=4/15

), by coordinate region;

\eta^{2}

and

\eta^{2}_{\text{corr}}

are equation 4 and equation 5 . “of

D

” is the region’s share of the parameter vector, which shifts from embeddings to blocks with scale.

pool	Mini	Small	Medium
within-anchor (same $\pi$ )	$0.166$	$0.179$	$0.189$
between-anchor (different $\pi$ )	$0.411$	$0.396$	$0.373$
singleton–singleton	$0.195$	$0.183$	$0.184$
anchor–singleton	$0.345$	$0.328$	$0.312$
all pairs	$0.250$	$0.236$	$0.232$
between/within anchor	$2.48$	$2.21$	$1.97$

Table 3: Pairwise next-token disagreement: fraction of

327{,}680

held-out positions at which two members predict a different argmax token, over the same four disjoint pools as Table 1 . Bold marks the extreme pools of each column.

	Mini (51M)		Small (124M)		Medium (355M)
source of weights	% $\Delta$ PPL	$\times$ ceil	% $\Delta$ PPL	$\times$ ceil	% $\Delta$ PPL	$\times$ ceil
replicate ceiling (seed noise)	1.01	1.0	1.09	1.0	0.97	1.0
$k=50$ 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 $\pi$	6.54	6.5	7.27	6.7	7.88	8.1

Table 4: Every source of weights on the same

85

held-out mixtures, in

\%\Delta\mathrm{PPL}

over the five-domain profile, averaged per mixture. The replicate ceiling (two real models at the same

\pi

) and

k=50

rank truncation provide an estimated lower bound. Retrieval is blinded to the same neighbourhood as the generator (see Table 30 for vertex mixtures.).

	stratum	real	flow	flow $-$ real ( $t$ )	flow $-$ rank $50$ ( $t$ )
Mini (51M)	interior	$-5.198$	$-5.178$	$+0.0204$ ( $+3.40$ )	$+0.0074$ ( $+1.62$ )
	vertices	$-5.228$	$-5.140$	$+0.0877$ ( $+0.77$ )	$+0.0870$ ( $+0.76$ )
Small (124M)	interior	$-4.685$	$-4.648$	$+0.0360$ ( $+4.27$ )	$+0.0085$ ( $+1.52$ )
	vertices	$-4.771$	$-4.611$	$+0.1601$ ( $+0.97$ )	$+0.1597$ ( $+0.97$ )
Medium (355M)	interior	$-4.148$	$-4.141$	$+0.0068$ ( $+0.62$ )	$-0.0150$ ( $-3.12$ )
	vertices	$-4.273$	$-4.101$	$+0.1720$ ( $+0.86$ )	$+0.1720$ ( $+0.86$ )

Table 5: Gold-continuation log score on HellaSwag,

2{,}000

questions, zero shot, ten held-out mixtures per scale, in nats/token; higher is better. real is the members pretrained at that mixture; flow is the leave-one-mixture-out generator.

\Delta

columns are paired over the five mixtures of a stratum and give the mean difference and its

t

(

\mathrm{df}=4

), with

|t|\geq 2

the separation threshold used throughout. rank

50

is the same real members’ codes truncated to rank

50

, the reconstruction floor, which involves no generator. The accuracy view of the same runs is Table 32 .

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 $\pi$	0.2056	0.1954	0.1967
Rank-50 floor	0.0931	0.0891	0.0935
Flow cohort spread	0.0825	0.1018	0.1257

Table 6: Next-token disagreement of generated models on interior pools (vertices are never pooled with the interior). Flow samples disagree with their target more than the rank-50 floor but less than 1-NN retrieval and other

\pi

, and less than original zoo members (Table 3 ). Vertex results are in Table 30 .

	CRPS, % $\Delta$ PPL			retrieval $-$ flow, % $\Delta$ PPL			marginal $-$ flow, CRPS
	flow	marginal	skill ( $\uparrow$ )	gap	$t$	won	gap	$t$	won
Mini (51M, n=32)	3.10	8.34	+0.628	$+1.70$	$+4.63$	61/85	$+5.23$	$+10.27$	84/85
Small (124M, n=32)	3.26	8.38	+0.611	$+1.62$	$+5.38$	61/85	$+5.12$	$+11.02$	83/85
Medium (355M, n=16)	3.51	8.83	+0.603	$+0.89$	$+3.13$	43/85	$+5.32$	$+12.44$	82/85

Table 7: All

85

held-out mixtures, one LOMO generator per cell, with rows indicating

n

samples/cell. CRPS equation 7 scores the whole cohort against the one real model at that mixture, which a round-trip error cannot do without being minimized by a point mass (Appendix B.7 ). Absolute round-trip errors are in Table 4 .

Figure 1: Advantage over blinded retrieval (a) and the proper score (b) against $L_{1}$ distance to the nearest mixture seen by the generator. Filled markers are equal-count bins over the $80$ interior cells with binned standard errors. The five vertices are open squares and are excluded from binning.

	round-trip rel. $L_{2}$			rel. $L_{2}$ 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

Table 8: An ablation on the number of Euler integration steps taken in generation of LOMO flows across all scales. Here round-trip error indicates the backward integration, followed by the reverse integration. Mean pairwise distances are reported as a dispersion reporting tool, additional metrics can be found in Table 26 .

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

	Mini	Small	Medium
$n_{\text{embd}}$	$512$	$768$	$1024$
$n_{\text{layer}}$	$8$	$12$	$24$
$n_{\text{head}}$	$8$	$12$	$16$
$D$	$51{,}475{,}968$	$124{,}439{,}808$	$354{,}823{,}168$
embedding share of $D$	$51.0\%$	$31.6\%$	$14.8\%$
tokens per block parameter	$40.8$	$29.3$	$23.5$

Appendix

Table 9: Model scales.

D

counts every trainable scalar, embeddings and biases included, and every model is trained on

20D

tokens ( Hoffmann et al., 2022 ) . Tokens per block parameter is the effective data budget on the transformer blocks alone; it swings

1.74\times

across the ladder at a fixed

20

tokens per total parameter.

Optimizer	AdamW ( Loshchilov and Hutter, 2019 )
$\beta_{\text{Adam}}$	$(0.9,\ 0.95)$
Weight decay	$0.1$ , applied uniformly
Gradient clipping	global norm $1.0$
Precision	bf16 mixed
Peak learning rate	$6\times 10^{-4}$
Warmup	linear, first $1\%$ of steps

Appendix

Table 10: zoo training hyperparameters, identical at all three scales.

Scale	$\beta$	total	trunk	branch
Mini	$0.15$	$7{,}854$	$6{,}676$	$1{,}178$
Mini	$0.30$	$7{,}854$	$5{,}498$	$2{,}356$
Small	$0.15$	$18{,}988$	$16{,}140$	$2{,}848$
Medium	$0.15$	$54{,}142$	$46{,}021$	$8{,}121$

Appendix

Table 11: Trunk and branch step split. Counts are exact.

\beta=0.30

exists only as the Mini-scale calibration point of Appendix A.4 .

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

Appendix

Table 12: Pretraining data domains. † Python-only because codeparrot-clean is Python-only. ‡ French only, one fineweb-2 per-language configuration. Evaluation split is SHA-1 of the UTF-8 document bytes, one in

64

Figure 2: (a) The mixture samples shown as entropy, $H(\pi)$ vs the maximum component-wise concentration in $\pi_{i}$ ; the design alone, on $\Delta^{4}$ , with no weights involved. (b) The Mini population by top-2 Gram-PCA components, colored by the $argmax\ \pi$ category. (c) The four disjoint pairwise pools of Table 1 in full $\mathbb{R}^{D}$ weight space. Cosine similarities are read off the displacement Gram equation 2 and are therefore exact at every coordinate, without rank truncation. (d) The same four pools in function space, as next-token disagreement (see Table 3 ).

	$\beta=0.15$	$\beta=0.30$	$\beta=0.60$
between/within distance ratio	$\mathbf{4.08}$	$3.98$	$2.88$
displacement from trunk	$16.5\%$	$28.7\%$	$55.9\%$

Appendix

Table 13: Weight-space geometry against

\beta

, measured at Mini with

N=12

. The between/within distance ratio is what a PCA basis is built from, and the largest

\beta

is the worst of the three on it.

	Mini		Small		Medium
	within	between	within	between	within	between
worst interior barrier	$-0.0067$	$-0.0669$	$-0.0105$	$-0.0642$	$-0.0161$	$-0.0271$
mean dip below chord	$-0.0165$	$-0.1568$	$-0.0239$	$-0.1587$	$-0.0353$	$-0.1251$
midpoint $-$ best endpoint	$-0.0100$	$-0.0919$	$-0.0177$	$-0.0283$	$-0.0305$	$\mathbf{+0.0541}$
midpoint beats both endpoints	$19/20$	$18/20$	$20/20$	$15/20$	$20/20$	$\mathbf{5/20}$

Appendix

Table 14: Linear-interpolation barrier, in nats of mean-over-domain log perplexity, maximized over the interior of the path. Negative values mean the interpolated model is better than the chord, so there is no barrier at any scale. The midpoint

-

best endpoint row is the model-averaging statement and is independent of the barrier.

Figure 3: Excess of the interpolant’s loss over the chord, in nats of mean-over-domain log perplexity, for all $40$ scored pairs at each scale. Members are interpolated in the full $\mathbb{R}^{D}$ weight space, $(1-\alpha)\theta_{i}+\alpha\theta_{j}$ on raw weights. Chord endpoints are excluded from reporting.

Rank $k$ , population statistics	$N-1=99$
Rank $k$ , generative experiments	$50$ (Appendix B.2 )
One-dimensional parameters	retained ( $\sim 0.1\%$ of $D$ )
Code normalization	one global scalar to unit training RMS
Synthetic weight-space noise	none

Appendix

Table 15: Weight-space representation hyperparameters.

	$500$ epochs	$8{,}000$	$20{,}000$
$k=50$	$27.6\%$	$90.2\%$	$\mathbf{101.1\%}$
$k=99$	$1.8\%$	$21.7\%$	—

Appendix

Table 16: Conditional recovery at width

(64,128,64)

, Mini scale, varying only rank and duration;

100\%

is a correct conditional. The only

k=99

configuration that converges used a network seven times wider, which Appendix B.6 shows to be collapse-prone.

	Mini	Small	Medium
flow round-trip	$4.05$	$4.17$	$5.11$
$-$ truncation floor	$1.80$	$1.93$	$2.46$
$=$ generator’s own excess	$\mathbf{2.25}$	$\mathbf{2.24}$	$\mathbf{2.65}$

Appendix

Table 17: Round-trip error decomposed into what the representation permits and what the generator adds,

\%\Delta\mathrm{PPL}

over

85

held-out cells.

Block	$\lambda$ share	within/ens.	of all within	$\eta^{2}$	$R^{2}(\pi)$
c0–c3	$0.803$	$0.001$	$0.4\%$	$0.9998$	$0.578$
c4	$0.038$	$0.004$	$0.1\%$	$0.9934$	$0.368$
c5–c9	$0.036$	$0.969$	$15.0\%$	$0.4390$	$0.182$
c10–c24	$0.036$	$4.178$	$65.2\%$	$0.0283$	$0.013$
c25+	$0.087$	$0.514$	$19.4\%$	$0.0144$	$0.001$

Appendix

Table 18: Code structure at

k=99

, pooled over the anchor groups. The rank must be quoted with this table, because the

\lambda

shares are shares of the retained trace. “within/ens.” compares the twenty anchors’ within-group variance against all

100

members’ total, so it exceeds

1

wherever the anchor groups are more dispersed than the ensemble.

\eta^{2}

is the anchors-only between-group share, bounded in

[0,1]

by construction; five one-hot groups saturate a one-way ANOVA on

\Delta^{4}

, so read it as a statement that the within-mixture noise is small in that block, not as evidence that

\pi

acts linearly — the

R^{2}(\pi)

column measures that.

same eigenvalues, embeddings only, read at lower $k$ :
	Mini (8 blk)		Small (12 blk)		Medium (24 blk)
region, $k=99$	eff. rank	var. share	eff. rank	var. share	eff. rank	var. share
whole stack	$4.466$	$100\%$	$5.637$	$100\%$	$8.725$	$100\%$
blocks only	$7.198$	$13.4\%$	$8.872$	$35.8\%$	$11.346$	$64.7\%$
embeddings only	$\mathbf{4.026}$	$86.6\%$	$\mathbf{4.084}$	$64.2\%$	$\mathbf{4.961}$	$35.3\%$
$k=49$	$3.669$		$3.709$		$4.419$

Appendix

Table 19: Member spectrum by coordinate region. Absolute effective rank

k

is stated because the normalized variant is not

k

-invariant. “var. share” is the region’s share of the total centered variance, which is the composition confound in one number — compare it to the region’s share of

D

in Table 2 .

		Mini		Small		Medium
region	design	interior	vertex	interior	vertex	interior	vertex
whole model	$\pi$	$-0.0874$	$-0.1919$	$\phantom{-}0.0539$	$-0.2537$	$\phantom{-}0.0297$	$-0.2815$
	$\pi,H(\pi)$	$\phantom{-}0.0560$	$\phantom{-}0.2209$	$\phantom{-}0.1576$	$\phantom{-}0.1874$	$\phantom{-}0.0979$	$\phantom{-}0.1002$
	$\pi$ quadratic	$\phantom{-}0.3251$	$\phantom{-}0.1170$	$\phantom{-}0.3420$	$-0.0043$	$\phantom{-}0.2174$	$-0.2518$
blocks only	$\pi$	$\phantom{-}0.0463$	$-0.0959$	$\phantom{-}0.0295$	$-0.1354$	$\phantom{-}0.0355$	$-0.2348$
	$\pi,H(\pi)$	$\phantom{-}0.1079$	$\phantom{-}0.1940$	$\phantom{-}0.0846$	$\phantom{-}0.1563$	$\phantom{-}0.0847$	$\phantom{-}0.1162$

Appendix

Table 20: LOMO predictive

R^{2}

of the branch delta on

\pi

equation 6 ,

85

folds over

100

members: interior scores the

80

singleton folds, vertex the

5

anchor folds; higher is better within

[-1,1]

. Embedding and function-space regions are in Table 21 .

		Mini		Small		Medium
region	design	interior	vertex	interior	vertex	interior	vertex
embeddings	$\pi$	$-0.1217$	$-0.2036$	$\phantom{-}0.0728$	$-0.3084$	$\phantom{-}0.0145$	$-0.3480$
	$\pi,H(\pi)$	$\phantom{-}0.0427$	$\phantom{-}0.2242$	$\phantom{-}0.2142$	$\phantom{-}0.2017$	$\phantom{-}0.1331$	$\phantom{-}0.0772$
	$\pi$ quadratic	$\phantom{-}0.3461$	$\phantom{-}0.1413$	$\phantom{-}0.4619$	$\phantom{-}0.0867$	$\phantom{-}0.4009$	$-0.0721$
function space	$\pi$	$\phantom{-}0.0362$	$-0.4358$	$\phantom{-}0.0073$	$-0.4089$	$\phantom{-}0.0025$	$-0.4004$
	$\pi,H(\pi)$	$\phantom{-}0.0900$	$\phantom{-}0.1125$	$\phantom{-}0.0577$	$\phantom{-}0.0754$	$\phantom{-}0.0434$	$\phantom{-}0.0325$

Appendix

Table 21: The embedding and function-space regions of Table 20 , on the same

85

folds and the same three designs. The embedding region is where the quadratic surface does best — it is the only region reaching

R^{2}>0.40

anywhere — and function space is where

\pi

alone does worst.

	Mini (51M)			Small (124M)			Medium (355M)
stratum	flow	retr.	gap	flow	retr.	gap	flow	retr.	gap
lower tercile ( $n=29$ )	2.60	3.63	+1.03	2.21	3.39	+1.18	3.20	3.72	+0.52
middle tercile ( $n=29$ )	2.38	3.94	+1.56	2.48	3.61	+1.13	3.24	3.66	+0.41
upper tercile ( $n=27$ )	7.40	9.96	+2.55	8.10	10.71	+2.61	9.18	10.99	+1.81
interior ( $n=80$ )	2.66	4.43	+1.77	2.57	4.26	+1.68	3.39	4.37	+0.98
one-hot vertices ( $n=5$ )	26.27	26.83	+0.56	29.79	30.38	+0.59	32.65	32.23	-0.42

Appendix

Table 22: The difficulty axis. Rows one to three are rank terciles of the

L_{1}

distance from a request to the nearest mixture its generator still saw, cut on Mini and reused verbatim at every scale so the same mixtures fall in the same row; the boundaries are

0.240

and

0.291

. Rows four and five split the same

85

mixtures by type instead, which the terciles cannot do. “gap” is retrieval minus flow, so positive favours the generator.

	CRPS, % $\Delta$ PPL			retrieval $-$ flow, % $\Delta$ PPL			marginal $-$ flow, CRPS
	flow	marginal	skill	gap	$t$	won	gap	$t$	won
Mini (51M)	3.10	8.34	+0.628	$+1.70$	$+4.63$	61/85	$+5.23$	$+10.27$	84/85
Small (124M)	3.26	8.38	+0.611	$+1.62$	$+5.38$	61/85	$+5.12$	$+11.02$	83/85
Medium (355M)	3.51	8.83	+0.603	$+0.89$	$+3.13$	43/85	$+5.32$	$+12.44$	82/85

Appendix

Table 23: All

85

held-out mixtures with one LOMO generator per cell, producing

32

samples per cell at Mini and Small and

16

at Medium. CRPS equation 7 scores the whole cohort against the one real model at that mixture (Appendix B.7 ); skill is CRPS against the

\pi

-blind marginal. Absolute round-trip errors are in Table 4 .

Type	rectified flow ( Liu et al., 2023a ; Lipman et al., 2023 )
Width	$(64,\ 128,\ 64)$
Parameters	$\sim\!51$ k
Integration	$50$ explicit Euler steps
Conditioning	$2$ -layer SiLU MLP, concatenated ( Perez et al., 2018 )
Conditioning dropout	$0.15$ ( Ho and Salimans, 2021 )
Guidance scale at evaluation	$1.0$

Appendix

Table 24: Conditional flow hyperparameters. The generator is deliberately small: it is what makes an

85

-cell leave-one-mixture-out sweep affordable, since each cell needs an independently trained flow.

	Mini		Small
$\sigma_{0}$	$a_{\pi}$	pooled	$a_{\pi}$	pooled
$0.125$	$0.316$	$0.950$	$0.591$	$0.992$
$0.25$	$0.521$	$0.966$	$0.485$	$0.952$
$0.5$	$\mathbf{0.877}$	$0.983$	$\mathbf{0.868}$	$0.999$
$0.75$	$0.774$	$0.983$	$0.908$	$0.990$
$1.0$ (committed)	$0.707$	$0.976$	$0.681$	$0.966$

Appendix

Table 25: Flow source width,

\sigma_{0}

, sweep at established hyperparameter values (see Table 24 ).

a_{\pi}

is the sampled spread at one mixture over that mixture’s own within-group spread (

[0,1]

, higher is not always better). Column “pooled” is the corresponding ratio against the whole training set - as such they cannot be cleanly compared to the LOMO generators.

	round-trip error		vs. 200-step sample		dispersion
steps	rel. $L_{2}$	$\times\,n$	rel. $L_{2}$	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

Appendix

Table 26: Small model scale ablation of varied Euler integration steps with 4 leave-one-mixture-out flows (32 samples). Round-trip error integrates a real latent backwards (

1\!\to\!0

) then forwards again (

0\!\to\!1

) and measures how far it lands from its start. Real latents should sit around

1

RMS, and this value is included to detect collapsed flows. This table expands on Table 8 with additional metrics and dispersion checks.

steps	$\Delta$ PPL %	CRPS %	CRPS skill	$a_{\pi}$
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

Appendix

Table 27: Small: all 85 LOMO cells, 8 samples per cell with identical seeds.

a_{\pi}

is the conditional dispersion ratio, where

1

is correct and

<1

is over-tight. We see that generated models survive drastic reductions in Euler integration steps in function space.

	Mini	Small	Medium
$I(z;\pi)$ lower bound vs. matched twin (nats)	$+18.5$	$+31.9$	$+4.9$
identification, top- $1$ (chance $1.2\%$ )	$\mathbf{96.0\%}$	$\mathbf{94.0\%}$	$\mathbf{97.0\%}$
identification, top- $5$	$99.0\%$	$98.0\%$	$99.0\%$
median rank of true $\pi$ (of $85$ )	$1.0$	$1.0$	$1.0$
tie fraction	$0.000$	$0.000$	$0.000$
C2ST, joint $(z,\pi)$	$0.976$	$0.976$	$0.976$

Appendix

Table 28: Code-space density at

k=50

. The mutual-information bound compares against a matched twin — same width, budget and source scale, conditioning ablated — so the difference is not a capacity comparison. Identification ranks the true

\pi

among all

85

candidates. The C2ST is quantised at

1/41

on its held-out split, so the three scales are identical, at

40/41

Figure 4: Table 4 on one axis per scale, with each bar’s multiple of the replicate ceiling beside it. Every bar is a function-space quantity,

\%\Delta\mathrm{PPL}

over the five-domain profile, but the arms reach it differently: the flow samples a code in

z

and decodes it to

\mathbb{R}^{D}

before evaluation, while the

\pi

-surfaces, the global mean and retrieval predict the perplexity profile from

\pi

directly and never materialize weights. The shape of the ladder holds across a

6.9\times

parameter range; what moves is the gap between the flow and retrieval.

Figure 5: One traverse from the barycentre of $\Delta^{4}$ to the multilingual vertex, $41$ requests at two samples each, with samples at a request averaged before plotting. The abscissa traverses $\Delta^{4}$ , the ordinate is function space, and each of the flow’s answers is a distinct code in $z$ decoded to $\mathbb{R}^{D}$ before evaluation. Retrieval returns five distinct real members over the whole traverse; the generator returns $41$ distinct models, and both arrive at very nearly the same place.

	Mini (51M)		Small (124M)		Medium (355M)
domain	real zoo	flow	real zoo	flow	real zoo	flow
web	$-0.633$	$-0.328$	$-0.754$	$-0.352$	$-0.821$	$-0.325$
code	$-0.635$	$-0.401$	$-0.615$	$-0.349$	$-0.618$	$-0.334$
math	$-0.744$	$-0.458$	$-0.804$	$-0.436$	$-0.861$	$-0.414$
books	$-0.361$	$-0.083$	$-0.431$	$-0.063$	$-0.465$	$-0.058$
multilingual	$-1.056$	$-0.720$	$-0.919$	$-0.535$	$-0.895$	$-0.512$

Appendix

Table 29: Ordinary least-squares slope of log perplexity on domain

d

against the requested

\pi_{d}

, with the Pearson correlation in parentheses, for the real zoo and for the generated cohort on the same requests. The final row is the mean over domains of the generated slope as a fraction of the real one.

Figure 6: The dose–response of Table 29 at Mini, measured entirely in function space: a code is sampled in $z$ at each requested $\pi$ and decoded to $\mathbb{R}^{D}$ before any perplexity is taken. Faint points are the generated cohort ( $32$ samples at each of $85$ mixtures) with the real zoo’s per-mixture profiles behind them; the lines carry the slopes of Table 29 . The real slope is fitted per member while the plotted real points are per-mixture means, since an anchor’s four replicates carry leverage their mean does not.

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 $\pi$	0.2744	0.2503	0.2427
replicate floor (same $\pi$ )	0.1657	0.1791	0.1892
rank-50 floor	0.0192	0.0133	0.0120
flow cohort spread	0.1302	0.1343	0.1585

Appendix

Table 30: Generated model next-token disagreement from vertex pools. Vertices and interior are never pooled, organization and comparison follow structurally from

\pi

. Similar results for zoo members are shown in Table 6 . A related Table 22 in the main body shows the same experiment run on the interior of

\Delta^{4}

. We again find that generated models show higher disagreement rates than noise, but lower next token disagreement rates than original models share (see Table 3 ).

scale	$n_{\text{test}}$	gen $\to$ tr	test $\to$ tr	$Z_{U}$
Mini	$9$	$0.352$	$0.435$	$-2.36$
	$4$	$0.511$	$0.213$	$+2.59$
	$9$	$0.357$	$0.613$	$-2.72$
	$9$	$0.331$	$0.292$	$-1.03$
Medium	$9$	$0.671$	$0.385$	$+1.99$
	$4$	$0.934$	$0.167$	$+3.39$

Appendix

Table 31: Three-sample data-copying test on the leading five components, at four blinded neighbourhoods and two scales. The mixture plan is identical across scales, so these are literally the same four mixtures.

Z_{U}<-3

is copying,

Z_{U}>+3

is underfitting, and

|Z_{U}|\leq 3

is generated points sitting where real held-out data sits.

	stratum	real	flow	$\Delta$
Mini (51M)	interior	$0.2521$	$0.2520$	$-0.0001$
	vertices	$0.2558$	$0.2519$	$-0.0039$
Small (124M)	interior	$0.2574$	$0.2580$	$+0.0006$
	vertices	$0.2616$	$0.2590$	$-0.0026$
Medium (355M)	interior	$0.2729$	$0.2715$	$-0.0014$
	vertices	$0.2720$	$0.2680$	$-0.0040$

Appendix

Table 32: HellaSwag accuracy for the runs of Table 5 ,

2{,}000

questions, zero shot. real is the members pretrained at that mixture, averaged within a cell before the five mixtures of a stratum are averaged; flow is the leave-one-mixture-out generator at

8

samples per cell. Accuracy is quantized at

1/2000=0.0005

and chance is

0.25

Weight-Space Learning

See all 27 papers

Data Mixing for Large Language Models Pretraining: A Survey and Outlook
Mar 25, 2026Zhuo Chen, Yuxuan Miao, Supryadi +1Data Mixture Optimization Language Model Pretraining
The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China
Position: Weight Space Should Be a First-Class Generative AI Modality
May 18, 2026Zhangyang Wang, Peihao Wang, Kai WangGenerative Modeling Weight-Space Learning
University of Texas at Austin · Tencent
Can an LLM Actually Reason and Program about Neural Network's Weights or is it just memorization on steroids? Exploring "Weight-Space Programming"
Nov 5, 2026Nazmus Sakib AhmedNeural Network Interpretability Memorization in Language Models
BRAC University, Dhaka, Bangladesh

Mar 25, 2026J/K move · Enter open · S save

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

Zhuo Chen, Yuxuan Miao, Supryadi, Deyi Xiong

The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China

Data Mixture Optimization Language Model Pretraining

Open paper AskPDF arXiv

Large language models (LLMs) rely on pretraining on massive and heterogeneous corpora, where training data composition has a decisive impact on training efficiency and downstream generalization under realistic compute and data budget constraints. Unlike sample-level data selection, data mixing optimizes domain-level sampling weights to allocate limited budgets more effectively. In recent years, a growing body of work has proposed principled data mixing methods for LLM pretraining; however, the literature remains fragmented and lacks a dedicated, systematic survey. This paper provides a comprehensive review of data mixing for LLM pretraining. We first formalize data mixture optimization as a bilevel problem on the probability simplex and clarify the role of data mixing in the pretraining pipeline, and briefly explain how existing methods make this formulation tractable in practice. We then introduce a fine-grained taxonomy that organizes existing methods along two dimensions: static versus dynamic mixing. Static mixing is further categorized into rule-based and learning-based methods, while dynamic mixing is grouped into adaptive and externally guided families. For each class, we summarize representative approaches and analyze their strengths and limitations from a performance-cost trade-off perspective. Building on this analysis, we highlight challenges that cut across methods, including limited transferability across data domains, optimization objectives, models, and validation sets, as well as unstandardized evaluation protocols and benchmarks, and the inherent tension between performance gains and cost control in learning-based methods. Finally, we outline several exploratory directions, including finer-grained domain partitioning and inverse data mixing, as well as pipeline-aware designs, aiming to provide conceptual and methodological insights for future research.

Abstract

Weight-Space Learning

Explore similar work

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

Position: Weight Space Should Be a First-Class Generative AI Modality

Can an LLM Actually Reason and Program about Neural Network's Weights or is it just memorization on steroids? Exploring "Weight-Space Programming"