Research map / Large language models / Language, safety and interpretability
Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives token-level, non-IID generalization bounds for LLMs via Azuma's inequality on martingale differences, and combines them with Monarch/Kronecker/LoRA parametrizations and QuIP post-training quantization to obtain non-vacuous BPD bounds for models up to LLaMA2-70B.
Problem
Compression-based generalization bounds for LLMs are computed at the document level with IID assumptions, so they are vacuous for billion-parameter models and only non-vacuous for heavily compressed models that generate low-quality text.
Approach
Treat each token (rather than each document) as a data point and use Azuma's inequality on the martingale difference sequence of token NLLs, so the complexity term is divided by the far larger token count and tolerates much less restrictive compression, including post-training quantization of pretrained models.
Key results
- Quantized GPT2 small (124M) achieves a token-level BPD bound of 7.61 on OpenWebText, versus 15.62 random guess, with text quality visibly better than SubLoRA's document-level best model.
- LLaMA2-7B 2-bit QuIP# quantized achieves a non-vacuous BPD bound of 4.28 on Amber (random guess 14.97), and LLaMA2-70B achieves 6.39 in the main table (6.14 in the appendix at 2 bits).
- Monarch-only parametrization gives the tightest token-level BPD bound of 7.65 among nonlinear compressions, versus 10.49 for the SubLoRA baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use token-level prediction smoothing with a learned token-dependent alpha instead of a single global alpha | Token-level generalization bounds · Token-level bound evaluation | BPD bound | Helped |
| Use Monarch matrices without subspace compression rather than SubLoRA | GPT2 Small compression · OpenWebText | BPD bound: 10.49 -> 7.65 | Helped |
| Remove subspace compression from Kronecker factorization (ablation) | GPT2 Small compression · OpenWebText | BPD bound: 10.02 -> 8.03 | Helped |
| Extend LoRA compression to all fully connected layers and train biases and layer normalization parameters | GPT2 Small LoRA compression · OpenWebText | BPD bound: 10.49 -> 7.85 | Helped |
| Increase GPT2 model size from 124M to 774M with post-training quantization | GPT2 · OpenWebText | BPD bound: 7.61 -> 10.47 | Hurt |
| Increase LLaMA2 model size from 7B to 70B with 2-bit post-training QuIP quantization | LLaMA2 · Amber | BPD bound: 4.28 -> 6.39 | Hurt |
| Fine-tune LLaMA2 for dialogue | LLaMA2 · Amber, 2-bit quantization | Bits per Dimension: 4.28 -> 4.54; 4.52 -> 4.77; 6.14 -> 6.40 | Hurt |
| Increase GPT2 context length from 1 to 2 tokens (ablation) | GPT2-S-Quantized · OpenWebText | BPD bound: 11.1 -> 9.0 | Helped |
| Use GPT2 rather than a Markov chain with context length 4 | GPT2-S-Quantized · OpenWebText | BPD bound: 22.4 -> 7.9 | Helped |
| Apply more aggressive post-training quantization (ablation) | GPT2 on structured versus random integer sequences · Structured and unstructured sequences | Integer prediction training accuracy | Mixed |
| Use token-level bounds instead of document-level bounds | Compression-based LLM generalization bounds · LLM next-token prediction | Bound tightness | Helped |
Surprising findings
- Subspace compression, which was essential for document-level bounds, becomes unnecessary or even harmful at the token level.
- LLaMA2-Chat models obtain looser generalization bounds than their base counterparts, suggesting dialogue fine-tuning hurts next-token generalization.
Limitations
- Bounds remain well above typical validation BPD, leaving a gap between the bound and empirical performance.
- Bounds conditioned on training contexts and resampled next tokens; they do not directly bound expected risk on new contexts, and rely on the assumption that next-token distributions are diffuse.
- Token-level bounds for LLaMA models use the Amber dataset as a proxy for unknown pretraining data.
Experimental setup
Datasets: OpenWebText Amber Observed Antibody Sequences (OAS) WebText-derived GPT2 zero-shot benchmarks
Benchmarks: OpenWebText BPD bound Amber BPD bound Top-1/10/100 error bounds OAS antibody bounds
Metrics: bits per dimension bound top-k error bound validation loss correlation with downstream accuracy and perplexity
Baselines: SubLoRA Random Guess GPT2 SubLoRA document-level bound (Lotfi et al.) Markov chains
Builds on: PAC-Bayes / finite hypothesis compression bounds SubLoRA QuIP QuIP# LoRA Azuma's inequality
Closest papers by meaning
- Non-Vacuous Generalization Bounds for Large Language Models (ICML 2024 · similarity 0.95)
- Compute-Optimal LLMs Provably Generalize Better with Scale (ICLR 2025 · similarity 0.92)
- On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.89)
- Language Modeling Is Compression (ICLR 2024 · similarity 0.89)
- Compressing LLMs: The Truth is Rarely Pure and Never Simple (ICLR 2024 · similarity 0.89)
- Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.