Chain-of-Model Learning for Language Model
CoLM enables progressive scaling and elastic inference; CoLM-Air reaches about 3× faster 1M-token prefilling.
Key Findings
Methodology
The paper introduces Chain-of-Representation (CoR), partitioning hidden states into chains, where output chain i accesses only input chains 1 through i. Chain-of-Layer (CoL) applies this causal constraint to Linear, Attention, FFN, and Normalization, yielding Chain-of-Model (CoM). Chain-of-Language-Model (CoLM) embeds CoM into Transformers. CoLM-Air computes Keys and Values only in the first chain and shares them across chains.
Key Results
- On SlimPajama with roughly 200B training tokens, 1.11B-parameter CoLM with C={16,16} and width 2560 obtains 45.27 average accuracy, slightly exceeding the 1.10B standard baseline at 45.11.
- After only about 8B tokens, Chain Expansion raises Tiny-LLaMA from 57.43 to 58.35 average accuracy and LLaMA-3.2-1B from 60.39 to 60.53, while freezing the first chain.
- For 1M-token prefilling, CoLM-Air is approximately 1.6× faster with C={16,16} and 3.0× faster with C={8,8,8,8}; combined with MInference, the reported maximum speedup reaches 27×.
Significance
The work turns model scale from a fixed architectural choice into an incrementally expandable and selectively activatable property. This addresses two persistent problems: scaling generally requires retraining from scratch, while deployment usually activates one fixed parameter budget. A single CoLM can expose small, medium, and large operating points for different latency, memory, and quality requirements. Scientifically, CoR, CoL, and CoM provide a useful abstraction for studying causal multi-scale representations. Industrially, the approach is particularly relevant to long-context prefilling, resource-constrained serving, and continual model growth.
Technical Contribution
The central contribution is causal sparsity along the hidden dimension: parameters for output chain i consume only preceding chains. Chain-of-Linear computes yi=Wi x≤i+bi, with Wi connected to cumulative input chains. Attention assigns dedicated heads to chains, requiring chain ratios to sum to the total head count. Multi-chain cross-entropy trains separate prediction heads for each scale. CoLM-Air restricts KV computation to the first chain and repeats KV following GQA-style sharing, enabling cache reuse, seamless scale switching, and cheaper prefilling.
Novelty
Unlike dense Transformers, which activate a fixed parameter set, and MoE systems, which rely on routing, CoLM explicitly encodes scale causality inside every representation. The authors position CoLM-Air as the first design enabling different-scale language models to share existing KV states without recomputation during switching. The deeper innovation is not ordinary pruning: the prefix of a large model is itself a runnable smaller model.
Limitations
- CoLM-Air computes KV only in the first chain, improving flexibility and prefilling but reducing representational capacity; its 1.11B configuration averages 44.80 versus 45.27 for full CoLM.
- Experiments use about 1B-parameter models, 50K steps, and approximately 200B tokens on A100 GPUs. Evidence for much larger models, longer training, and real online decoding economics remains limited.
- Multi-chain cross-entropy adds training overhead, so pretraining uses ordinary cross-entropy and applies the multi-scale objective mainly during later fine-tuning.
Future Work
Future research should test chain expansion at substantially larger scales and over longer training runs, and explore nonuniform chain widths, dynamic chain activation, and optimized block-sparse kernels. Systematic comparisons with GQA, MQA, LoRA, quantization, and Chain Tuning are needed. Real serving studies should jointly measure quality, latency, memory, communication, and cost under heterogeneous workloads.
AI Executive Summary
Large language models have largely advanced by scaling Transformer parameters. Yet this strategy has two structural weaknesses: expanding a model commonly requires retraining from scratch, while deployment typically activates one fixed model size. Chain-of-Model (CoM) proposes to place several usable scales inside one model and to grow them incrementally.
Its foundation, Chain-of-Representation (CoR), partitions hidden states into chains. Output chain i may read only chains 1 through i, creating Chain-of-Layer and then Chain-of-Language-Model (CoLM) when applied to Transformer Linear layers, attention, FFNs, and normalization. CoLM-Air goes further: it computes Keys and Values only in the first chain and shares them across later chains, allowing cache reuse when changing model scale.
Experiments use SlimPajama, a 600B-token corpus, LLaMA-2 tokenization, 32 A100 GPUs, and 50K pretraining steps. A 1.11B CoLM reaches 45.27 average accuracy versus 45.11 for a 1.10B baseline; Chain Expansion improves Tiny-LLaMA from 57.43 to 58.35. At a 1M-token context, CoLM-Air provides roughly 1.6×–3.0× faster prefilling, and up to 27× with MInference. KV sharing slightly lowers accuracy, and validation remains small-scale. Still, the paper presents a compelling shift: model size becomes a composable and switchable property rather than a single immutable setting.
Deep Analysis
Background
LLM development has centered on dense Transformers, alongside MoE, elastic inference, and continual training. Dense models are predictable but activate a fixed parameter budget; MoE reduces computation but introduces routing and load-management complexity. Existing continual-scaling methods often fail to preserve a usable old scale. CoM reframes scale at the representation level, allowing multiple submodels to coexist within one parameterization.
Core Problem
The goal is to train multiple model sizes jointly and expand an existing model without destroying prior capability. This is difficult because matrix multiplication, attention dot products, and normalization naturally mix information from different scales. KV caches create a second bottleneck: caches are usually tied to a specific model configuration, so switching from a small to a large model requires recomputing the entire prefix.
Innovation
First, CoR partitions a vector into ordered chains and defines scale by prefix length. Second, CoL imposes the causal rule that output chain i depends only on input chains ≤i. Third, CoLM applies this rule throughout the Transformer. Fourth, CoLM-Air computes KV only in the first chain, enabling cache reuse and scale switching. Fifth, Chain Expansion treats a trained model as the first chain and appends capacity rather than restarting training.
Methodology
- �� Linear: partition inputs and outputs; compute yi=Wi x≤i+bi, with Wi connected only to preceding chains.
- �� Attention: allocate dedicated attention heads to each chain; chain ratios sum to the total head count, preventing cross-scale mixing within a head.
- �� FFN/Normalization: replace all Linear layers with Chain-of-Linear and normalize each chain separately.
- �� Embedding: keep training embeddings unchanged; inference at scale i uses the first i chain dimensions.
- �� Objective: attach a classification head to each scale and use multi-chain cross-entropy; ordinary cross-entropy is used during pretraining for efficiency.
- �� KV sharing: CoLM-Air computes first-chain KV and repeats it across query heads in a GQA-like manner.
Experiments
Pretraining uses the 600B-token SlimPajama corpus, LLaMA-2 tokenizer with vocabulary 32,000, sequence length 4096, BF16, AdamW at 1.5×10^-4, 32 NVIDIA A100 40GB GPUs, and effective batch size 1024. Training runs for 50K steps, approximately 200B tokens. The baseline is a single-chain LLaMA-3.2-1B-style model. Zero-shot evaluation uses EleutherAI Language Model Evaluation Harness and acc_norm on HellaSwag, OBQA, WinoGrande, ARC-e/c, BoolQ, and PIQA. Additional studies test expansion, elastic inference, and prefilling.
Results
The standard baseline averages 45.11. A 1.11B CoLM with C={16,16} averages 45.27; 0.86B and 1.18B variants average 44.00 and 44.51. Corresponding CoLM-Air results are 44.80 and 43.41. Chain Expansion improves Tiny-LLaMA from 57.43 to 58.35 and LLaMA-3.2-1B from 60.39 to 60.53. In elastic inference, the first 0.33B chain averages 41.13, while the 0.86B two-chain model reaches 43.90.
Applications
Cloud services can select chain count according to latency, cost, and quality targets; edge devices can run the first chain while servers activate more capacity. CoLM-Air can prefill long contexts with the first chain and switch to a larger model for generation. Continual training and Chain Tuning can freeze old chains and update later chains, reducing adaptation cost and catastrophic forgetting.
Limitations & Outlook
KV sharing sacrifices some representation capacity, and CoLM-Air underperforms full CoLM on most reported averages. The evidence is limited by model size, training duration, and task coverage, so trillion-scale behavior is unknown. Multi-chain execution also requires specialized block-sparse kernels and communication optimization; multi-scale losses add complexity. Future work should evaluate larger models, dynamic chain routing, nonuniform widths, and end-to-end serving metrics.
Plain Language Accessible to non-experts
Imagine a factory that can be expanded one production line at a time. A conventional large language model is like a giant factory built all at once: every order activates the whole building, even when a small job would do, and increasing capacity often means designing and training a new factory from scratch. CoLM divides the factory into ordered lines. The first line can complete basic work; the second reads its output and adds more capability; later lines build on everything before them. Running the first line gives a small model, while opening more lines gives a larger one without discarding the original ability.
CoLM-Air adds a shared inspection record. The first line checks the incoming materials and later lines reuse that record instead of inspecting the entire shipment again. This is why very long prompts can be prepared faster; at one million tokens, some settings are about three times faster.
The trade-off is accuracy: sharing one record can be less precise than letting every line create its own. The overall idea is to make model size behave like modular factory capacity—available in several levels and expandable over time—rather than a single permanent building.
ELI14 Explained like you're 14
Think of a video game base that helps your character solve missions. A normal big language model is like a giant base built at the start: it can do impressive things, but every mission may require the whole base to run, which is expensive on a phone. If you want more power, you often have to rebuild and retrain everything.
CoLM splits the base into connected levels. Level one handles basic jobs. Level two sees level one’s results and adds smarter tools. Level three uses the first two and adds even more power. Turn on only level one and you get a small model; turn on more levels and you get a stronger model. Importantly, adding a level does not erase what the old levels learned.
CoLM-Air also saves a shared record of a long story. Level one reads the story and writes the record; later levels reuse it instead of reading the whole story again. For a one-million-token context, some versions prepare the text about three times faster, and the combination with MInference reaches up to 27× in the reported setting.
But shortcuts have costs! A shared record may miss details that separate records would capture, so accuracy drops slightly. The tests are also around one-billion parameters, so researchers still need to see whether this works equally well for giant models.
Glossary
Chain-of-Representation (CoR)
A hidden vector is partitioned into ordered subrepresentations called chains. A prefix of chains defines a usable model scale.
CoR is the representation-level basis of CoLM’s multi-scale behavior.
Chain-of-Layer (CoL)
A layer satisfies CoL when output chain i depends only on input chains 1 through i. This supports causal computation and prefix activation.
Linear, Attention, FFN, and Normalization are redesigned as CoL components.
Chain-of-Model (CoM)
A model whose layers all satisfy the CoL property. It contains multiple submodels at different scales and can be expanded by adding chains.
CoM is the paper’s general learning paradigm.
CoLM-Air
A CoLM variant with KV sharing: Keys and Values are computed only in the first chain and shared by later chains.
It enables cache reuse, scale switching, and faster prefilling.
Chain Expansion
A trained model becomes the first chain, after which additional chains are added to increase capacity. Earlier chains may be frozen to preserve knowledge.
The paper applies it to Tiny-LLaMA and LLaMA-3.2-1B.
Elastic Inference
Inference dynamically selects computational scale according to latency, memory, or quality requirements. CoLM provides scales through different chain counts.
The paper evaluates this property with zero-shot commonsense benchmarks.
Open Questions Unanswered questions from this research
- 1 It remains unknown whether CoLM scales reliably beyond 70B parameters; larger systems may amplify communication, load imbalance, and optimization problems.
- 2 The paper does not establish whether selective KV, hierarchical caching, or distillation can remove the accuracy cost of KV sharing.
- 3 Code, reasoning, multilingual, and real production workloads are underexplored; end-to-end latency and cost measurements are still needed.
Applications
Immediate Applications
Elastic cloud serving
Providers can run the first chain for latency-sensitive requests and activate more chains for difficult tasks. This requires multi-scale training and optimized inference kernels, but can balance quality, memory, and cost within one deployment.
Long-context prefilling
CoLM-Air can process historical context with the first chain and switch to a larger scale during generation. This is useful for document QA, agents, and long conversations where prompt prefilling dominates runtime.
Long-term Vision
Continually expandable foundation models
Future models could grow like modular systems: freeze existing chains to retain general knowledge and train only new chains for domains or tasks. This may reduce retraining cost and mitigate catastrophic forgetting.
Abstract
In this paper, we propose a novel learning paradigm, termed Chain-of-Model (CoM), which incorporates the causal relationship into the hidden states of each layer as a chain style, thereby introducing great scaling efficiency in model training and inference flexibility in deployment. We introduce the concept of Chain-of-Representation (CoR), which formulates the hidden states at each layer as a combination of multiple sub-representations (i.e., chains) at the hidden dimension level. In each layer, each chain from the output representations can only view all of its preceding chains in the input representations. Consequently, the model built upon CoM framework can progressively scale up the model size by increasing the chains based on the previous models (i.e., chains), and offer multiple sub-models at varying sizes for elastic inference by using different chain numbers. Based on this principle, we devise Chain-of-Language-Model (CoLM), which incorporates the idea of CoM into each layer of Transformer architecture. Based on CoLM, we further introduce CoLM-Air by introducing a KV sharing mechanism, that computes all keys and values within the first chain and then shares across all chains. This design demonstrates additional extensibility, such as enabling seamless LM switching, prefilling acceleration and so on. Experimental results demonstrate our CoLM family can achieve comparable performance to the standard Transformer, while simultaneously enabling greater flexiblity, such as progressive scaling to improve training efficiency and offer multiple varying model sizes for elastic inference, paving a a new way toward building language models. Our code will be released in the future at: https://github.com/microsoft/CoLM.