Chain-of-Model Learning for Language Model
Authors: Xiaohua Wang
Kaitao Song
Xu Tan
Huiqiang Jiang
Chengruidong Zhang
Yongliang Shen
Cen LU
Zihao Li
Zifan Song
Caihua Shan
Yansen Wang
Kan Ren
Xiaoqing Zheng
Tao Qin
Yuqing Yang
Dongsheng Li
Lili Qiu
Venue: NeurIPS 2025 poster ()
Keywords: Foundation Model
Next-Gen Language Models
Chain-of-Model
Dynamic Activation
Efficiency
Primary area: deep_learning

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). 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.

Chain-of-Model Learning for Language Model
Xiaohua Wang 1,2,+
Kaitao Song2**
Xu Tan"
Huiqiang, Jiang?
Chengruidong Zhang?
Yongliang Shen
Cen 2.56
Zihao Li
Zifan Song".. 2.8
Caihua Shan
Yansen Wang?
Kan Ren2.4
Xiaoqing Zheng'
Tao Qin
Yuqing Yang
Dongsheng Li²
Lili Qiu
Fudan University! Idiap Research Institute"
Microsoft Research2
Zhejiang University
ShanghaiTech University Tongji University
EPFL"
UlUC
Abstract
this paper, we propose a novel learning paradigm, termed "Chain-of oof---oe--- (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) 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 clastic inference by using different chain numbers. Based on (CoLM), which incorporates the idea of CoM into each layer of Transformer architecture. Based on CoLM, we further introduce CoLM-Air by introducing 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 CoL.M family can achieve comparable performance the standard Transformer, while simultaneously enabling greater flexiblity, such progressive scaling improve training efficiency and offer multiple varying model sizes for clastic inference, paving new way toward building language models. Our code will be released the future at: petps://githh com/ woom//icrrooof//o......
this principle, we devise Chain-of-Lanye haart/-ff-anuugge--aaa
Introduction
With the advent of large language models (LLMs) [1-6], scaling up Transformer architecture [7-9] has been considered the de facto approach to revolutionize existing AI landscape and achieve state-of-the-art performance across a wide range of different tasks. Therefore, there has a growing trend in both industry and academia to explore strategies for scaling up Transformer models. Under these circumstances, the parameter size of LLMs has been an exponential growth, scaling from the billion to the trillion level. As result, its explosive parameters also impose extremely expensive burden on training and cannot offer varying inference usage for different deployment environments. In light of this increasing trend of scaling laws, how to develop and harness LLMs effectively address user instructions in various scenarios has been an open and critical challenge among the whole community.
The first three authors have equal contributions. Corresponding author: Kaitao Song Stillkeeprr@@oallook.com
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Motivated by elastic inference [10, 11] and continual training [12, 13], we claim that existing scaling strategies for LLM architectures still remain some inherent issues:
Unlike human intelligence, which acquires new knowledge incrementally, existing scaling policies cannot retain existing scales and always requires training from scratch, leading inefficiencies.
Existing LLM architectures (e.g., Dense or MoE [14]) always activate fixed scale of parameters, lacking mechanisms to dynamically adapt problem-solving capabilities.
In this paper, we propose new concept, termed "Chain-af-Reprreenttaiii
Scale 1
(CoR), that generalizes the scope of representation paradigm a broader range. Specifically, we observe that any representa- tion can always be viewed as a combination of multi- ple sub-representations at the hidden dimension level. Therefore, we define this composition as Chain-of- Representation, that each sub-representation corre- sponds to one chain. Based on this definition, by using different number of preceding chains, its cor- responding features can be used to encode different knowledge (we call it "scale"), just as shown in Fig- ure 1. Therefore, how to build the connections among CoR features to ensure feature transformation across each scale has been critical step.
(a) Representation
Stale
Scale 2
Scale
Chiin
Chain
Chkka
[b] Chain-of- Represeerrettinn.
Figure 1: An example of CoR with chains.
To this end, we introduce "Chain-of-Model" (CoM), novel learning paradigm for modeling CoR features. Its core idea is to incorporate causal dependencies across different scales, ensuring that each scale can only use information from its preceding scales at the representation level. To fulfill this, we propose "Chain-of-Layee
(CoL), to reformulate current network layers based on CoR features. Specifically, we require each output chain is only conditioned on the input chains from to i. Figure 2 illustrates an example of applying CoL. into Linear layer. CoL. is capable of three important characteristics: Generality, Causality and Compositionality. That means, 1) Any layer can be viewed as CoL (chain is I); 2) To obtain the features at scale i, CoL allows to only activate the parameters of its correlated chains; 3) If two layers both conform to the CoL settings, their composition also adheres to the CoL property. Following this principle, we can stack multiple CoL layers to generalize from the layer level to the model level, i.e., Chain-of-Modee (CoM). The scope of CoM can be generalized to any ML models all models can be considered as a special case of CoM that chain is set as 1. Besides, the output features of CoM also conform to the characteristic of CoR, CoM can offer multiple sub-mot sub-models at different scales by using different number of chains. Based on these characteristics, CoM can exhibit greater flexibility to build foundation models.
On the basis of CoM framework, we re-formulate the architecture of language models by applying the idea of CoL into each layer within Transformer, and term it as "Chain-of- Laaiiioo--LLnnggg-M-aaaa
(CoLM). CoLM can integrate multi-scale training within single forward propagation, to offer multiple sub- models for elastic inference. also enables us to scale up language models via chain expansion (i.e., increase the number of chains), while simultaneously preserving the capability to use previous scales. Moreover, based on the CoL criteria, we further introduce a KV sharing mechanism in attention module that requires all keys and values to be computed within the first chain, and term it as CoLM-Air. Based on this mechanism, CoLM-Air offers greater extensibility and ficxibility, including allowing us to seamless switching any different scales of L.LMs without any additional re-computations (e.g.. keys and values), accelerating prefilling within the first chain and so on. Experimental results on multiple benchmarks also demonstrate our CoLM family can achieve comparable performance while both exhibit better extensibility and flexibility. Overall, the contributions of our paper can be summarized as below:
We point out the inherent issues of existing LLM scaling approaches, and formulate the concept of chain-of-ne cf representation (CoR), that integrates multi-scale capabilities within single representation.
To model the relationships among CoR features, introduce chain-of-model (CoM) learning and further extend it into language model, termed CoLM. CoLM is able to provide multiple sub-models varying scales within unified foundation model, enabling more flexible and scalable extensions. Experimental results on multiple benchmarks demonstrate that our CoL.M family can achieve com- parable performance to existing LLM frameworks, while simultaneously exhibit better extension extensibility and flexibility in different scenarios (e.g., prefilling, tuning and so on).
2

---

2 Chain-of-Model Learning
In this section, first introduce the concept of CChuin-of-Reppreenttttorrr (CoR), which generalizes the paradigms of existing representation to broader scope. After that, we introduce of (CoL) and "Chain-of- of-Model" (CoM), model CoR from the layer level the model level.
2.1 Representation
Definition 1
For representation x ER,, can always be equiva- lently formulated as concalenation of n. sub-representations, denoted as E(x, n) {x1,-... Lnh. where 24 E IRd and >i=1 d4 D. We term this as chain-of-rerrssenaaiin (CoR) of x.
According to Definition 1, each chain corresponds to each sub-representation (i.e., 2y.) in CoR. By activating the first few chains (i.e., I<ze). it can be used encode information of scale with dimension of sum() Hence, CoR allows us to encode n different scales within one representation. If I, CoR identical the original representation. Figure illustrates the CoR.
2.2 Layer
However, challenge arises is how to design layers to build connections between CoR inputs and CoR outputs, enabling multi-scale feature transformation while both preserving the output features to follow the criteria of CoR in Definition 1. Therefore, we need to maintain each scale can only utilize information from all its preceding scales, and introduce Chain-of Layeer to incorporate causal relationships into hidden states of CoR, as below:
Definition 2 (Chain-of-Layer) For layer = fe(x). where its input x and output can both expressed as the Shain--ffrepreeenttaiin E(x,1) n) and Ely, n). We term it as a Chain-of-Layerr (CoL) where fo(-) satisfies that each ys is only conditioned on xçi.
CoL is capable of three fundamental properties generality, causality and compositionality:
Corollary 2.1 (Generality) Any layer can be viewed as a case of chain-of-layer when n. 1.
In 1, CoL can be considered as the standard network layer to process the conventional represen- tation. In other words, any network layer can be generalized to the form of CoL. Therefore, we can gradually increase layer dimension by introducing additional chains upon the existing chains.
Corollary 2.2 (Causality) Ifa layer fx) satisfies the property of laain-of-layer, that means
its weights 0 can be partitioned into n. independent parts {0i,.... On), where each Og is used to compute ys based on Te... Therefore, y allows the causal computation, i.e,, to obtain feature ye we only need to compute 0zi bused T-gi-
In the CoL design, each scale can aggregate information of all its previous scales to increase scalability and effectively avoid catastrophic forgetting [15]. Therefore, to obtain information of Scale;, this design allows us to only compute its preceding chains (i.e., Ogi), without calculating the whole weights. As a result, Col. can offer multi-scale information within one representation and allow us to dynamically calculate features at varying scales.
Corollary 2.3 (Compositionality) Assume we have two layers y= fi (x) and = f2(y), where x = {(x, n). y = E(y, n) and z = E(z,n) n). If both fi and f2 adhere to the requirements of CoL. their composition z = (f(((()) also satisfy the setting of CoL, i.e., each 2; only depends xei.
More important, CoL also supports compositionality, meaning that stacking multiple CoL. layers also retains the characteristics of CoL.. This property allows us to generalize the scope of Col. from the layer level to the model level.
3

---

2.3 Model
Definition 3 (Chain-of- Modell) For a model 0 with layers, we define it as a chain-of-model (CoM) if each of its layers, from the first to the final, conforms to the chain-of-layer property.
Based on Definition 3, if a model satisfies the criteria of CoM, it also inherits all CoL. properties, such as generality and causality. In other words, any model can be viewed as kind of CoM (i.e., n 1). And CoM can integrate multiple sub-models at different scales within one model, and enable to scale up model from previous capabilities. This capability directly empowered foundation models with better extensibility and flexibility, In the next section, we will investigate how to build language model based on CoM framework.
Architecture
In this section, we describe the details about how apply CoM into language model, including Linear, each module (e.g., embedding, self-attention, feedforward, normalization) within Transformer, and the objective function. Here, term it ofain--o-Langugg--e MooeeeeeeeMelll (i.e., CoLM). Moreover, we further introduce a KV sharing mechanism based on CoLM framework, and term it as CoLM-Air, which provides better flexibility, such prefilling speedup and seamless LLM switching, which a always needs to re-calculate keys and values..
O
O0000000
(a) Linear.
(b) Chain Linear.
Figure 2: Comparisons between Linear and Chain--of-Linee layer.
3.1 Linear
Linear layer is fundamental layer in neural network, which offers linear operation to transform the input features into the output features. The mathematic expression linear layer is as y = Wx + b, where W R D, xD, and b RD, Here, each neuron of is dense connected to the input feature x, so that we always need to compute all elements of W and to preserve the integral semantics yf Following the CoL. setting, introduce
hyper-parameter into the linear layer, termed "Chain", to determine the size of each chain r and y as represented by CoR. Here, we set the hyper-parameter C = (cI...., Vi E RDr, where Dx, c D& and Dy,
the base ratio compute the dimension of each x; R Px and
Ee Dy. Therefore, it is calculated as:
(yi) = Il_ (Wixe1 + by),
(1)
where W; E RD, b; E Dy; and || denotes the concatenation operation. Given this design, the linear layer also satisfies the CoL criteria, and is identical to the standard linear when 1. Here, we term it as Mifin-of-Linear or Linear (Chain) layer, which serves as the foundation to build language model. Figure depicts the comparisons between Linear and Chain-of-LLinaa layer. To facilitate large-scale pre-training, we further provide an efficient implementation with well-designed block-wise sparse kernel, optimizing both data access and all-reduce communication, as detailed in Appendix A.

---

08900668
n00006n0
Qutput Linar
(CURY
Scaled Dot-Product Attention
Scaled Dot-
Scaled Dat-Produrt
repeat
oppel
Q Lnnea
K Unneer
Vilnear
Liier ohhit
(Chemi
Vüinarr
X
x
X
X
(a) Allention.
(b) Chain of Attention.
(c) Chain of Allention KV sharing.
Figure 3: Differences between Attention, Chain of Attention and Chain of Attention with KV sharing.
3.2 Transformer
3.2.1 Multi-head Attention
In MHA, it applies linear operation to first obtain the query, key and value based on the input features, and split each into multiple heads. And then, we perform the attention function over each head (i.e.. query, key and value) and concatenate all heads, followed by linear transformation to produce the final output. Overall, the mathematic expression of MHA can be formulated as:
9ike
(HA = - O k, vi)) = 01ll softmax
vd Vi),
where || = Q(x), l K(x), I1 vi = V(x),
where h is the head number, dy is the head dimension, and Q(-), K(-), V(-), O(-) represent Linear layer to obtain query, key, value and output. To support CoR modeling in the MHA, first directly replace all Linear layers (i.e., Q, K, V. O) by using Chain-of-Lineee layers described in See 3.1. However, due to the dot product operations within Attention(q, k, v), if single head (i.e., q. k,, w)) contains information from 2 or more chains, its output will blend these information and then distribute to each chain, which cannot meet the setting of Col.. Therefore, we expect each head to only contain single scale information. To fulfill this, we design a simple trick that requires the sum of is equal to so that each o; represents how many heads will be allocated for chain This design enables each chain to retain their dedicated query, key, and value, ensuring our attention can satisfy the criteria of CoL, and we term as Sluuin-of-Affention Attention (Chain) module. Figure illustrates the differences between Attention and Chain-of-Attention.
3.2.2 Feed-Forward Network
Feed-Forward Network (FFN) is another crucial component of the Transformer, which composed of multiple linear layer interleaved with non-linear activation functions (e.g.. ReLU [16] or SiLU [17]). To generalize the CoM framework into Transformer, FFN should also comply with the setting of CoL. To achieve this, we just need to substitute each Linear of FFN by using Chain-of-Lineaar layer, and adopt the same of FFN can also conform to the requirements of CoR. The implementation and illustrative example about FFN can refer Appendix A.6.3.
C used Attention (Chain) module. In this way, the output features
3.2.3 Normalization
home))
Normalization is technique to stabilize layer- wise training. In Transformer, we perform normalization [18, 19] the feature level, before attention or FFN layers. To preserve the CoR proper- ties, we apply the normalization function to each chain individu- ally. Specifically, we use Root Mean Square Layer Normalization (RMSNorm) [19] for each chain to align with the L.LMMA archi- tecture. This chain- wise operation not only guarantees that nor-
Adl Namsausttos
Nome
Nomi:
Narm.)
Figure 4: Normalization (CoR).
5

---

malized features maintain their causal structure but also slightly improves training efficiency by enabling group- wiie operations, which can reduce the number of all-reduce communications. Figure 4 depicts the design of normalization layer.
3.2.4 Embedding
Embedding E RVxD usually serves as the first layer in language model that converts the raw data I as the input embeddings E DR. We do not involve any modifications at the Embedding during the training, but when encoding information at scale i, we just need to use embeddings x(: sum) Dxy] corresponding to the first i chains. A detailed example can be found in the Appendix A.6.4.
3.3 KV Sharing
Based on our above designs, Transformer can provide multiple capability across different models scales, enabling flexible deployment for various applications. However, in attention module, each chain also preserves its unique set of keys and values, and cannot bridge the connections between different scales. For example, when transitioning from a small language model (SLM) to a LLM for generation, typically requires re-computing all keys and values for the preceding content. To this end, we further propose a ambitious technique, called KV sharing, where all keys and values are computed only in the first chain, and then shared to all other chains. If the number of keys and values is smaller than query heads, we follow the practices of GQA [20] by repeating keys and values to match the query heads, just as shown in Figure 3c. The design also meets the criteria of CoL. ! Although it slightly affect the performance, this design also exhibits some unique characteristics, e.g., prefilling speedup, and allowing us to seamlessly switch LMS at different scales from CoL.M for generation. To the best of our knowledge, this is the first work to fulfill this functionality, and we term our CoLM with KV sharing as CoLM-Air. More detailed implementation are in Appendix A.6.2.
3.4 Objective Function
Based on the above designs of each component, the output feature x E RD of Transformer network will adhere to the CoR form {(z,,) And finally, we apply a classification layer W E R DxV, to map the feature dimension into the vocab size. Generally, we can adopt the cross entropy loss as the objective function for CMLM training. However, to enable multi-scale prediction, we also need to build individual classification heads for each scale. Motivated by MRL [21], we propose multi-chain
cross- cross-entropy loss to compute each scale as: Loss; = x1), where W e (R XV is subset of W, and L is the cross entropy function. Although the computations of different logits (i.c.. xix can be shared without any additional cost, calculating multiple cross-entropy losses will adopt
still bring some computational overhead during the training. Due to resource limitations, two-stage strategy inspired by MRL. We first pre-train the entire CoLM network using a standard cross-entropy loss calculated only on the output of the final, largest scale. This allows the model backbone to converge efficiently. After pre- -raitraining, we freeze the model backbone and efficiently fine-tune the classification heads for all scales simultaneously.
4 Experiments
4.1
Performance
In our experiments, utilize the SlimPajama dataset [22] the pre-training corpus, comprising 600 billion tokens. Here, we use LLama-2-tokenizerr" to process corpus, with vocab size of 32000. We leverage Fully Sharded Data Parallel (FSDP2) for distributed training, facilitated by the TorchTitan framework [23]. Our training is built upon a cluster of 32 NVIDIA A100 40GB GPUs with gradient accumulation of 4. yielding an effective batch size of 1024. 1 is adopted for training, and the sequence length is set as 4096 tokens. We choose Adam V optimizer [24] with a learning rate of 1.5 10 Due to resource limitations, our model is pre-trained with 50K steps, nearly 200 billion tokens. For model configuration, we use chains C and dimensions D to determine different models. The baseline use C (32), and adopt same configuration as LLaMA-3.2- IB, so it is equivalent to
It can be considered as y is only conditioned on Iı, which belongs to a subset of Definition 2. hnttps://hugiiggtee co/meta- 1lw.laaaaaaaaaa aaamaaaaaaa-----mm 7b

---

the standard language model. For the baseline, we also increase the dimension of our CoLM as Chain-of- than Linear layer. We use the FleutherAI Language Model Evaluation Harness [25] to evaluate models on commonsense tasks in a zero-shot setting. More details can refer to Appendix A.I.
we use C = {16, 16) and C = {8, 8888,8888 8, To align
takes fewer parameters
Our results are reported in Table I, and have these observations: 1) CoLM-Air with KV sharing mechanism will slightly affect model performance, the first chain. However, this design introduces greater flexibility, which will be discussed in the following section; 2) using {16, 16) can achieve better performance when compared with {8. 8, 8, 8). When using {16, 16} with wider size, our method can achieve comparable performance to baseline. can be considered
only calculates keys and values within
trade-off as using more chains can offer more sub sub-models for usage. Overall, CoLM achieves results competitive with baseline while offering significantly faster prefilling speeds and greater flexibility, which will discussed in subsequent sections.
Table 1: Performance of baseline and CoLM commonsense reasoning tasks. Models trained for 50K steps. Each task is reported by acc acc__oce norm metric.
Chain(C)
Dims
Params
HellaSwagt
Obqat
WinoGranda"
ARC-et
ARC-et
Boolgi
Piqat
Avgt
(32)
2048
1.10B
40.01
31.19
52.72
43.52
23.63
57.13
67.30
16.11
CoLM
[16, 16]
2560
1.118
40.25
31.39
32.41
43.73
23.81
58.01
67.30
45.27
(16, 16)
2048
0.86B
37.12
29.18
51.07
40.82
22.61
61.74
65.48
44.00
8,
3072
1.18B
38.63
31.99
51.62
12.80
23.89
56.70
66.00
44.51
18, 8, 8, 8]
2048
0.74B
34.49
27.57
50.20
39.02
22.78
57.65
63.00
42.10
CoLM Air
[16, 16]
2560
1.11B
39.85
31.19
32.09
44.30
23.63
56.72
66.76
44.80
(16, 16]
2048
0.86B
36.82
28.
51.62
49.19
22.70
61.31
65.94
43.90
8,
3072
1.18B
36.98
29.78
49.80
41.46
24.57
55.81
65.51
43.41
8,
2048
0.74B
33.77
27.97
49.41
38.93
22.61
57.65
62.10
41.77
4.2
Chain Expansion
Table 2: Performance of baseline versus expanded models on commonsense reasoning tasks.
Model
HellaSwagt
Obqui
WinoGranda'
ARC-et
ARC-et
Boolgt
Piqui
Seiqt
Avg"
Tiny -LMMAIL
Baseline
61.47 61.66
36.62 36.62
59.43 60.62
55.47 56.27
32.68 32.94
55.99
73.56
84.20
57.43
Expansion
58.44
74.05
86.20
58.35
LLaMA-3.2-1B
Baseline
63.78
36.42 37.02
50.83
60.56
36.03
63.70
74.37
88.40
60.39
Expansion 63.19
60.31
36.69
63.55
74.65
88,20
60.53
Given the generality and causality of CoM, any model can regarded as a special case of CoM when chain is 1, and can be extended to multiple chains. Therefore, we introduce Chain Expansion, which use the well-trained model as the initial chain and then expand it with new additional chains. Conceptually, this is similar to progressive network architectures [12], that allowing us to increase additional capacity while retaining prior knowledge. However, CoLM has extended it to Transformer for training. Overall, Chain Expansion allows us to incrementally increase model scales based on the existing scales, avoid training from scratch.
To validate this point, we choose two LLaMA variants (i.e., LanyLLM [26] and LLaMA-3.. IB [3]) as the first chain for expansion. Since these two models both use group query attention [20], where query number is 32 and kv head is 8. Following the Attention setting in Section 3.2.1, We set c 32 and subsequently introduce a second chain with c2 8, incurring nearly 0.8B parameters. We perform zero initialization for the expanded classification head, so that its output logits is started from the original trained model. Due resource limitations, our expanded models are trained with batch size of 1024 for 20K steps, nearly 8 billion tokens. Besides, we also freeze the first chain to preserve the original knowledge and speedup training. More details can refer Appendix A.1.
All of results are presented in Table 2. We achieve improvements of 0.92 and 0.14 over the Tiny- LLaMA-v1.1 and LLaMA-3.2-II. respectively. Since LLaMa-3.2-1BB is a stronger baseline,
Each chain is equal slightly improve the training efficiency.

---

LLama 3.2 (Same Parameter)
ELarmaJ3 (Seme Paramatert
Parameter
LLamb-3.23 3.1-8-3--8 (Same Parameter)
Eveine
26
30K
73X
100
D4C
235
5125
120k
25.6
$126
(1) Clama-
(2K))
(c)
La (128K
Figure Comparisons on Prefilling speed on LLaMa-IB and LLaMa-8B settings (2K 1M).
requires more computation to achieve significant gains but our method still can improve it under the limited computations. Overall, these results also indicate the effectiveness of our method in improving the baselines, even with constrained resources.
4.3 Elastic Inference
Elastic inference is to provide dynamic inference usage meet the requirements from different deployment scenarios. Compared with previous methods [10, 11] that involved additional sampling policies, CoLM integrates multiple sub-models varying scales within one model, enabling unified pre-training via single forward pass. To obtain sub-models across different scales for prediction, we further pre-train CoLM model using the loss function described in Sec 3.4. We choose the setting of CoLM-Air, where C = {16, 16} and Dims as 2048, for validation. Following prior work [21]. we freeze the Transformer backbone and only train the classification head accelerate the training. The results in Table 3 also highlight the potential of CoLM for enabling elastic inference.
Table Elastic Inference offer multiple sub-models by using different number of chains.
Scale
Params
HellaSwag
Obqa
WinoGranda
ARC-e
ARC-c
Boolg
Piga
Avg
CoI.M-Arr, C = (16, 16), Dims = 2048
Chain +
0.86B
36.82 29.31
28.77 25.75
51.62 52.96
40.19
22.70
61.31
65.94
43.90
Chain 1
0.33B
34.50
22.01
62.23
61.15
41.13
4.4 Prefilling
Prefilling [27] is to generate keys and values in Transformer to understand the prompt of user instruction. Existing LLM paradigms always requires us to deploy the full model for calculation, introducing substantial computational costs. For CoLM-Air, all keys and values computed within the first chain, so that we just need to calculate the first chain according to the causality of CoM. Please note the first chain is still standard language model, techniques (e.g, MInference [28]). Here, we conduct two scales (I.I ama-3.22-1 and a. for evaluation, with a context length from 2K IM. We choose two settings where C {8, 8, 8, 8} and C (16, 16} with similar parameters as baseline. Our results are shown in Figure 5. More experimental details and results can refer Appendix 1.2.
further apply any inference-time
From Figure 5, We observe that CoLM-Air achieves faster prefilling speed when compared to LLaMa with similar parameter sizes. With the increasing of sequence length. CoLM-Air can achieve faster speedup at the prefilling. In particular, to process 1M tokens, CoLM-Air yields nearly 1.6x and 3.0x faster prefilling when is {16, 16) and (8,8, 8, 8,88. 8). And, when combining with MInference [28], CoL.MAiir achieves at most 27 x speedup (1920 - 70.0) in Figure 5d. These results all demonstrate that our M-Ar can effectively speedup prefilling.
4.5 Chain Tuning
Benefiting from the causality of CoM architecture, CoLM is composed of multiple chains, that each chain can leverage capability from previous chains. Therefore, based on this property, we propose Chain Tuning, which aims to fine-tune the latter chains while freezing the first few chains. Chain

---

Tuning can partially reduce the tuning costs and mitigate catastrophic forgetting via preserving the first few chains. Moreover, when we adopt the CoLM-Air setting and freeze the first chain, the keys and values from the fine-tuned models can seamlessly transfer to the original model, without any additional computations. Here, we utilize the expanded models (with two chains) trained in Section 4.2 as the baseline, and then fine-tune the expanded chain.
In addition, compared with Parameter-Efficieen Fine-Tuning (PEFT) methods (e.g., LoRA [29]), our Chain Tuning is fully compatible with these PEFT methods, while further enable KV sharing to switch between the fine-tuned models and the original model. In other words, we can apply LoRA to the latter chains for even greater parameter efficiency. To demonstrate this capability, we choose GLUE benchmark [30] for faster validation. We report the results within Table 4, and more experimental details can refer to Appendix A.1. From Table 4, we observe Chain-tuning can boost model performance by just fine-tuning approximately 42% of the model parameters.
Table 4: GLUE dev set results. Each task is reported by accuracy. "+ CT means Chain Tuning, and the values in brackets mean the fine-tuned parameters and the total parameters.
Model
SST-2
COLA
MNLI
MRPC
QNLI
QQP
RTE
WNLI
Avg.
Baseline
67.89
32.98
35.87
49.02
51.77
51.65
54.87
43.66
48.46
LoRA
91.74
58.46
81.32
61.27
73.20
54.92
55.96
52.11
66.12
CT
92.32
54.87
82.26
62.01
82.87
55.22
59.21
53.52
67.79
4.6
Generative Performance
To assess performance on open-ended generation, we instruction-uun our expanded TinyLLaMA- v1.1 model (from Section 4.2) the Alpaca dataset [31] and evaluated MT-Bench [32]. The results in Table show that our model, which leverages KV sharing, outperforms the similarly tuned baseline, particularly in reasoning and humanities tasks. This validates that the CoM framework and KV sharing mechanism are effective for enhancing instruction-folloingg and conversational abilities.
Table 5: MT-Bench scores for instruction-tuned models. "Ours" refers our expanded model. Both models are fine-tuned on the Alpaca dataset. Higher scores are better.
Model
Writing
Roleplay
Reasoning
Math
Coding
Extraction
STEM
Humanities
Avg.
TinyLLaMA-v/
2.85
3,20
2.05
1.10
1,30
1,40
2.35
3,20
2.18
Ours
2.75
3.40
2.70
1.00
1.10
1,30
2.35
3.65
2.28
5
Related Works
Foundation Models play critical component in machine learning for effectively modeling relation- ships among data. From the perspective of architectural level, we posit that existing Large Foundation Models (LFM), can be conceptually decomposed into two parts, Base Architecture and Scaling Archi- tecture. For Base Architecture, refers to the backbone network (e.g., Transformer [7], Mamba [33]. RWKV [34]) to perform feature transformation from the original data. For Scaling Architecture, it represents the scaling strategies (e.g.. Increasing layers, dimensions or experts [35, 14]) to improve model capabilities. To some extent, CoM can also be regarded as a scaling architecture via increasing the number of chains, so that it can be applied to any base architectures. Morcover, compared with Dense or MoE architectures, the core idea of CoM is to integrate capabilities across varying scales and allow us to leverage inherited knowledge from previous models to improve scaling efficiency and avoid catastrophic forgetting.
Elastic Inference to offer multiple sub-models tailored handle diverse user instructions under different deployment constraints. Current LLM architectures can only provide a fixed inference capacity, and thus requires us to pre-train multiple models at different sizes, leading to inefficiencies. To this end, Once-for-All [10] build super-net and then randomly sample specialized sub-network for training. Some other works [11, 36] introduce some nested architecture designs into network blocks to offer clastic inference across different model granularity, but also require sampling strategy

---

of each sub-network. By incorporating causal dependencies among sub-models within CoR, CoM can integrate multi-scale training within single forward pass, demonstrating better efficiency.
Continual Training is to continue the training process from the pre-trained models to boost its capability and avoid catastrophic forgetring. To leverage the capability of existing pre-trained models, the most simple method is using depth-up scaling method (i.e., increase the number of layers), such as Solar 10.7B [37]. Depth-up scaling can effectively reuse weights from previous layers for initialization, but fails to retain the output logits (i.e., probability) of the classification layers. CoM can be regarded as a width-up scaling method, which can perfectly preserves the knowledge of the entire model from the first layer to the final layer. Please note that CoM is also compatible with the depth-up scaling approach. Besides, Progressive neural network (PNN) [12] is seminal work in continual training, which also motivates us to develop CoM. PNN extends the MLP layers by incorporating new parameters via lateral connections, and then apply to the reinforcement learning tasks. Compared to PNN, CoM extends its scope to broader range and effectively demonstrates the feasibility of using Transformers for pre-training.
6 Conclusion
In this paper, we introduce the concept of information within hidden representation. Subsequently, we further propose "Chain-of-Modde -f--Meee" (CoM) learning to capture the causal relationship across varying scales among CoR features of each layer within the model. The model built upon the CoM framework can provide multi-scale functionality and retain capability from previous scales for usage. Therefore, we devise (CoM) paradigm into each layer of the Transformer
(CoR) to encode multi-scale
(CoLM) by integrating the architecture. As result, CoLM can offer multi-scale sub-model at varying scales within a unified language model. Moreover, we propose a KV sharing mechanism that computes all keys and values within the first chain. This design significantly unleashes the flexibility and extensibility of language models, including prefilling acceleration, dynamic LLM switching and so on. Experimental results across multiple datasets demonstrate that our CoLM family matches the performance of standard Transformers, while offering superior extensibility and adaptability across a wide range of scenarios, paving new way to build next- coxt-generation foundation models.
Limitations
While Chain-of- Model (CoM) introduces flexible and efficient paradigm for scaling language models, some challenges and areas for future exploration remain. First, the Chain-of-Linearr layer, in its current form, presents complexities for Tensor Parallelism (TP), a key technique for training extremely large models. Although compatible with other parallelism strategies, optimizing CoM for TP-centric infrastructures requires further investigation. Second, in the CoLM-Air architecture, the first chain's quality is critical, as it generates the keys and values for all subsequent chains. Determining the optimal configuration and training strategy for this foundational chain is crucial for maximizing the overall model's 's potential. We provide more detailed discussion of these limitations and outline potential future directions in Appendix A.5.
References
[1] OpenAL. GPT-4 technical report. CORR, abs/2303.08774, 2023.
[2] Google Gemini Team. Gemini: A family of highly capable multimodal models. CoRR, 12 2023.
[3] Meta LLama Team. The Ilama herd of models. CoRR, abs/2407.21783,, 2024.
[4] Alibaba Qwen Team. Qwen2.5 technical report. CoRR, abs/2412.151155 2024.
[5] Deepseek Team. Deepseek-v3 technical report. CoRR, abs/2412. 2024.
[6] Mistral AI Team. Mixtral of experts. CoRR, abs/2401.04088. 2024.
10

---

[7] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems 30: Annual Conference on Neural Information P Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.
[8] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08235 2020.
[9] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas. Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. CoRR, abs/2203. 15556, 2022.
[10] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize it for efficient deployment. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[11] Devvrit, Sneha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit S. Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi. Sham M. Kakade, Ali Farhadi, and Prateek Jain. Mat- former: Nested transformer for elastic inference. CoRR, abs/2310.07700 2023.
[12] Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. CoRR, Nass1606..4671, 2016.
[13] Liyuan Wang, Xingxing Zhang, Hang Su, and Jun Zhu. A comprehensive survey of continual learning: Theory, method and application. IEEE Trans. Pattern Anal. Mach. Intell., 46(8):5362- 5383, 2024.
[14] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Comput., 3(1):79-87, 1991.
[15] James Kirkpatrick, Razvan Pascanu, Neil C. Rabinowitz, Joel Veness, Guillaume Desjardins, An- drei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks. CoRR, abs/1612.00796. 2016
Demis
[16] Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on Machine Learning (ICML. 10), June 21-24, 2010, Haifa, Israel, pages 807-814, 2010.
[17] Noam Shazeer. GLU variants improve transformer. CoRR, abs/2002.05202, 2020.
[18] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, abs/1607.06450. 2016.
[19] Biao Zhang and Rico Sennrich. Root mean square layer normalization. In Advances Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 12360- 12371, 2019.
[20] Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pages 4895-4901, 2023.
[21] Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyded Kaifeng Chen, Sham M. Kakade, Prateek Jain, and Ali Farhadi. Matryoshka representation learning. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NewrIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
11

---

[22] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Joel Hestness, and Nolan Dey. SlimPajama: 627B token cleaned and deduplicated version of RedPajama. s1 jmampaae amaaaaaaaaaaa a- a-tajaaaaaa---t-oeeeeeaa toke-t - ce of offd ama,
Jacob R Steeves,
https: https://cerebraat ai /
of
June 2023.
[23] Wanchao Liang. Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, and Stratos Idreos. Torchtitan: One-stop pytorch native solution for production ready LLM pre-training. CoRR, abs/2410.06511,, 2024.
[24] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 3rd International Conference on Learning Representations, ICLR 2015, San Diego. CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[25] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
[26] Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. CoRR, abs/2401.02333 2024.
[27] Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ramachandran Ramjee. SARATHI: efficient LLM inference by piggybacking decodes with chunked prefills. CoRR, abs/2308. 16369, 2023.
[28] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Cei Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context ilms via dynamic sparse attention. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[29] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation large language models. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April
25-29, 2022. OpenReview.net, 2022.
[30] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[31] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https:// //gi thub. ctoo/tatsu-111 haabbstaafordd ataaaaod_aapaca, 2023.
[32] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang. Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric P. Xing, Hao Zhang. Joseph E. Gonzalez, and Ion Stoica. Judging llm-as-a-judde with mt-bench and chatbot arena. In Alice Oh, Tristan Naumann. Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[33] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CORR, abs/2312.00752 2023.
[34] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, Xingjian Du, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Jiaju Lin, Krishna Sri Ipsit Mantri, Ferdinand Mom,
12

---

Atsushi Saito, Guangyu Song, Xiangru Tang. Johan S. Wind, Stanislaw Wozniak, Zhenyuan Zhang, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. RWKV: reinventing mns for the transformer era. In Findings of the Association for Computational Linguisties: EMNLP 2023, Singapore, December 6- 6-10, 2023, pages 14048-14077, 2023.
[35] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gater mixture-of- experts layer. CoRR, 0tos/1701.065344 2017.
[36] Abhinav Shukla, Sai Vemprala, Aditya Kusupati, and Ashish Kapoor. Matmamba: A matryoshka state space model. CoRR, abs/2410.06716 2024.
[37]
Sanghoon Kim, Dahyun Kim, Chanjun Park, Wonsung Lee, Wonho Song. Yunsu Kim, Hyeon- woo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, Changbae Ahn, Seonghoon Yang, Sukyung Lee, Hyunbyung Park, Gyoungjin Gim, Mikyoung Cha, Hwalsuk Lee, and Sunghun Kim. SOLAR 10.7b: Scaling large language models with simple yet effective depth up-scaling. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies: Industry Track, NAACL 2024, Mexico City, Mexico, June 16-21, 2024, pages 23-35. Association for Computational Linguistics, 2024.
[38] Tri Dao. Flashattention-2:: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
[39] Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. Distserve: Disaggregating prefill and decoding for gpodput-optimizee large language model serving. In 18th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2024, Santa Clara, CA, USA, July 10-12, 2024, pages 193-210. USENIX Association, 2024.
[40] Pratyush Patel, Esha Choukse, Chaojie Zhang, Aashaka Shah, fñigo Goiri, Saeed Maleki, and Ricardo Bianchini. Splitwise: Efficient generative LLM inference using phase splitting. In 5/st ACM/IEEE Annual International Symposium on Computer Architecture, ISCA 2024, Buenos Aires, Argentina, June 29- Juuly 3, 2024, pages 118-132. IEEE, 2024.
[41] Ruoyu Qin, Zheming Li. Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu. Mooncake: A kvcache-centric disaggregated architecture for LLM serving. CoRR, abs/2407.00079, 2024.
[42] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott E. Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In IEEE Conference Computer Vision and Pattern Recognition, CVPR 2015, Boston, MA, USA, June 7-12, 2015, pages 1 1-9. IEEE Computer Society, 2015.
[43] William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Oogan-Kelley, Reducing transformer key-value cache size with cross-layer attention. In Ad- vances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[44] Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, and Alham Fikri Aji. MLKV: multi-layer key-value heads for memory efficient transformer decoding. CoRR, 02t/2446609297,, 2024.
[45] Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for lan- guage models. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NewrIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
[46] David Raposo, Samuel Ritter, Blake A. Richards, Timothy P. Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of-depths: Dynamically allocating compute in transformer-basse language models. CoRR, abs/2404.02258, 2024.
13

---

[47] Dmitry Lepikhin, HyoukJoong Lee. Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang. Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. CoRR, abs/2006. 1772006.16668,, 2020.
[48] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng. Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-Ou caaere-ff---pee experts language models. CoRR, abs/2401 1.00666,, 2024.
14

---

A Appendix
A.1 Experimental Details
Datasets SlimPajama [22] is a deduplicated, multi-corpora, open-source dataset for training LLMs. We use SlimPajama dataset exclude books corpus as the pre-training corpus, result in approximately 600 billion tokens. The composition of our used dataset is shown in Table 6.
Dataset
Ratio
Commoncrawl
52.2%
C4
26.7%
GitHub
5.2%
ArXiv
4.6%
Wikpedia 3.8%
StackExchange
3.3%
Table 6: Dataset statistical of SlimPajama dataset.
Hyperparameters Table presents the hyperparameters that were used in pretraining. Besides, P [38] is applied to speed up attention computation. For the Expansion experiment, we set the training step to 20,000, while preserving the same configuration as the pre-training. Table presents the hyperparameters used for Chain Tuning.
Evaluation In this project, we utilized the Eleuther AI Language Model Evaluation Harness library [25]. We used the default evaluation configuration. For chain tuning, we set the prompt template the default from the evaluation scripts.
Table 7: Hyperparameters used in pretraining.
Hyperparameter
Value
Learning Rate
2e 4
Sequence Length
4096
Batch Size Per Device
Gradient Accumulation Steps
8
Optimizer
Adamw
LR Scheduler Type
Linear
Warmup Steps
2000
Max Norm
1.0
Activation checkpointing
full
Mixed Precision for Parameters biloat16
Mixed Precision for Reductions
float32
Table 8: Hyperparameters used in Chain Tuning.
Hyperparameter
Value
Learning Rate
2e 4
Epochs
Batch Size Per Device
Gradient Accumulation Steps
32
Optimizer
Adamw
LR Scheduler Type
Cosine
Warmup Ratio
0.1
Sequence Length
4096
Model Configuration Table 9 presents the detailed configurations for all the models used in our experiments. Each model identified by name and characterized by configurations such as the Chain setting, dimension (dim), hidden dimension (hidden dim), the number of layers, attention
15

---

Table 9: Model Configuration.
Model Name
Chains
Params
Dim
Hidden Dim
Layers
Heads
Heads
sspansion-llamats BBB-BB
(32. 8]
1.93B
2560
10240
16
40
10
Sspansion-tinyllann
(32, 8)
1.44B
2560
7040
22
40
Llama
[32]
1.10B
2048
8192
32
CoLM 16-16
(16. 16]
0.86B
2048
8192
32
CoLM_8
24}
0.92B
8192
32
CoLM 8-8-16
[8. 16)
0.80B
2048
8192
32
CoLML 8-8-8-8
(8, 8, 8]
0.74B
2048
8192
32
Col 16-16-6ame (6-same-params
{16, 16]
1.11B
2560
8192
32
CoLM 8-8-8-8-same-parree
(8, 8,
1.18B
3072
8192
16
32
heads, key-value heads and the number of parameters. Figure 12 illustrates the structure of the Chain of Transformer, designed for multi-scale representation learning.
A.2 Full Results
A.2.1 Prefilling
Our prefilling is evaluated on single NVIDIA A100 GPU with batch size of 1. Specifically, from 128K to IM, we perform chunk-wise prefilling (e.g., chunk operation for layer calculation) to avoid the out-of-memory (OOM) issue, and we also evaluate and validate our method when combining with MInference [28]. Here, we evaluate six configurations by using different chain settings ([16, 16]. {8, 8, 8, 8). {8, 8, 16). {8, 24)) and increase the number of layers ensure the parameter sizes comparable. Table 10 and Table 11 report all results of our CoLM-Air setting (IB and 8B) sequence length from 2K to 64K. And Table 12 and Table 13 report all results of our CoLM-Air setting (1B and 8B) of a sequence length from 128K to IM. From these results, we can find that CoLM can significantly speed up the prefilling stage in caleulating long-context instructions.
Moreover, our CoLM is also extremely suitable for PD disaggregation architecture [39-41], as only requires us to deploy the entire weights of the first chain on the prefilling server to calculate KV caches, which significantly reduces the computation loads of the prefilling server.
Table 10: Prefilling results on LLama-3. 3.B-IB setting (2K - 64K). All results are reported in millisec- onds. xx, represents the dimension of the first chain.
Length
Method
L
D(Dx,)
2K
4K
8K
16K
32K
64K
LLama-3.2-1B
16
2048 (2048)
51.0
62.0
132.0
313.0
997.0
2,600
Same configuration (CoLM-Air)
(16.16)
16
2048 (1024)
47.0
47.6
60.0
140.0
380.0
1,189
(8, 8, 8, 8) 16
2048 (512)
48.0
50.0
51.0
71.0
193.0
613.0
Similar parameters (CoLM-Air)
(16,16)
22
2048 (1024)
66.0
69.0
78.0
183.0
503.0
1,587
[8, 8, 8, 8)
26
2048 (512)
76.0
76.0
78.0
105.0
293.0
929.0
(8,8,16)
24
2048 (512)
68.0
69.0
73.0
99.0
272.0
863.0
(8,24)
21
2048 (512)
62.0
60.0
65.0
87.0
243.0
768.0
A.2.2 CIFAR-10 MIP Experiments
We compare baseline fully-connected MLP against several Chain-of-Modd MLP (CoM-MLP) variants on CIFAR-10. The baseline model uses a 4-layer architecture with hidden sizes [768, 1024, 512, 256] (an input embedding of dimension 768, followed by three dense layers and a 10- classes output). CoM-MLP models augment each layer with multiple sequential conditional modules specified by the "Chain-setting" vector. Models were trained using the Adam optimizer with initial
16

---

Table 11: Prefilling results on LLama-3.1 setting (2K 64K). All results are reported in seconds.
Dx, represents the dimension of the first chain. Dx, represents the dimension of the first chain.
Length
Method
D(Dx,)
2K
4K
8K
16K
32K
64K
LLama- -Lama-3.1144 3.1-8B
32
4096 (4096)
0.19
0.38
0.79
1.72
4.18
11.55
Same configuration (CoL.M-Air)
(16,16)
32
4096 (2048)
0.09
0.14
0.28
0.64
1.64
4.89
{8, 8, 8, 8)
32
4096 (1024)
0.09
0.10
0.12
0.28
0.72
2.30
Same parameters (CoLM-Air)
[16,16)
43
4096 (2048)
0.13
0.18
0.37
0.85
2.17
6.49
[8, 8, 8, 8)
52
4096 (1024)
0.13
0.16
0.19
0.43
1.06
3.61
(8,8,16)
48
4096 (1024)
0.13
0.15
0.17
0.34
0.94
3.35
[8,24)
40
4096 (1024)
0.11
0.14
0.15
0.85
2.17
2.82
Table 12: Prefilling results on LLama-3.2-1B setting (128K IM). All results are reported in seconds.
The value in the bracket is using the MInference trick. Dx represents the dimension of the first chain.
Length
Method
D(Dx,)
128K
256K
512K
1M
a.mma-3.2-18
16
2048 (2048)
8.6 (5.5)
32.4 (10.3)
132.0 (21.4)
522.0 (45.8)
Same configuration (CoLM-Air)
(16, 16)
16
2048 (1024)
4.0
15.4 (4.5)
63.2 (9.7)
265.0 (20.1) 124.0 ( 9.9)
(8, 8, 8, 8)
16
2048 (512)
1.9 (1.0)
7.5 (2.2)
30.6 (4.8)
Same parameters (CoLM-Air)
(16, 16) (8, 8, 8, 8)
22
2048 (1024)
(2.9)
(6.2)
88.8 (13.1)
352.0 (28.4)
26
2048 (512)
(1.6
12.5)
50.7 ( 7.4)
213.0 (15.5)
(8, 16)
24
2048 (512)
(1.5)
11.3 (3.1)
46.6 ( 6.7)
188.0 (14.5)
(8, 24)
21
2048 (512)
(1.3)
9.9 (2.8)
40.7 ( 6.0)
165.0 (13.3)
learning rate of 10-3, a batch size 64, for a total of 20 epochs. Reported accuracies are the mean + standard deviation over independent runs.
Table 14 summarizes the test aceuracy and parameter count for the baseline MLP and CoM-MLP configurations. While all models achieve similar accuracy (around 54-55%), the CoM-MLP with Chain-sctting (4, 12) achieves the best accuracy (55.11 + 0.21%) with fewer parameters compared to the dense bascline.
A.3 Training Loss Curves
We conducted several small-scale experiments (batch size 8, 15K steps) to study how different CoLM configurations and settings affect training convergence. The goal was to understand the impact of CoM design choices on convergence speed and final loss. Figures 6 plot the training loss over time for each experiment.
Comparison of CoLM Configurations Figure 6(a) shows the training loss curves for four chain
configurations: equal four-chain (8, 8, 8, 8), two 16-unit chains (16, 16), and two mixed settings (8, 24) and (8, 8, 16). The 16-16 configuration converges more quickly and reaches the lowest loss, indicating more efficient training, which appears yield faster convergence and lower loss under these conditions, suggesting trade-off in chain design.
17

---

Table 13: Prefilling results on LLama-3.1-88 setting (128K 1M). All results are reported in seconds.
The value in the bracket is using the MInference trick. Dx, represents the dimension of the first chain.
Length
Method
L
D(Dx,)
128K
256K
512K
1M
LLama-3.1 3.1-BB
32
4096 (4096)
36.0 (21.6)
134.0 (49.5)
512.0 (106.0)
1920 (186.0)
Same configuration (Col.M-Air)
[16, 16)
32
4096 (2048)
15.6 (6.4)
60.0 (13.3) 28.1 ( 5.1)
244.8 (28.5)
927.0 (89.0)
18. 8) 32
4096 (1024)
7.3 (2.5)
117.3 (10.9)
472.0 (44.7)
Same parameters (CoL.M-irr
[16, 16) [8, 8,
43
4096 (2048)
20.6 (10.8)
77.7 (23.5)
314.8 (45.6)
1293 (128.0)
52
4096 (1024)
12.0 6.2)
46.4 (12.5)
192.4 (25.4)
738 ( 70.1)
{8. 8. 16)
48
4096 (1024)
11.0 ( 5.9)
40.5 (10.7)
167.4 (23.9)
680 65.0)
{8, 24)
40
4096 (1024)
9.2 ( 4.2)
35.4 ( 9.7)
146.0 (19.6)
591 ( 56.1)
Table 14: CIFAR-I0 classification accuracy (mean + std) and parameter count for the baseline MLP and CoM-MLP variants.
Model
Chains
#Params
Accuracy ( (%)
MLP (Dense)
3,806,218
54.93 0.23
CoM-MLP
[4, 4, 4, 3,263,754
54.60 0.16
CoM-MLP
[8, 8)
3,443,978
54.68 + 0.04
CoM-MLP
[4, 12)
3,534,090
55.11 0.21
CoM-MLP
[4,4, 8]
3,353,866
54.55 + 0.36
Loss Across Different Heads We also examined head-wise training by optimizing each chain output head together in the [8, 8, 8, 8 CoLM setting. Figure 6(c) plots the training loss for each of the four heads. The results reveal that larger chains effectively leverage the outputs of preceding chains, resulting in greater capacity and improved performance. This trend demonstrates that deeper chains benefit from the cumulative computations earlier chains, enhaneing their ability to model complex features. The intermediate heads show intermediate performance, reflecting gradual increase in capacity along the chain depth. These findings highlight the hierarchical structure of CoLM, where later chains achieve stronger capacity by building the representations provided by earlier ones.
Effect of KV Sharing Figure 6(c) compares the {8, 8) setting with and without the KV-sharing mechanism (CoLM-Pluss. While KV sharing results in a slightly higher loss for the same head, it
Teating
Cumparisnn CoL.M Configurations
Lose Actune
Effect of
Figure 6: (a) Training loss for different CoLM chain configurations. (b) Training loss for each output head the (8, 8, 8) model (heads 1-4). (c) Training loss for the 88, 8, 8) configuration, with versus without KV sharing.
18

---

Generation (Understand)
Understand
Understand
Encoder
Encoder
Decoder
(a) Encoder
(b) Secoder-Deccoder
Generation
Generation
(Understanc)
Understand
Decoder
Chain 7
Decoder
(c) Decoder
(d) CoLM-Air
Figure 7: Illustrative comparisons between Encoder, Encoder-Decoder, Decoder and CoI.M-Air, in term of understanding and generation capabilities.
significantly improves prefilling speed, as discussed in Section 4.4. This trade-off highlights the efficiency gains of KV sharing without a substantial impact on performance.
A.4 Discussion
Question 1 Chain-of-model architecture seems to be a major improvement in the width of language models (i.e., dimension). Can it be applied the depth of language models (i.e., layers)?
Yes, the existing mechanism of chain-of-mode design requires us to first determine the depth of large foundation models, and then allows us to extend its capability from the perspective of the dimension (i.e., width). Actually, we also conduct preliminary experiments explore the feasibility of applying the idea chain-of-moot -odel across the depth (i.e., layers) of the language model. Here, when applying CoM into the model depth, also require it to meet two characteristics to ensure the alignment with the settings of our current version (i.c., CoLM):
1. The model should encompass multiple sub-models to enable dynamic selection and inference; 2. Each sub-model can seamlessly reuse past states (e.g., keys and values in Transformer) that are generated by different sub-models.
To meet these requirements, we have made two major designs based on the Transformer architecture: (a) we select some intermediate layers and then append the classifier to each as sub-models; (b) we limit keys and values are only produced by the first sub-model (the first few bottom layers), which are shared across the subsequent layers. In our internal experiments, the setting (a) can guarantee that the model can be optimized and reach convergence, while can also obtain multiple nested sub-models. Previous researches (c.g., GoogleNet [42]) also validate the availability of such solution. But this design also introduces additional complexity, extremely cumbersome and redundant. Although we can tie the weights of each classifier together to reduce model parameters, it will also cause interference to the model performance and cannot any computations to calculate each classification head. However, when further integrated with the setting (b), which delivers the keys and values from the bottom layers the top layers, we find that severely degrades model performance. In other words, only using bottom-level states (i.e., keys and values) for all Transformer layers is not enough to create high-level semantics for next-token prediction. Existing attempts [43, 44] also indicate that we can only restrict key and value sharing adjacent layers. which identifies the importance of depth abstract high-level feature states (i.e.. keys and values). Besides, some recent works (c.g., YOCO [45]) have explored reusing features from the outputs of half of the layers, which introduces modifications that combine efficient self-attention and cross-atts attenttioonionn mechanisms. However, this technique is restricted to KV cache reduction during prefilling and does not enable multiple sub-model inference. Some other works (e.g.. MoD [46]) introduce routers to dynamically skip certain layers for faster inference. But this design requires us to reset the keys and values to zero for the skipped layers, which potentially introduces discrepancies between training and inference. In our CoM method, we can retain all keys and values at each layer, enabling inference at multiple scales. Generally speaking, abstract high-level features in neural network, depth usually matters more than width and and preserving invariant depth (i.c., layers) for
need to maintain multiple classifiers, which is
19

---

foundation models is extremely necessary. This can also explain why our current version can only be upplied to width, rather than depth. We will continue exploring solutions that extend the idea of chain-of-model to both width and depth in the future.
Question 2 What are the differences between the Chain-of-Model (CoM) and Histure-oof-xperii (MoE) architectures?
Table 15: Comparisons between MoE and CoM from different aspects.
Mixture-of-Epper (MoE)
Chain-of-Model (CoM)
Expert Capability
Equivalent
Weak - Strong
Expert Range
FFN
Full Model
Expert Activation
Sparse
Nested
The concept of Mixture of Experts (MoH) was first proposed by [14], which aims to create multiple expert networks and then use gating mechanism for selective activation. As Transformer has exhibited huge potential in scalability, some works [35, 47] have investigated how to apply the MoE architecture into the FFN layers within Transformer to build the giant network. On the basis of this, some works [6, 48] have successfully scaled up language models with MoE architectures to a substantial size. We deem that MoE can also be considered as a kind of scaling architecture. Existing MoE architectures in Transformer are designed to create multiple experts with similar capabilities at the FFN layers, while only few experts are activated for each token prediction. In the CoM structure, it aims to create sequence of nested experts, progressing from weak to strong capability. Specifically, each expert contains the capability of all its preceding experts (i.e., weaker experts). In Table 15, we also illustrate the comparisons between CoM and MoE architectures. Therefore, CoM and MoE are designed from different perspectives to construct experts. In other words, CoM completely orthogonal to MoE, meaning both of them can be applied within the same architecture and inherit advantages from each. We will explore how to combine CoM and MoE architecture, and leave it as future work.
Question 3 When do we need the Chain-of-Marn Modeel architecture?
Generally, Chain-of-Model (CoM) architecture can be considered as an innovative design to optimize the scaling architecture (e.g., increasing model capabilities) of foundation models, rather than altering the backbone network (e.g., Transformer). Therefore, it can be generalized to any model with different architectures, such as CNNs or parallelized RNNs. However, this does not mean we can apply CoM architecture to any scenario. Based on the designs of CoM architecture, it can be considered as a composition of multiple nested sub-models. Each sub-model will contribute the subsequent models. More specifically, the first chain (sub-model) plays a critical role in CoM architectures as it directly determines the understanding capability of the entire model. Therefore, the capacity of the model is not sufficient to exhibit the power of the scaling law or lacks enough generalization, it unnecessary to use the architecture the chain-of-model method. In other words, the CoM structure is better suitable for large-scale model sizes, rather than smaller ones. This also indicates that CoM is an advanced design to optimize the scaling architecture
the capability of
A.5 Limitations
From the perspective of methodology, Chain-of-Model (CoM) introduces an innovative and brand- new solution to revolutionize the scaling architecture of large-scale foundation models. However, CoM still remains some ongoing challenges when applying it to practical applications:
Although the CoM architecture offers exceptional flexibility in utilizing LLMs, also poses new challenges for training a large-scale CoM architecture at the infrastructure level. Specifically, Chain-of- layer is compatible with data, pipeline, and context parallelism (DP, PP, CP), but not well-suited for tensor parallelism (TP). In the naive implementation (Please see Figure 2b or Algorithm 1), the Chain-of-Linnar layer can be viewed as a composition of multiple sub-linear layers. As result, if we apply tensor parallelism to each sub-linear within the Chain-of-Linn Linear
20

---

layer, it will significantly increase the number of all-reduce operations and data access when compared with the standard Linear. Although block-wise sparse kernel optimization can help us to alleviate these issues, we still need to devise new pipeline mechanism when handling large tensor inputs. To some degree, Chain-of-Linnarr layer can be considered as imbalanced split-wise tensor, which requires to specifically design tensor parallelism mechanism for it. Here, we leave as future work.
Empowered by chain-of-ne model design, CoLM is novel architecture that combines multiple nested language models at varying scales. Benefiting from such a design, each model will inherit all its nested sub-models and extend their capability. In other words, each nested model will affect its subsequent model. More specifically, CoLM further introduces a K-rreduction mechanism to guarantee that the understanding parts (e.g.. keys and values in Transformer) are all produced by the first chain. That means the first chain plays very important role building CoLM as determines the understanding capability of the entire model. So, compared to the other chains, expect the first chain to pay more attention to expressive and generalized content states (i.e., keys and values), in relative to its generation. Therefore, how to explore the optimal setting for the first chain could be very important in the CoLM framework. However, due to resource limitations, we leave it future works.
Algorithm 1 Pseudo Code for Chain-of Lfnear Layer (Linear chain.py)
class Linear Module)
def -..itt__
self, dim: int, hidden_dine int, chains: List[int]] bias False
)
Dimension of input features hidden_ dim: Dimension of output features
dim:
chains:
Base of each chain whether using bias or not
bias:
super ()___intt__
self num chain len (chains)
self. in_dims [dim с // sum (chhaiiss)) for in np. cumsum cumsum (chains)] // sum sumhciinns) for с in chains]
self nut_dims [hidden_dim self self.mlps = mn ModuleList([ Linear (indim, outdim, bias=bias)
for indim, outdim in zip(self in_dims, self self.out_dis out_oinnnnn
])
def forward(self,, x: torch. Tensor, chain: int None) -> torch. Tensor:
Input features, shape (B, L, The number of activated chains ([O, num chain 1])
chain:
if chain is None:
chain self num num_chair
assert (
<= chain self num_chain
), "The chain id should be in [o, (self. num_chain)).
outputs
self wwt.mmpsill. (x[..., :indim])
for i, indim in cumerate(self in dims chain 1])
return torch. cat (outputs, dim=-1)
21

---

Weight Tensor W
Weight Tensor
Input Gradient
xigpes
suppal
Skiped
Impul Tensor
Oulpul Tensor Y
Oulput Gradient
Input Gradient
Oulpul Gradienl
Weight Gradienl
(a) Forward.
(b) Backward for inpul gradient.
(c) Backward for weight gradient.
Figure 8: Block-sparse GeMM kernel for Linear (Chain) layer.
Speed comparisons between different kernels. The "fwd", "bwdw" " and "bwd;" repre- sent the speed of forward, backward (weight) and backward (inputs). Chain-of-Linear (Naive) refers to Algorithm I, and Chain-of-Linear (Triton) corresponds to our block-wise sparse kernel implementation.
Table 16:
Computation
Implementation
Dim
Params
fwd
bwd
bwd:
C 8,,8,8, 8, lensor size is (4096, 4096).
Linear
4096
167K
0.67
0.62
0.64
Chain-of-LLneen (Naive)
4096
104K
0.58
0.49
0.72
Chain-of-L (Triton)
4096
104K
0.47
0.46
0.56
44,4, 4, 4, 4, 4, 4), tensor size (8192, 8192).
Linear
8192
671K
1.56
1.61
1.56
Chain-of-LLinaa (Naive)
8192
419K
0.96
0.84
1.82
Chain-of-LLnnaa (Triton)
8192
419K
0.77
0.76
0.90
A.6
Implementation
In this section, we will provide detail implementation about each module used in CoLM, including Linear, Attention, FFN, normalization and so on.
A.6.1 Linear
In Algorithm provide the naive PyTorch implementation of Chain-of- layer. Generally, the proposed Chain-of-LLined layer can be viewed as a composition of multiple sub-linear layers. Under the same dimension, our implementation (i.e., Algorithm 1) reduce more weights when compared with the standard Linear layer, but it will add the number of data loading and all-reduce communication. Therefore, further devise block-wise sparse kernel to accelerate the computation of Chain-of-Linear structure, which is introduced in below.
Block-wise Sparse Kernel Acceleration We developed a block-wise sparse matrix multiplication kernel for both forward and backward paths of Chain-of-Linear layer. To conserve GPU memory, we store the weights of Linear (Chain) layer in block-compressed seaddd-par---ooo now (BCSR) format and perform matrix multiplication only on the step-shaped mask accelerate computation. The kernel -aat--MMat--iipllllle
implementation is based on the Triton both sparse input and sparse output. For sparse input, each thread block processes row the BCSR matrix, skipping unmapped data from the other dense input matrix. For sparse output, thread blocks in the masked area will exit immediately. Figure illustrates the details of our block-wir sparse kernel for acceleration. We also report the computation comparisons between our kernel with naive
(GeMM) kernel and supports
22

---

Algorithm 2 Pseudo Code for Attention (Chain) Layer.
from flash_attn import Baash_attn_fume
from Linear_ chain import Linear
class Attention(m) (neettioo(n.MMdulee))
def __init__ self, dim: int, n heaed int, n_kv_head: int, chains: List [int]
): super()...inite
assert sum
n_head
for chain in chains:
assert chain n_kv_head n_head
self headddiadd dim // n_head
self xq
Linear (dim, dim, chains)
self.wt wk Linear (dim, n_kv_head self www.heeaimdn chains)
self.w Linear (dim, n_kv_head * self head_dims, chains)
self no Linear (dim, dim, chains)
def forward
x: torch. Tensor, freqs_cis: torch. Tensor, chain: int = None) -> Tensor::
baz, seqlen, dim x. shape
q, k, v self (x chain), self.wk(x, chain), self (x chain)
q q. v.viewbszz seqlen, -1, self head_dim) k k. v.view (bsz, seqlen, -1, self.head head_dim)
visw seqlen, -1, selfhhead dim)
9, x apply.rotary_emle k, freqs_cis)
0 flash_attn_funce k, v, causal=True)
o o. ovview seqlen, -1) return self wo chain=chain)
implementation in Table 16. From Table 16, we observe that the naive implementation exceedingly slow in computing backpropagation gradients for the input data, while our kernel significantly reduces its computation for input data, and then improves the speed of both the forward and backward stages. Besides, our kernel can also reduce the number of all-reduce communications during the distributed training. Benefiting from this design it allows us to conduct large-scale pre- pre-training efficiently.
A.6.2 Attention
Just as aforementioned in Section 3.2.1, we can replace all linear layer (i.e., Q, K. V, O) by using Chain-of-Linear layer. For MHA setting [7], we use C = {ci,..., C. where Σn C, = h, determine how many heads (i.e., qs, k, v;) should be assigned to each chain. For the GQA setting [20], we require c; x hbv should be multiple of sum(C). For Chain-of-Attenin KV sharing mechanism, we require each e; should be multiple of fikes as keys and values are computed within the first chain and then shared across the remaining chains. And then, we will replicate keys and values per chain to guarantee that their number to match query. Figure 9 illustrates the differences between GQA and Attention (Chain) with the KV sharing mechanism. We also provide the Py Porch implementation of -haie-of-Attention in Algorithm 2 and its version with KV sharing in Algorithm 3.
with
number of
23

---

Algorithm 3 Pseudo Code for Attention (Chain) Layer with KV sharing.
import torch. nn as nn
from flash_attn import Flash_attn_fte from Linear_chain import Linear
class Attention (Mn..el)
def __init__ self, dim: int, n_ head: int,
int, chains: List [int]
): super()._.
assert sum(chains) n_head
for chain in chains:
assert chain n_kv_heat self.head thead_dim dim // n_head
self kv_mi
chains [0] self head_dim
self pq Linear (dim, dim, chains)
self wk nn. Linear (self kv_dim, n_kv_head head self head_dims,, chains)
self.w wv [Linear(self kv_dim, head self head _imm, chains)
self wo Linear (dim, dim, chains)
def forward(self,
x: torch. Tensor, freqs_cis: Tensor,
chain: int None) -> torch. Tensor:
bsz, seqlen, dim shape
x0 x[..., :self.kv_di
q, k, v self wx chain), (xlf), self.we ..f.wwvxx) q q. vbsz, seqlen, -1, self head_ dim)
k k. visu seqlen, -1, self half.head_dee head_dim)
v. vieiewbb seqlen, -1, self helf.headddim)
9, apply_rotary_emb(g k, freqs_cis)
n_ rep s.apap [2] // shape[2]
k k. repeat (1, 1, n_rep, 1)
v v.repea 1, n_rep,
O = flash_ Mash_atttn_uuncqq k, causal=True)
o o. (iew seqlen, -1)
return self wo chain=chain)
A.6.3 FFN
In Figure 10, illustrate the differences between the original FFN and our proposed FFN layer, which just needs to replace the Linear layer by using Chain-of-Lineau layer. They are equivalent when the number chains up the FFN layer, so that the output features of the Transformer block (i.e., Attention FFN) also follow the criteria of CoL. Definition 2. We also offer a Py Torch implementation of FFN (GeLU) in Algorithm 4 and FEN (SwiGLU) in Algorithm 5 for reference.
1. We use the same hyperparameter of the Autention module to set

---

Algorithm Pseudo Code for FFN implementation (GeLU).
import torch. mn. functional F
Import Chain-of Manin-of-Linear Layer
from Linear_chain import Linear
class FFN Monuule))
def __init__ self, dim: int, hidden_dimi int, chains: List [int]
Dimension of input features hidden, dim: Dimension of output features chains: Base of each chain
dim:
super () initt__
self. w1 Linear (dim, hidden_dim, chains=chains) self w2 Linear (hidden_ding dim, chains chains)
def forward(self, x: torch. Tensor, chain: int None) >> torch. Tensor: return self. w2 C.....gllusslff.eee w1 wx(x, chain)), chain)
Algorithm 5 Pseudo Code for FFN implementation (SwiGLU).
import torch. mn functionnl as F Import Chain-of-Lineds Layer from Linear_chain import Linear
class FFN (nn. Module):
def __init__
self, dim: int, hidden_diag dim: int, chains: List [int]
)
mnn
Dimension of input features hidden, dim: Dimension of output features
dim:
chains: Base of each chain
super () init__
self. w1
Linear (dim, hidden_dim, chains)
self x2 Linear (hidden_dde dim, dim, chains) Linear (dim, sidden_dim, chains)
self w3
def forward x: Tensor, chain: int = None) -> torch. Tensor: return self. (2 silu(self, wi wi(x, chain)) self. w3 (x,, chain), chain)
25

---

41
42
43
94
9s
96
47
41
42
43
44
95
96
97
qa
K2
33
kg
kg
k2
kg
k4
11
1/2
1/3
F4
V1
22
1/
(a) Group Query Attention.
(b) Attention (Chain) w/ KV sharing.
Figure 9: Differences between Group Query Attention (GQA) and Attention (Chain) with KV sharing in the attention implementation.
Chain
Re GAU
Reily GRIU
Uiear
Chain of Linar
(a) FFN (ReLU/ GELU).
Ch Linear
SEU
SEU
Tiwer
Lineer
Chan i
Chain of linear
(b) FFN (SwiGLU).
Figure 10: FFN for CoR. We replace all Linear by using Chain-of-LLineal layer.
A.6.4 Embedding
Figure 11 illustrates how the Embedding layer is utilized in our CoI .M architecture by activating differ- ent numbers of chains. To use different scales, we only need to activate the dimension corresponding to its all preceding chains.
26

---

Algorithm 6 Pseudo Code for Embeddings.
import torch.me nn as nn
class Embedding (nn. Embedding):
def __init__
self,
dim: int, vocab_size: int,
chains: List [int],
padding.idx: idx: Optional [int] None,
):
super() ivoc_ size, dim, padding_idx idx)
self dims [dim с // sum (chains))) for in np. cumsum [chainn)
self nuainn (chaiin)
def forward(self,, x: torch. Tensor, chain: int None): if head is None: return embedding (x, self weight, self padding-idx
else:
assert chain self www.uuuumiiaaan chains
dim self dims [chain]
return embedding (x, self self weight,,, :dim], self faiffpaddding.ixx.
Inputs
Embedding
How
are
you
use first chains
Figure 11: Example of Embedding in CoLM architecture. In this example, design the number of chains as 3, while the dimensions for each chain are set as {2, 2, 43. we expect to use the first chains, only need activate the first neurons for each word in the embedding layer.
27

---

Output
Output
Output_2
Scale_0
: 1Scale_1
Softmax
Scale_2
Feed loorward (Chain) SwiGLU
norm
norm()
norm
Nx
Attention (Chain) with KV Sharing
K
norm(-)
norg/))
porm(-)
Embeddings
Inputs
Figure 12: The architecture of our CoLM. The model use Chain-of Representation to build each layer, including embedding, self-attention, feed-forward network, normalization. Each scale contributes progressively refined outputs (e.g., OutputO, Outputl, Output2).
28

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the paper's contributions, scope, and assumptions, which are consistent with the presented results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes. See Section Limitation for details.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]

---

Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The model architecture and training procedure are fully specified in Section 4, enabling reproduction.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
30

---

Answer: [Yes]
Justification: Yes. The code and instructions for reproducing the main experimental results are publicly available on GitHub (link provided in the abstract).
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: ggtgg/uids////// wwweesuumissiooPaee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes. Training details, hyperparameters, and data splits are described in Section 4 and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: No. Due to the high computational cost, experiments were run with a single seed.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
31

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Yes. Compute resources and training time are reported in Section 4.1.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips: http://pp....pppppppccaaeaaeeeea
Answer: [Yes]
Justification: Yes. We have read the NeurlPS Code Ethics and confirm that this research complies with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is primarily methodological and foundational, without a direct or immediate path to real- cal-world deployment identifiable societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
32

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not release models or datasets that pose a high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All existing code, data, and models used in this work are properly cited the paper, and their licenses and terms of use respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
33

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The released code documented the GitHub repository, including usage instructions and license information.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
This work does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
34

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were only used for writing and editing purposes and did not affect the core methodology or results.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //terr//s...../eeee ttteurip..cc/oooffeeeeeeeeeeeeee /ttt////eenneee/225/////////// 20LL/LLM for what should or should be described.
35


Reviews:
[{"review_id": "KYauc9G4su", "reviewer": "Reviewer_fgKQ", "summary": "The paper proposes Chain-of-Model (CoM) learning: split every hidden vector into ordered “chains,” force each output chain i to depend only on the first i input chains (Chain-of-Layer, CoL) and stack such layers to obtain a network that contains several nested sub-models. The instantiated Chain-of-Language-Model (CoLM) keeps the first chain identical to a standard Transformer and adds extra chains that can be trained later (chain expansion) or dropped at inference time (elastic inference) without recomputing KV caches thanks to a KV-sharing trick (CoLM-Air) .", "questions": "1- Which exact normalization scheme do you use (RMSNorm vs. LayerNorm)? Please include the details of how did you adapt the layer norm in your method.\n\n2- Does CoLM (with and without KV sharing) maintain accuracy in long sequence scenarios?\n\n3- How does the method behave at ≥7 B parameters? Do the small accuracy gap of CoLM-Air persist?", "limitations": "The paper has no separate limitation section. Including a short “Limitations & Societal Impact” subsection covering the points such as Scalability, Long-context quality etc. would make the submission more transparent and align with NeurIPS best practices.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n1- The paper is well-written and it is easy to understand. The application is technically sound and potentially useful as it shows an architectural innovation with tangible speed.\n\n2- The method is innovative and tries to tackle an interesting problem.\n\n3- Demonstrates that a single checkpoint can serve multiple width targets with negligible quality loss.\n\nWeaknesses:\n\n1- Some low level details are missing. For example in section 3.2.3 the details of normalisation scheme has not explained.\n\n2- Performance does decline as more chains are added, and even the two-chain CoLM-Air lags the dense baseline by ~0.3 points. Although that margin is still modest (well under one point), it indicates that the KV-sharing trick may exact a small but measurable accuracy cost, which could widen at larger model scales or on harder benchmarks. Demonstrating results at ≥7 B parameters (or providing significance tests to show whether 0.3 is within noise) would help establish how serious the penalty really is.\n\n3- Based on the speed-up curves in Figure 5, latency gains emerge only when the pre-fill prompt exceeds roughly 4 K tokens; below that, CoLM is nearly as slow as the dense baseline. This implies the architecture is chiefly advantageous in long-context settings. Yet all reported generalization scores rely on short-context LM-Harness benchmarks or GLUE tasks. To demonstrate practical value, the authors should also measure long-context accuracy (e.g., on Needle-In-a-Haystack or LongBench) to verify that KV-sharing retains quality.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments. Below are our response to address your concerns:\n\n1. **Missing the details of normalization. (W1 & Q1)**:\n\n- Thank you for pointing out this missing details. Generally, for each chain-of-representation feature, we apply normalization at the chain-wise level to maintain its Chain-of-Layer property. That means, each chain will have its individual normalization layer, which could be Layer Norm, RMS Norm or others. In our experiment, we use RMSNorm to aligh the Llama setting, which is also widely used in existing LLMs. We apologize for the missing information and will fix it in the final version.\n\n2. **Demonstration Results over 7B parameters or significant tests on CoLM-Air. (W2 & Q3)**:\n\n- Thanks for your constructive comments. We also agree that the experiments over 7B will be better to demonstrate the potential of CoLM-Air. Besides, since the first chain of CoLM is still the standard language model, it naturally maintains the scalability of existing language models. To validate this point, we conduct experiments using CoLM-Air (chain=[8, 8, 8, 8]) configurated with 4096 dims and 32 layers as LLama-7B configuration, resulting in a model with 4.1B parameters (Chain-of-Layer requires fewer parameters that dense layer). Besides, to make a fair comparison, we also set a dense model (4096 Dims and 16 Layers), with nearly 4.1B parameters. Due to resource limitations (32x 40GB A100 GPUs), the batch size is set as 2 per GPUs for 4.1B model and 1 per GPUs (accum_step=2) for 7B models, seq length is 2048. We apply FSDP and the total training step is 15K steps during the rebuttal period, and we report PPL over SlimPajama datasets per 5K steps to verify its convergence. The results are reported in below:\n\n    | Model |Bsz| Memory | Dim | Layer | Params | 5K | 10K | 15K |  \n    |-|-|-|-|-|-|-|-|-|\n    | CoLM (8-8-8-8) | 2|26GB | 4096 | 32 | 4.1B | 22.56 | 17.01  | 15.89 |\n    | CoLM-Air (8-8-8-8) |2| 26GB | 4096 | 32 | 4.1B | 22.74 | 17.45 | 16.48 |\n    | Dense |2| 39GB | 4096 | 20 | 4.1B | 23.29 | 17.12 | 15.80 | \n    | Dense |1| 32GB | 4096 | 32 | 7.3B | 22.26 | 16.39 | Ongoing |\n\n    Due to resource limitations, the total pre-training is still under-fit, and we will continue to seek more GPU resources in the future to fully pre-train a larger size LLM over 7B for usage. Moreover, as shown in the above table, the CoLM configurations require less GPU memory. Therefore, within the same computation, our COLM allows more training steps. But to make a fair comparison, we still compare different settings under the same steps. We also provide codes in the supplemental materials to reproduce our experiments. \n    \n- For significant test, we also calculate p-value of CoLM-Air, where p $\\approx$ **0.02** < 0.05. These results also indicate the effectiveness of our method. Since our current method is only pre-trained with 200B tokens, we will continue to pre-train our model to see its robustness when using more computational resources. Besides, just as mentioned in Line 200 ~ 204, KV sharing mechanism used in CoLM-Air can be considered as a trade-off technique built based on the standard CoLM architecture, to balance performance and additional flexibility / efficiency. For the standard CoLM architecture, our method can achieve comparable performance to the baseline model. Since CoLM takes fewer GPU memory / flops, and KV sharing mechanism can significantly reduce the prefilling speed (27x) with seamlessly LLM switching, we can increase the training budget of CoLM-Air to further improve its capability. Overall, in the future, we will continue to seek more GPU resources to explore its scalability.\n\n3. **Prefilling below 4K and Long-context accuracy (W3 & Q2)**:\n\n- For the prefilling, we report the results of CoLM-Air under a batch size of 1 for evaluation (Please see Appendix A.2.1). Therefore, our method cannot fully utilize GPU under 4K with a batch size of 1. However, for sequence length below 4K, our CoLM-Air allows larger batch size when compared with the baseline, as we only need to activate the first chain. Here, our CoLM-Air (8-8-8-8 setting) enables 4x batch size when compared with baseline under 4K settings. We will add these details about prefilling into our final version to make a clear understanding.\n- Besides, in the original experiments, we mainly generate KV caches by using sliding window mask. To extend its long-context ability, we also need to pre-train CoLM-Air on long-context data. Here, we freeze the first chain and then fine-tune its latter chains (i.e., Chain Tuning) on long context, so that it can ensure all keys and values are identical to the original model and thus perserve the original results under short context (only switching to tuned chains when requiring long-context scenarios). We use LongFormer trick [1] to continue pre-train our model on all chains except the first one. And the results on LongBench is shown as:\n\n    | Text Length | Single-Doc QA | Multi-Doc QA| Summarization | Few-shot Learning | Synthetic Task | Code Completion |  \n    |-|-|-|-|-|-|-|\n    | 8k | 0.50 | 0.51| 0.52| 0.55 | 0.55 | 0.66 |\n    | 16k | 0.51 | 0.54 | 0.53 | 0.58 | 0.51 | 0.66 | \n    | 32k | 0.53 | 0.58 | 0.56 | 0.61 | 0.53 | 0.67|\n    | 64k | 0.55 | 0.6 | 0.56 | 0.61 | 0.54 | 0.66 |\n\n- These results indicate our method can effectively extend its long-context capability and maintain short-context preformance. In the future, we will investigate better long-context training techniques on CoLM-Air to unlock its potential on long-context tasks. \n    > [1] Longformer: The Long-Document Transformer.\n\n4.  **Limitations Section**\n\n- We apologize for the misplacement of the Limitation section, which is currently in Appendix A.5. We will move it to the main paper in the final version."}, {"review_id": "RhowMrubuC", "reviewer": "Reviewer_LJrW", "summary": "In summary, the authors propose a novel learning paradigm called \"Chain-of-Model\" (CoM). Within this paradigm, the hidden representations are divided into multiple sub-representations, each of which represents a different scale. With such division of hidden representations, the authors propose the concept of \"Chain-of-Layer\" (CoL), where the scales are learned causally within a layer. That is, the scale $i$ of the layer output only depends on the scales $\\leq i$ of the layer input. By stacking layers that satisfy the CoL property, the entire model becomes CoM. Based on this concept, the authors redesign the common layers for language models to satisfy the CoL property and propose Chain-of-Language-Model (CoLM). The proposed CoLM enjoys several properties: 1) The existing LLMs can be treated as chain=1 and easily extended by adding another chain without training from scratch. 2) A model trained with CoLM can be deployed with dynamic inference capacity. 3) CoLM models with KV sharing achieve better pre-filling latency. These properties are evaluated on 1B scale model pertaining with several common-sense reasoning tasks.", "questions": "1. The proposed CoLM is mainly tested on Llama variants. Is the proposed CoLM applicable to other model series such as Qwen and DeepSeek MLA? \n\n2. For the chain expansion experiment, each model is only expanded with 1 chain. It would be good to see the performance of extending more chains (e.g., $c_2=4, c_3=4$) and discuss the pros/cons.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. This paper studys an important problem as scaling LLM models is extremely resource-consuming. The proposed CoLM could extend an existing model with more chains (more capacity) without training from scratch, which provides a valuable solution for the community\n2. Though there are several weaknesses, the proposed method is well-presented in the paper and looks quite novel to me\n\nWeakness:\n1. The proposed CoLM is mainly validated on models with 1B parameters, which seems not enough to guarantee that it will remain the same performance for large scales such as 3B, 8B or even larger. Though it is understandable that pre-training 3B/8B model is resource-restricted, it would be nice to see experiments where the proposed CoLM is used to extend some larger models such as Llama-3.2-3B and Llama-3.1-8B. \n2. In Table 2, although the CoLM expansion introduces additional 0.8B parameters, the improvement of LLaMA-3.2-1B looks trival to me given that it is trained with 8B tokens. \n3. Since the proposed CoLM makes changes to the layer structure. It is unclear whether the proposed method is compatible to the existing large-scale training techniques (e.g., Different parallelism).\n4. As the author mentioned, each scale should have its own loss function and that potentially limits the training efficiency when the number of chains is huge \n5. With the KV sharing mechanism, the key/value heads only contain the information of the first scale. When the model is expanded with multiple rounds, the model capacity will be limited by the key/value heads.", "was_revised": true, "final_justification": "In general, this paper proposes a novel learning scheme to scale up pre-trained models. During the rebuttal, most of my concerns were resolved. The remaining concern is mainly about the scalability of the proposed method to larger models. So I will maintain my recommendation as borderline accept.", "rebuttal": "Thank you for your valuable comments. Below are our response to address your concerns:\n\n1. **It is better to see experiments on 3B, 8B or larger (W1)**:\n\n- Thanks for your constructive comments. We also agree that the experiments on 3B, 8B will be better to demonstrate the potential of CoLM-Air. Besides, since the first chain of CoLM is still the standard language model, it naturally maintains the scalability of existing language models. To validate this point, we conduct experiments using CoLM (chain=[8, 8, 8, 8]) configurated with 4096 dims and 32 layers as LLama-7B configuration, resulting in a model with 4.1B parameters (Chain-of-Layer requires fewer parameters that dense layer). Besides, to make a fair comparison, we also set a dense model (4096 Dims and 16 Layers), with nearly 4.1B parameters. Due to resource limitations (32x 40GB A100 GPUs), the batch size is set as 2 per GPUs for 4.1B model and 1 per GPUs (accum_step=2) for 7B models, seq length is 2048. We apply FSDP and the total training step is 15K steps during the rebuttal period, and we report PPL over SlimPajama datasets per 5K steps to verify its convergence. The results are reported in below:\n\n    | Model |Bsz| Memory | Dim | Layer | Params | 5K | 10K | 15K |  \n    |-|-|-|-|-|-|-|-|-|\n    | CoLM (8-8-8-8) | 2|26GB | 4096 | 32 | 4.1B | 22.56 | 17.01  | 15.89 |\n    | CoLM-Air (8-8-8-8) |2| 26GB | 4096 | 32 | 4.1B | 22.74 | 17.45 | 16.48 |\n    | Dense |2| 39GB | 4096 | 20 | 4.1B | 23.29 | 17.12 | 15.80 | \n    | Dense |1| 32GB | 4096 | 32 | 7.3B | 22.26 | 16.39 | Ongoing |\n\n  Due to resource limitations, the total pre-training is still under-fit, and we will continue to seek more GPU resources in the future to fully pre-train a larger size LLM over 7B for usage. Moreover, as shown in the above table, the CoLM configurations require less GPU memory. Therefore, within the same computation, our COLM allows more training steps. We also provide codes in the supplemental materials to reproduce our experiments. We will also put these experiments into our paper to highlight the scalability of our model.\n    \n2. **The improvement of LLaMA-3.2-1B looks trival (W2)**:\n\n- Thanks for your comment and for closely examining our experimental results. The main reason for the marginal improvement over LLama-3.2-1B is because we only use SlimPajama (670B tokens) as the pre-trained dataset and only pre-train 20K steps for pre-training with total 8B tokens. In the original LLama-3.2-1B, it uses high-quality dataset with over 15T tokens plus massive pre-training. Therefore, it will be a little difficult to signficantly improve LLama performance with total 8B tokens training, but it still achieves nearly 0.15 points improvement. To this end, we also offer results on TinyLLama-v1.1, which is built upon SlimPajama plus Starcoder dataset, to make a fair comparison. In this case, we observe that our model can improve the baseline of TinyLlama-v1.1 with nearly 0.92 points within 8B tokens training. This results also indicate our method can effectively increase model scales in a progressive way. In the future, we will also continue to increase our pre-training budget to further increase the capability of our model.\n\n3. **Whether the proposed method is compatible to the existing large-scale training techniques. (e.g., Different parallelism) (W3)**:\n\n- Thanks for your comments. Generally, our method is fully compatible with existing parallelisms, including data, pipeline, context and expert parallelism (i.e., DP, PP, CP and EP). In our experiments, we also use Fully Sharded Data Parallel (FSDP) to pre-train our model. For Tensor Parallelism (TP), our method actually is also compatible with it by using TP over each chain based on Algorithm 1. It will not introduce any extra computation but increase some communication cost (e.g., more all-reduce to access data). Therefore, we use triton to implement a block-wise sparse kernel to improve our training efficiency in Appendix A.5. We will also explore more efficient parallelism methods (e.g., chain parallelism) to scale up our model in the future.\n\n4. **Each scale should have its own loss function and that potentially limits the training efficiency when the number of chains is huge. (W4)**:\n\n- Thank you for this astute observation. We admit this point and thereby adopt the MRL [1] trick to increase the training efficiency for multiple chains (Line 213-215). MRL trick is to freeze the backbone network (pre-trained) and then only fine-tune the classification layer to provide multi-scale prediction. Motivated by MRL [1], we also first pre-train the whole network, and then freeze the backbone and only fine-tune classification layer with multiple chains to increase training efficiency. We will provide more details about this part to make a clear understanding in the final version.\n    \n    > [1] matryoshka representation learning. NeurIPS 2022\n\n5. **When the model is expanded with multiple rounds, the model capacity will be limited by the key/value heads. (W5)**:\n\n- Thank you for this insightful question. We admit your concern, just as our discussion in the Limitation part (Linear 650-661), that the first chain plays a critical role in CoLM-Air setting and how to determine the optimal setting for the first chain is important. Specifically, KV sharing can be considered as a trade-off technique (Line 200-204) built based on the standard CoLM architecture, to balance performance and extra flexibility / efficiency. For the standard CoLM without KV sharing, our method can achieve comparable performance to the baseline. Specifically, as CoLM takes fewer GPU memory / flops and KV sharing mechanism can significantly reduce the prefilling speed (27x) with seamlessly LLM switching, we can increase the training budget of CoLM-Air to further improve its capability. In the future, we will also continue to explore the optimal setting of the first chain in our CoLM-Air.\n\n6. **Q1: Is the proposed CoLM applicable to other model series such as Qwen and DeepSeek MLA?**:\n\n- Yes, the idea of CoM framework can be applied to any model series including Qwen and DeepSeek MLA. Based on the Corollary 2.1, any model can be considered as a kind of CoM when n = 1. Therefore, CoM is a general scaling paradigm, which just requires each layer to meet the criteria of chain-of-layer. For Qwen or DeepSeek series model, we just need to replace all Linear layer, attention and normalization by using our setting in Section 3.2. When the number of chain is 1, it is identical to the original LLM settings like Qwen or DeepSeek.\n\n\n7. **Q2: For the chain expansion experiment, each model is only expanded with 1 chain. It would be good to see the performance of extending more chains and discuss the pros/cons.**:\n\n- Thanks for your constructive comment. Due to resource limitations, in our original experiments (Sec 4.2), we only extend one chain over 1B model to verify the effectiveness of our method in chain expansion, as it introduces more parameters (e.g., 0.8B parameters) for computations. We agree that adding more chains would better highlight its pros and cons. Therefore, we provide multi-chain experiments with 4 chains (C = [8, 8, 8, 8]) on our pre-training settings (Section 4.1). We also provide the detailed training curve about different chains in Figure 6(b) in Appendx, to offer the visualization of using more chains. These results also indicate that using more chains allows model to reduce its training loss. Generally, the pros / cons can be summarized as below:\n\n    - **Pros**: When we use multiple chains, it enables us to conduct elastic inference and efficient progressive scaling to offer different model sizes to meet the requirements from diverse devices. Moreover, by introducing KV sharing mechanism, it can significantly reduce the prefilling speed. Besides, each new chain would build upon the cumulative knowledge of all preceding chains (including the original model), leading to a more powerful model.\n\n    - **Cons**: Just as mentioned in Line 651 - 661, each added chain will reuse knowledge from previous chains. Therefore, if the former chains are not well-trained, they will also affect the latter chains to make prediction. Therefore, how to ensure high-quality train over each chain is still very imporant.\n\n    Overall, multi-chain expansion can be considered as a flexible strategy to extend model and offer elastic inference. We will also add these discussions into our final version to make a clear understanding."}, {"review_id": "I5LLXy7bp9", "reviewer": "Reviewer_xjCS", "summary": "This paper proposes a novel learning paradigm, termed \"Chain-of-Model\" (CoM), which aims to address the core challenges in the model scaling and deployment of current Large Language Models (LLMs). Its core idea originates from \"Chain-of-Representation\" (CoR), which treats the representation vectors in the model's hidden layers as a concatenation of multiple \"sub-chains\". The authors  apply this paradigm to the Transformer architecture, designing the \"Chain-of-Language-Model\" (CoLM). Experimental results show that this method performs slightly worse than the vanilla Transformer in terms of effectiveness, but has advantages in incremental learning, elastic inference, and efficient deployment.", "questions": "See Weaknesses", "limitations": "None", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- The \"Chain-of-Representation\" idea proposed in this paper is highly innovative. The resulting model is inherently capable of multi-scale adaptation, which can solve many problems in incremental learning, elastic inference, and efficient deployment.\n\n- The paper provides a thorough discussion on various aspects of the Chain-of-Representation concept, including model training  inference, etc.\n\n- The experimental results demonstrate that the proposed method achieves good performance.\n\n**Weaknesses:**\n\n- Potential Impact of KV Sharing on Generation Quality: The paper mentions that the KV sharing mechanism \"slightly affects the performance\" , a claim verified by the accuracy on commonsense reasoning tasks. However, the specific impact of this on open-ended text generation tasks (e.g., in terms of coherence, logic, and creativity) remains unclear. Since higher-order chains rely on the semantic information (K and V) generated by lower-order chains, this could lead to an information bottleneck. Providing some qualitative analysis or results on generation-focused benchmarks would allow for a better assessment of the trade-offs of this design.\n    \n- Compatibility Issues with Tensor Parallelism (TP): The authors acknowledge in the appendix that the naive implementation of the Chain-of-Linear layer is not well-suited for Tensor Parallelism (TP), which is a very practical obstacle for training very large-scale models. Although the authors mention optimization via a custom-developed kernel, the paper does not provide a clear answer as to whether the additional communication overhead brought by TP can be fully resolved. This could be a significant barrier to the widespread adoption of this architecture in the industry.", "was_revised": true, "final_justification": "I think this paper should be accepted to Neurips 2025. Most of my concerns regarding the comprehensiveness of experiments and the practical application of this method have been addressed through the authors' detail responses.", "rebuttal": "Thank you for your valuable comments. Below are our response to address your concerns:\n\n1. **Potential Impact of KV Sharing on Generation Quality. (W1)**:\n\n- We sincerely thank you for providing these constructive comments for us. We also admit that adding generation results can better demonstrate the effectiveness of our method. To resolve your concerns, we also provide our results on generation tasks. Specifically, we took the TinyLlama-v1.1 model as our baseline and applied our expanded model (used in Section 4.2), which use KV sharing to reuse KV cache from the TinyLlama. Both baseline and our model are fine-tuned on Alpaca datasets for post-training. We report the generation results on MT-Bench[1], a standard benchmark for assessing conversational and instruction-following abilities. The results are shown below:\n\n    | Model            | Writing | Roleplay  | Reasoning  | Math  | Coding  | Extraction |STEM | Humanities  | Avg.  |\n    |------------------|-------|-------|-------|-------|-------|-------|-------|-------|--------|\n    | Tinyllama         | 2.85 | 3.20 | 2.05 | 1.10 | 1.30 | 1.40 | 2.35 | 3.20 | 2.18  |\n    | Ours             | 2.75 | 3.40 | 2.70 | 1.00 | 1.10 | 1.30 | 2.35 | 3.65 | **2.28** |\n    \n    These results also indicate that by using KV sharing, our model can significant enhance performance over the original model. We will add these generation results into our final version to make a clear understanding.\n\n- Besides, just as mentioned in Line 200 ~ 204, KV sharing mechanism used in KV sharing can be considered as a trade-off technique built based on the standard CoLM architecture, to balance performance and additional flexibility / efficiency. For the standard CoLM architecture, our method can achieve comparable performance to the baseline model. Since KV sharing mechanism can significantly reduce the prefilling speed (nearly 27x) and enable seamlessly LLM switching across different scales, we can increase the training budget of CoLM-Air to further improve its capability. In the future, we will continue to search the optimal setting for the first chain to design our CoLM-Air. Overall, we will highlight this detail in our final version to make a clear understanding.\n\n2. **Compatibility Issues with Tensor Parallelism. (W2)**: \n\n- Thanks for your comments. Generally, our method is fully compatible with existing parallelisms, including Data Parallelism (DP), Pipeline Parallelism (PP), Context Parallelism (CP) and so on. In our experiments, we also fully utilized Sharded Data Parallel (FSDP) to pre-train our model. For Tensor Parallelism (TP), our method can also use it to each chain actually by using Algorithm 1 in Appendix A.6.1. It will not introduce any additional computation cost, but increases some communication cost in data access during the distributed training. Therefore, just as discussed in Appendix A.5, we propose a block-wise sparse kernel based on triton to mitigate this problem. In the future, we will also continue to devise novel parallelism methods (e.g., chain parallelism) to enable more efficient scalabilty for our method."}, {"review_id": "3FKB1IABEb", "reviewer": "Reviewer_Uwqq", "summary": "This paper proposes **Chain-of-Model (CoM)**, a progressive learning paradigm inspired by Progressive Neural Networks, where newly added neurons are causally dependent on the preceding ones. CoM is applied to progressively expand the **width** of Transformer models. The authors provide detailed implementations for each module of the Transformer architecture (e.g., FFN, multi-head attention, normalization, and embedding). Experimental results demonstrate the effectiveness of CoM in model expansion, elastic inference, and continual learning.", "questions": "See Weaknesses", "limitations": "See Weaknesses", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- The proposed method is intuitive and well-motivated for scaling the width of Transformers. The design is supported by experimental evidence.\n- The paper provides a clear and detailed description of the method, including pseudocode and implementation details.\n- The experiments cover a variety of scenarios, such as model expansion, tuning, and inference.\n- The paper is well-written and easy to follow.\n\n\n**Weaknesses**\n\nWhile the method is novel and well-implemented, its **superiority over existing approaches remains unclear**. Although the experiments cover multiple settings, **comparisons with strong baselines are missing**, making it difficult to assess the practical advantages of CoM.\n\n1. **Comparison in 'Chain Expansion' (Section 4.2)**:\n    The authors should compare CoM with other model expansion techniques, such as vanilla width expansion.\n\n   > For example, when expanding the representation space from NNN to N+MN+MN+M, CoM adds a matrix of shape (N+M,M)(N+M, M)(N+M,M). How would this compare to simply adding a matrix of shape (M,M)(M, M)(M,M) without causal constraints?\n   >  This ablation could help validate the core contribution of CoM—namely, the necessity of causal dependencies across chains.\n\n2. **Comparison in 'Chain Tuning' (Section 4.5)**:\n    The paper should include comparisons with parameter-efficient fine-tuning methods such as **LoRA**, to show how CoM performs in terms of efficiency and performance in limited-resource settings.\n\n3. **Lack of comparison with MoE**:\n    The authors claim that CoM is compatible with Mixture-of-Experts (MoE) models. However, a natural question arises:\n\n   > Is it more effective to use CoM to increase the capacity of each expert, or to increase the number of experts directly?\n   >  Both approaches aim to scale model capacity—some comparison or analysis would be valuable.\n\n4. **Concerns about evaluation tasks**:\n    All evaluated benchmarks appear to be **multiple-choice tasks**. It would be beneficial to include results on **generative tasks**, which may better reflect the model's performance in real-world use cases.\n\n\n\n\n**Overall**\n\nThis work addresses an important problem: how to progressively scale Transformer models without retraining from scratch. The method is well-structured and the engineering implementation is solid. However, to fully validate its contribution, the paper should include **comparisons with relevant baseline methods** and consider **broader evaluation settings**, including generative tasks.", "was_revised": true, "final_justification": "The authors' response addresses my concerns. I recommend accepting this paper.", "rebuttal": "Thank you for your valuable comments. Below are our response to address your concerns:\n\n1. **Comparison CoM with other model expansion techniques, such as vanilla width expansion. (W1)**:\n\n- Thank you for this insightful question. Generally, vanilla width expansion (i.e., CoM without constraints) can be considered as creating an individual expert (model) and combining it into the final prediction layer. Thus, this expert cannot leverage the capability from the previous model, which could limits its upper bound. To resolve your concerns, we also offer some experiments to compare with vanilla width expansion (i.e., CoLM without constraints). The results are shown in below:\n\n    | Model                   | HellaSwag | Obqa | WinoGranda | ARC-e| ARC-c | Boolq | Piqa | Sciq | Avg   |\n    |------------------------|------------|-------|-------------|--------|--------|--------|-------|-------|--------|\n    | Baseline               | 61.47      | 36.62 | 59.43       | 55.47  | 32.68  | 55.99  | 73.56 | 84.20 | 57.43  |\n    | + Vanilla Width Expansion | 61.20      | 35.41 | 60.22       | 55.33  | 32.08  | 57.61  | 73.72 | 85.00 | 57.57  |\n    | + Expansion (Ours)            | 61.66      | 36.62 | 60.62       | 56.27  | 32.94  | 58.44  | 74.05 | 86.20 | **58.35** |\n\n- Besides, to some degree, vanilla width expansion (i.e., without constraints) can also be considered as a specical case of Chain-of-Model, that $y_i$ is only conditioned on $x_i$ which belong to a subset of Definition 2. In our CoLM, our normalization also remove this causal constraint to slightly improve training efficiency. Therefore, we deem that for some specific module, we can use vanilla width expansion (CoM without constrains) to add some efficiency. In the future, we will also continue to seek more GPU resource to explore the optimal setting for our CoLM architecture.\n\n2. **Comparisons with parameter-efficient fine-tuning methods such as LoRA, to show how CoM performs in terms of efficiency and performance in limited-resource settings. (W2)**:\n\n- Thank you for this excellent suggestion to clarify the relationship between Chain Tuning and established PEFT methods as LoRA.\n    Chain Tuning is not merely another PEFT method, but rather a novel fine-tuning paradigm that leverages the inherent causal structure of our Chain-of-Model (CoM) architecture. As described in lines 290-292, our Chain Tuning strategy involves freezing the initial chain(s) and only fine-tuning the subsequent ones. Therefore, its perfectly preserve the previous module to generate keys and values. Consequently, in KV sharing, the KV cache generated by the original pre-trained model can be perfectly and seamlessly reused by any model that has been fine-tuned on the latter chains. To the best of our knowledge, this is the first work that allows us to seamlessly switch different LLMs without re-computing KV caches. In contrast, the motivation of LoRA is designed to reduce the heavy fine-tuning cost over the huge model, which is also compatible with our method, as we can apply LoRA to our latter chains only. The below Table also summarizes the comparisons between our Chain Tuning and LoRA.\n    | Model            | SST-2 | COLA  | MNLI  | MRPC  | QNLI  | QQP   | RTE   | WNLI  | Avg.  |\n    |------------------|-------|-------|-------|-------|-------|-------|-------|-------|--------|\n    | Baseline         | 67.89 | 32.98 | 35.87 | 49.02 | 51.77 | 51.65 | 54.87 | 43.66 | 48.46  |\n    | LoRA             | 91.74 | 58.46 | 81.32 | 61.27 | 73.20 | 54.92 | 55.96 | 52.11 | 66.12  |\n    | + CT             | 92.32 | 54.87 | 82.26 | 62.01 | 82.87 | 55.22 | 59.21 | 53.52 | **67.79** |\n\n    In the future, we will also add these discussion into our final version to make a clear understanding and highlight our biggest advantages of chain tuning.\n\n3. **Comparison with MoE. (W3)**:\n\n- Thanks for your question. We want to highlight that CoM and MoE are two strategies for model scaling from different dimensions, just as discussed in Question 2 of A.4 in the Appendix. CoM focuses on progressive scaling and offers multi-scale abilities. MoE is to increase model capability via building giant model with sparse activation among one scale. That means, MoE is good at enhancing model capability within single scale, while CoM is to add new scale over existing scales. For example, if we want to extend an additional chain (scale) over existing models, we can still use MoE within this chain (scale), but bridge the connections among multiple scales via CoM. We will add more discussions about this part to further highlight the differences between CoM and MoE.\n    \n4. **Concerns about generation tasks. (W4)**:\n\n- We are grateful to the reviewer for this insightful suggestion. To address this point, we have conducted new experiments focusing on generative instruction-following. Here, we took the TinyLlama-v1.1 model as our baseline and use our chain expansion model for comparison. Both the baseline and our expanded model were then instruction fine-tuned on the Alpaca dataset.\n    We evaluated the generative performance of both models on MT-Bench [1], a standard benchmark for assessing conversational and instruction-following abilities. The results are reported in below:\n    | Model            | Writing | Roleplay  | Reasoning  | Math  | Coding  | Extraction |STEM | Humanities  | Avg.  |\n    |------------------|-------|-------|-------|-------|-------|-------|-------|-------|--------|\n    | Tinyllama         | 2.85 | 3.20 | 2.05 | 1.10 | 1.30 | 1.40 | 2.35 | 3.20 | 2.18  |\n    | Ours             | 2.75 | 3.40 | 2.70 | 1.00 | 1.10 | 1.30 | 2.35 | 3.65 | **2.28** |\n    \n    These results also indicate that our method can also achieve better performance on generation tasks than baseline, which validates the generalization of our proposed method. We will add these results into our final version to further highlight the effectivenss of our method on generation tasks.   \n    > [1] Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena, NeurIPS 2025"}]

Author rebuttal:

Meta-review: All reviewers agreed this paper should be accepted: it addresses an interesting problem, it is well-written, and the proposed idea creatively addresses the problem. The main concerns were around experimental comparisons: e.g., reviewer Uwqq wanted to see comparisons against other model expansion techniques and parameter-efficient fine-tuning methods. The authors provided these comparisons in the rebuttal. This satisfied the reviewer and they increased their score. In fact, the authors responded to all reviewers extensively, often running detailed experimental comparisons. This addressed additional reviewer concerns: the paper is a clear accept. Authors please remember to include all reviewer requests into the camera-ready paper. Thank you!
Decision: Accept (poster)