Research map / Architectures and efficiency / Attention and transformers
Chain-of-Model Learning for Language Model
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.
Proposes Chain-of-Model (CoM) learning, where hidden states are split into causally ordered chains so models can scale by adding chains and yield nested sub-models, instantiated as CoLM and CoLM-Air (KV sharing) for language modeling.
Problem
Existing LLM scaling requires training from scratch for each size and activates fixed parameter counts, so there is no incremental scaling or elastic inference within one model.
Approach
Represent each layer's hidden state as a concatenation of chains with causal (lower-triangular) connectivity so chain i only sees chains up to i; this lets a prefix of chains act as a smaller sub-model, enables chain expansion from a trained model, and with KV sharing lets all sub-models reuse first-chain keys and values.
Key results
- CoLM {16,16} at 1.11B reaches 45.27 average vs 45.11 for the 1.10B baseline on commonsense tasks (50K steps).
- Chain expansion of LLaMA-3.2-1B with an extra 0.8B chain improves average from 60.39 to 60.53, and TinyLLaMA-v1.1 from 57.43 to 58.35.
- CoLM-Air prefilling at 1M tokens is about 1.6x and 3.0x faster for {16,16} and {8,8,8,8} on LLaMA-1B settings; with MInference up to 27x speedup (1920 to 70.0 s).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CoLM with chains {16, 16} and dim 2560 instead of the standard Transformer | Language model · Commonsense tasks (zero-shot) | Avg. acc_norm: 45.11 -> 45.27 | Helped |
| Add KV sharing to CoLM with chains {16, 16} and dim 2560 (ablation) | CoLM · Commonsense tasks (zero-shot) | Avg. acc_norm: 45.27 -> 44.80 | Hurt |
| Use {16, 16} chains rather than {8, 8, 8, 8} at matched wider dimensions (ablation) | CoLM · Commonsense tasks (zero-shot) | Avg. acc_norm: 44.51 -> 45.27 | Helped |
| Expand the pretrained Tiny-LLaMA-v1.1 with an additional chain | Tiny-LLaMA-v1.1 · Commonsense tasks | Avg.: 57.43 -> 58.35 | Helped |
| Expand the pretrained LLaMA-3.2-1B with an additional chain | LLaMA-3.2-1B · Commonsense tasks | Avg.: 60.39 -> 60.53 | Helped |
| Fine-tune the expanded chain while freezing the first chain | Expanded model · GLUE dev | Avg. accuracy: 48.46 -> 67.79 | Helped |
| Use CoM-MLP with chains {4, 12} instead of a dense MLP | MLP · CIFAR-10 test | Accuracy: 54.93 ± 0.23 -> 55.11 ± 0.21 | Helped |
| Use CoM-MLP with chains {4, 4, 4, 4} instead of a dense MLP (ablation) | MLP · CIFAR-10 test | Accuracy: 54.93 ± 0.23 -> 54.60 ± 0.16 | Hurt |
| Use CoM-MLP with chains {8, 8} instead of a dense MLP (ablation) | MLP · CIFAR-10 test | Accuracy: 54.93 ± 0.23 -> 54.68 ± 0.04 | Hurt |
| Use CoM-MLP with chains {4, 4, 8} instead of a dense MLP (ablation) | MLP · CIFAR-10 test | Accuracy: 54.93 ± 0.23 -> 54.55 ± 0.36 | Hurt |
| Use only the first chain rather than both chains (ablation) | CoLM-Air, {16, 16}, dim 2048 · Commonsense tasks (zero-shot) | Avg.: 43.90 -> 41.13 | Hurt |
| Use CoLM-Air chains {16, 16} rather than the baseline for 1M-token prefilling | LLaMA-3.2-1B · 1M context | Prefilling time (seconds): 522.0 (45.8) -> 265.0 (20.1) | Helped |
Surprising findings
- Applying CoM along model depth (selecting intermediate layers as sub-models with KV sharing) severely degraded performance, so the method currently only works along width.
- CoLM-Air's KV sharing slightly reduces accuracy relative to CoLM across configurations, in exchange for prefilling speed.
Limitations
- Chain-of-Linear layers are poorly suited to tensor parallelism, requiring extra all-reduce and data access overhead.
- Performance of the first chain is critical, and the optimal first-chain configuration is unexplored; the paper also notes results are only at ~1B scale with limited compute.
- Extension to model depth degraded performance, so the method currently applies only to width.
Experimental setup
Datasets: SlimPajama GLUE
Benchmarks: HellaSwag OpenBookQA WinoGrande ARC-e ARC-c BoolQ PIQA SciQ GLUE CIFAR-10
Metrics: acc_norm accuracy average accuracy prefilling latency speedup GLUE accuracy
Baselines: LLaMA-3.2-1B TinyLLaMA-v1.1 Standard Transformer (LLaMA-style 1.1B)
Builds on: Transformer Progressive neural networks Matryoshka representation learning GQA FlashAttention-2
Closest papers by meaning
- From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency (ICLR 2025 · similarity 0.87)
- Eliciting Chain-of-Thought in Base LLMs via Gradient-Based Representation Optimization (AAAI 2026 · similarity 0.87)
- Markovian Transformers for Informative Language Modeling (ICLR 2026 · similarity 0.86)
- Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.86)
- C3oT: Generating Shorter Chain-of-Thought Without Compromising Effectiveness (AAAI 2025 · similarity 0.86)
- Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA (ICLR 2025 · similarity 0.86)
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.