Code for the paper "Requential Coding: Pushing the Limits of Model Compression with Self-Generated Training Data".
Requential coding compresses a generative model by coding a training process built from self-generated data. A student model samples its own training batches, a teacher training on real data selects them via relative entropy coding, and the code length is the cumulative teacher-student KL, independent of both parameter count and data entropy.
train.py
: the requential training loop (teacher on real data, student on teacher-sampled data, per-step MC KL, optional iso-loss projection and PAC-Bayes bound logging)model.py
: GPT-2-style transformer (flax NNX)bounds.py
: PAC-Bayes bound quantitiesdata.py
: dataset sckpt.py
: checkpoint save/resumeconfigs/default.yaml
: all hyperparameters with the paper defaultsdata_prep/
: dataset preparation scriptssweeps/
: one wandb sweep per paper experimentnotebooks/
: one self-contained notebook per paper figure
conda create -n requential python=3.10 -y
conda activate requential
pip install -U "jax[cuda12]" flax optax hydra-core omegaconf wandb tqdm numpy
pip install pandas matplotlib seaborn jupyter huggingface_hub datasets # data prep + figures
Everything runs on GPUs and, with jax[tpu]
, on TPUs. The batch is sharded
over all visible devices, and both single- and multi-GPU setups work. Use
CUDA_VISIBLE_DEVICES
to pick GPUs. The paper's experiments ran on TPU v6e-8
hosts, and the accumulation factors in the sweeps are calibrated for that
memory. The FineWeb sweeps use A=8
, and only the largest width-2752 model
needs A=16
. If you run out of memory, raise the gradient-accumulation
factor A
(the per-step microbatch is B/A
) or set
model.gradient_checkpointing=true
. Note that the FineWeb shuffles
data in microbatch-sized blocks, so changing A
changes the data order. The run is statistically equivalent but not step-for-step identical.
python data_prep/prepare_fineweb.py 201 # FineWeb, GPT-2 tokens (201 chunks = 20B)
python data_prep/prepare_open.py # OpenWebText, character-level (vocab 96)
python data_prep/prepare_cifar5m.py '/path/to/cifar5m_part{}.npz' # CIFAR-5M pixels
Datasets land under datasets/<name>/
. The synthetic noise
dataset (fresh
uniform-random tokens, unlearnable) and repeat
dataset (one token repeated, trivially learnable) need no preparation.
A single requential run, with the teacher on real data and the student distilled on teacher-sampled data, both with EMA smoothing:
CUDA_VISIBLE_DEVICES=0 python main.py ds_path=open model.width=256 T=100_000_000 \
B=256 opt.warmup_tokens=16_384_000 wandb_project=requential
Sampling from the teacher dominates the runtime, since every step
autoregressively generates a full synthetic batch. To maximize sampling
throughput, use a large batch size B
and train on multiple devices. Generation is parallelized across all visible devices, with each device sampling its own shard of the batch.
Key flags (see configs/default.yaml
for all of them):
| flag | meaning |
|---|---|
T / tpp |
|
| real-token budget D, absolute or in tokens per parameter | |
model.width |
|
| model size axis (depth fixed at 8 in the paper) | |
isoloss_proj |
|
| iso-loss projection: periodically reset the teacher to the student, then retrain it on real data (student d) until its loss recovers | |
teacher_ema / student_ema |
|
| EMA smoothing timescales (teacher smoothing off at 0) | |
ensemble_size |
|
| E students sharing one teacher and one synthetic stream | |
unique_tokens |
|
| cap on unique real tokens (multi-epoch regime) | |
log_bound |
|
log PAC-Bayes bound quantities (req_bound/* , ptq_bound/* ) each eval |
|
train_student |
|
false = plain LM training (prequential/PTQ baseline only) |
|
ckpt_dir |
|
| local directory to save/auto-resume full training state |
Every run logs the per-token teacher-student kl
(nats), the code lengths
L_req
(requential, bits) and L_preq
(prequential, bits), and
ema_{student,teacher}_eval_loss
on held-out data. With log_bound=true
it
also logs ema_{student,teacher}_train_loss
and the bound families
req_bound/{loss,c,sigma,bounded_loss}
(the student, priced with the
requential code) and ptq_bound/{...}
(the teacher, priced as an idealized
lossless 4-bit quantization). The bound's empirical risk and variance
statistic are measured on a fixed train prefix, and the token denominator is
D = unique_tokens
when multi-epoching.
Each experiment is a wandb sweep. Launch it and point one agent per node at it:
wandb sweep -p requential sweeps/scaling_fineweb.yaml
CUDA_VISIBLE_DEVICES=0,1,2,3 wandb agent <entity>/requential/<sweep-id>
| sweep | figure | experiment |
|---|---|---|
scaling_{open,cifar5m,fineweb}.yaml |
||
| Figs 5, 7, 12 | width sweep × iso-loss projection, bounds on | |
ensemble_fineweb.yaml |
||
| Fig 6 | E ∈ {1,2,4,8} ensembles, shared teacher | |
overfit_fineweb.yaml |
||
| Fig 8 | multi-epoch over 1e8 unique tokens | |
synthetic_controls.yaml |
||
| Fig 9 | noise / repeat controls | |
hundredM_{open,cifar5m,fineweb}.yaml |
||
| Figs 4, 9 | ~100M-param matrix |
After the sweeps finish, run the matching notebook in notebooks/
. Each is
self-contained, fetches your runs from wandb by sweep tag (set PROJECT
/
ENTITY
at the top), caches locally, and writes figures/*.pdf
:
fig5_scaling_model.ipynb
: larger models compress to smaller sizesfig6_ensemble.ipynb
: larger ensembles are more compressiblefig7_generalization_bound.ipynb
: PAC-Bayes bounds vs idealized PTQ (also draws the Fig 12 OWT panels)fig8_overfitting.ipynb
: U-shaped bound under data repetitionfig9_learnable_info.ipynb
: learnable information across datasets
@article{qiu2026requential,
title={Requential Coding: Pushing the Limits of Model Compression with Self-Generated Training Data},
author={Qiu, Shikai and Finzi, Marc and Zheng, Yujia and Zhang, Kun and Wilson, Andrew Gordon},
year={2026}
}