Chain-of-Model Learning for Language Models

A new architecture for efficient scaling and elastic inference.

Read paper View code

We propose a novel learning paradigm, termed “Chain-of-Model” (CoM), which incorporates a causal, chain-like structure into the hidden states of each layer. This approach introduces powerful new efficiencies in model training and unprecedented flexibility for deployment.

The foundation of CoM is “Chain-of-Representation” (CoR), which formulates the hidden state at each layer as a combination of multiple sub-representations, or “chains.” Critically, each chain in the output can only see its preceding input chains, enforcing a causal dependency. This structure allows a single model to house multiple sub-models of varying sizes.

As a result, models built on the CoM framework can be progressively scaled by adding new chains to existing ones, and they can offer elastic inference by activating a variable number of chains. Our experimental results demonstrate that our CoLM family of models achieves performance comparable to standard Transformers while offering these significant advantages in extensibility and adaptability.

How CoLM Works: An Interactive Look

CoLM applies the chain principle to the core components of the Transformer architecture.

Chain-of-Representation (CoR) is the fundamental concept of CoLM. It formulates a standard representation as a concatenation of multiple "chains." Each chain is a sub-representation, and by activating a different number of chains, the model can operate at different "scales" of complexity and knowledge.

For example, activating only the first chain (Scale 1) uses the smallest version of the model. Activating the first two chains (Scale 2) leverages a larger, more capable model that builds upon the first. This structure is what enables a single trained CoLM model to provide multiple sub-models for elastic inference, a key advantage over standard architectures.

Chain-of-Linear

Causal, grouped connections.

Standard Linear

Dense, all-to-all connections.

Chain-of-Attention

Attention heads are partitioned into chains, maintaining causal structure.

Comparison of standard Attention and Chain-of-Attention.
CoLM-Air (KV Sharing)

Key/Value pairs are computed in the first chain and shared across all subsequent chains.

Diagram showing KV sharing in CoLM-Air.

Competitive Performance, Unmatched Flexibility

CoLM delivers results on par with standard Transformers of similar size, proving that architectural flexibility and efficiency can be achieved without sacrificing performance.

Benchmark
Baseline (1.10B)
CoLM (1.11B)
CoLM-Air (1.11B)
HellaSwag
40.01
40.25
39.85
Obqa
31.19
31.39
31.19
WinoGranda
52.72
52.41
52.09
ARC-e
43.52
43.73
44.30
ARC-c
23.63
23.81
23.63
Boolq
57.43
58.01
56.72
Piqa
67.30
67.30
66.76
Average Score
45.11
45.27
44.80

The Breakthrough: Blazing-Fast Prefilling

By computing keys and values only in the first chain, our CoLM-Air variant radically accelerates the prefilling stage—a major bottleneck in long-context inference.

Charts showing CoLM-Air achieving significantly faster prefilling speeds.

Achieve up to 27x faster prefilling speeds with MInference.

The Complete CoLM Architecture

This diagram illustrates the full data flow within a CoLM layer. It shows how embeddings are processed through chain-aware normalization, attention with KV sharing, and a chained feed-forward network to produce progressively refined outputs for each scale.

Full architecture diagram of the CoLM model.