Latent-Kernel Flows (LKF)

Reference implementation of the Latent-Kernel Flow, a discrete flow map whose one-step kernel is a mixture of $M$ factorized components sharing a single sequence-level latent. The code trains and evaluates the text models on LM1B and WikiText-103, runs the synthetic $M$-sweep that supports the theory section, and reproduces the decoding ablations in the appendix.

Layout

data/          dataloaders for the pre-tokenized text caches and the synthetic corpora
models/        the LKF backbone (shared trunk + latent branch + router) and the AR anchor
sampling/      the reverse samplers, including the analytic sampler used for every reported number
training/      the training loops for LKF and for the AR anchor
eval/          the evaluation harnesses that write the normalized eval.json files
metrics/       generative perplexity, sample entropy, captured information and support checks
configs/       one YAML per reported run
scripts/       data preparation, the multi-GPU synthetic sweep, and the decoding ablations

Installation

The code was developed against Python 3.10 and CUDA 12.1 on two NVIDIA H100 80GB accelerators.

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

There is no custom CUDA dependency. Attention runs through PyTorch's own scaled_dot_product_attention, so any recent PyTorch build works.

Data

Everything reads from one root directory, which is taken from the LKF_DATA_ROOT environment variable and defaults to ./data_root.

export LKF_DATA_ROOT=/path/to/data

The text corpora are HuggingFace save_to_disk Arrow caches in the layout MDLM's dataloader writes, so LKF trains on the identical tokenization as the MDLM, SEDD and ReDi baselines and the comparison in the paper is a same-tokenizer one.

python scripts/prep_text_caches.py --datasets lm1b wikitext103

That writes $LKF_DATA_ROOT/text/lm1b/lm1b_{train,test}_bs128_unwrapped.dat and $LKF_DATA_ROOT/text/wikitext103/wikitext103_{train,validation}_bs128_wrapped.dat. LM1B is unwrapped, meaning one padded sentence per row, so roughly three quarters of its tokens are [PAD] and the trainer excludes those positions from the loss. WikiText-103 is wrapped, meaning documents are concatenated and cut into blocks. We verified that the script reproduces the caches the reported runs consumed token for token.

The synthetic corpora are generated locally and each one stores its analytic total correlation in meta.json, which is what the exact metrics in the theory section are checked against.

python scripts/gen_synthetic.py --only parity_d16 hidden_agree

Training

Each config in configs/text/ is the exact configuration behind the corresponding row of the main table, at 200k steps and batch size 512. Training writes checkpoints under the config's out_dir.

python -m training.trainer --config configs/text/lm1b_M8.yaml

Top-level fields can be overridden from the command line, for example --override steps=1000 batch_size=64, which is handy for a quick smoke test.

The convenience driver trains and then scores a config in one call, which is the fastest way to reproduce a single row.

bash scripts/train_text.sh configs/text/lm1b_M8.yaml
EVAL_ONLY=1 bash scripts/train_text.sh configs/text/lm1b_M8.yaml   # score an existing checkpoint

Evaluation

The text harness reports the ELBO perplexity, the generative-perplexity-vs-NFE curve, the sample entropy and the captured information in a single normalized eval.json. Every number in the paper uses the analytic reverse sampler, which is the MDLM-equivalent update, and the commit-$M$ decode in which the latent is drawn once and frozen for the whole trajectory.

python -m eval.eval_text --ckpt runs/lm1b_M8/ckpt_0200000.pt --out runs/lm1b_M8/eval.json \
    --sampler analytic --nfe 1,2,4,8,16,32 --gen_samples 512 --gen_ppl_model gpt2-large

The synthetic harness computes the exact conditional mutual information $I(k; X_s \mid X_t)$ and the effective total correlation by enumeration, both of which are bounded by $\log M$ and $L \log M$ respectively.

python -m eval.eval_synthetic --ckpt runs/M_sweep/parity_d16/M4_seed0/ckpt_0008000.pt \
    --dataset parity_d16 --out runs/M_sweep/parity_d16/M4_seed0/eval.json --do_ablations