GitHub - shikaiqiu/requential-coding

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.

Repository Structure

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 quantities
data.py: dataset loaders
ckpt.py: checkpoint save/resume
configs/default.yaml: all hyperparameters with the paper defaults
data_prep/: dataset preparation scripts
sweeps/: one wandb sweep per paper experiment
notebooks/: one self-contained notebook per paper figure

Environment Setup

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 loader shuffles data in microbatch-sized blocks, so changing A changes the data order. The run is statistically equivalent but not step-for-step identical.

Dataset Preparation

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.

Training

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 paused) 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.

Reproducing the Paper Experiments

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

Reproducing the Figures

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 sizes
fig6_ensemble.ipynb: larger ensembles are more compressible
fig7_generalization_bound.ipynb: PAC-Bayes bounds vs idealized PTQ (also draws the Fig 12 OWT panels)
fig8_overfitting.ipynb: U-shaped bound under data repetition
fig9_learnable_info.ipynb: learnable information across datasets

Citation

@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}
}