Chain-of-Model Learning for Language Model

Thank you for your valuable comments. Below are our response to address your concerns:

1. Comparison CoM with other model expansion techniques, such as vanilla width expansion. (W1):

• 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:

  Model	HellaSwag	Obqa	WinoGranda	ARC-e	ARC-c	Boolq	Piqa	Sciq	Avg
  Baseline	61.47	36.62	59.43	55.47	32.68	55.99	73.56	84.20	57.43
  + Vanilla Width Expansion	61.20	35.41	60.22	55.33	32.08	57.61	73.72	85.00	57.57
  + Expansion (Ours)	61.66	36.62	60.62	56.27	32.94	58.44	74.05	86.20	58.35

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

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

• Thank you for this excellent suggestion to clarify the relationship between Chain Tuning and established PEFT methods as LoRA. 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.

  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

  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.

3. Comparison with MoE. (W3):

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

4. Concerns about generation tasks. (W4):

• 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. 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:

  Model	Writing	Roleplay	Reasoning	Math	Coding	Extraction	STEM	Humanities	Avg.
  Tinyllama	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

  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.

  [1] Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena, NeurIPS 2025

Thank you for your valuable comments. Below are our response to address your concerns:

1. Potential Impact of KV Sharing on Generation Quality. (W1):

• 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:

  Model	Writing	Roleplay	Reasoning	Math	Coding	Extraction	STEM	Humanities	Avg.
  Tinyllama	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

  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.

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

2. Compatibility Issues with Tensor Parallelism. (W2):

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

Dear Reviewer xjCS,

The authors have responded to your concerns. How does their response change your view of the paper? If it does not, please clarify what the authors can do to address your concerns. If it does, please consider adjusting your score based on their response.

Your AC

Thank you for your valuable comments. Below are our response to address your concerns:

1. It is better to see experiments on 3B, 8B or larger (W1):

• 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:

  Model	Bsz	Memory	Dim	Layer	Params	5K	10K	15K
  CoLM (8-8-8-8)	2	26GB	4096	32	4.1B	22.56	17.01	15.89
  CoLM-Air (8-8-8-8)	2	26GB	4096	32	4.1B	22.74	17.45	16.48
  Dense	2	39GB	4096	20	4.1B	23.29	17.12	15.80
  Dense	1	32GB	4096	32	7.3B	22.26	16.39	Ongoing

  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.

2. The improvement of LLaMA-3.2-1B looks trival (W2):

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

3. Whether the proposed method is compatible to the existing large-scale training techniques. (e.g., Different parallelism) (W3):

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

4. Each scale should have its own loss function and that potentially limits the training efficiency when the number of chains is huge. (W4):

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

  [1] matryoshka representation learning. NeurIPS 2022

5. When the model is expanded with multiple rounds, the model capacity will be limited by the key/value heads. (W5):

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

6. Q1: Is the proposed CoLM applicable to other model series such as Qwen and DeepSeek MLA?:

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

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

• 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:

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

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

  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.

Thank you for your valuable comments. Below are our response to address your concerns:

1. Missing the details of normalization. (W1 & Q1):

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

2. Demonstration Results over 7B parameters or significant tests on CoLM-Air. (W2 & Q3):

• 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:

  Model	Bsz	Memory	Dim	Layer	Params	5K	10K	15K
  CoLM (8-8-8-8)	2	26GB	4096	32	4.1B	22.56	17.01	15.89
  CoLM-Air (8-8-8-8)	2	26GB	4096	32	4.1B	22.74	17.45	16.48
  Dense	2	39GB	4096	20	4.1B	23.29	17.12	15.80
  Dense	1	32GB	4096	32	7.3B	22.26	16.39	Ongoing

  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.

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

3. Prefilling below 4K and Long-context accuracy (W3 & Q2):

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

• 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:

  Text Length	Single-Doc QA	Multi-Doc QA	Summarization	Few-shot Learning	Synthetic Task	Code Completion
  8k	0.50	0.51	0.52	0.55	0.55	0.66
  16k	0.51	0.54	0.53	0.58	0.51	0.66
  32k	0.53	0.58	0.56	0.61	0.53	0.67
  64k	0.55	0.6	0.56	0.61	0.54	0.66

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

  [1] Longformer: The Long-Document Transformer.

4. Limitations Section

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

Dear Reviewer fgKQ,

We hope our latest response can address your concerns. Besides, to further understand the mechanism of our method in improving prefilling speed, we additional provide the computational complexity of CoLM-Air in sequence calculation.

For a batch (size = B) of sequences with length of L, the complexity of Transformer (Dimension=$D_{x}$) to process this batch is as $(BL^2{D_{x}} + BL{D_{x}}^2)$, where $BL^2{D_{x}}$ represents the complexity of attention and $BL{D_{x}}^2$ represents the complexity of FFN. So, for CoLM-Air, since it only requires the calculation of the first chain to obtain keys and values. Therefore, the Transformer complexity of CoLM-Air for prefilling is as $(BL^2{D_{x_1}} + BL{D_{x_1}}^2)$, where $D_{x_1}$ denotes the dimension of the first chain, just as shown in below Table:

As $D_{x_1} << D$, for the prefilling tasks, CoLM-Air enables longer sequence processing within the same batch, or enables larger batch size with processing same sequence length. In other words, CoLM-Air offers a different solution to optimize complexity from the perspectives of Dimension (i.e., $D$), rather than $L$. Therefore, the current architecture of CoLM-Air is not specific designed for long-context setting, and our original experiments mainly use sliding window mask tricks and highlight its advantages in prefilling speed. But just as discussed in Line 274, the first chain is still the standard language model. Hence, you can apply any long-context tricks (e.g., YaRN, Sparse Attention) to extend its capability. Therefore, we conduct some experiments to verify its compatibility with long-context settings, when compared with vanilla model (as it can maintain short-context capability and then improve its long-context capability). As CoLM-Air architecture is not designed for long-context training (i.e., optimize $L$), it still requires us to explore the optimal setting when combining CoLM-Air with existing long-context techniques to achieve STOA results. As CoLM is a new scaling architecture, we will consider this task as the future work to explore its potential and we will add more discussion into our final version to highlight the differences of our method between prefilling and long-context capability.

Thank you once again for your insightful feedbacks.

Dear Reviewer fgKQ

Thank you for your valuable feedback. We are pleased that our previous responses have addressed your first two concerns, and we apologize for the misunderstanding regarding the long-context setting in CoLM-Air. Below are our response to address your final concern:

1. Generally, the design of current CoLM-Air setting can mainly improve the prefilling speed as it only needs to activate the first chain to calculate keys and values. However, its current architecture is not specific designed for long-context training, which still requires other training tricks (e.g., sparse attention, streaming llm) to enable long-context. In the original paper, our first chain still use the standard language model settings (Line 274), and we report the prefilling speed from 4K to 1M by using sliding window mask tricks (in our response) to demonstrate its efficiency in prefilling speed. Therefore, in the original experiment setup, we mainly validate the efficiency of CoLM-Air in long-context prefilling (via sliding window mask trick), rather that its long-context performance. Just as our previous response, CoLM-Air enables faster prefilling above 4K tokens and supports larger batch size below 4K. We will add these details in the final version to make a clear understanding.
2. Besides, in our previous response, we also provide additional long-context experiments and just want to verify that our method can also be combined with existing long-context techniques to seamlessly extend its long-context capability via chain tuning. We also sorry for not reporting the baseline of our previous response, as we want to illustrate that our method can seamless extend its long-context capability via chain tuning by leveraging existing long-context tricks (e.g., sparse attention). Below are our results for comparisons, which reports acc under different length training over CoLM-Air by using sparse attention with chain tuning. Vanilla is the baseline.

These results can guarantee our method can improve its long-context capability when combining with existing technologies. Besides, to demonstrate the compatibility of our method with existing long-context techniques, we also presents the combining results with StreamingLLM on PG19 and arXiv.

We want to highlight that to enhance the long-context ability of CoLM-Air, it still requires the integration of existing long-context tricks (e.g., YaRN, Sparse Attention, StreamingLLM and so on). Therefore, we admit that to achieve STOA results, it is still necessary to explore how to combine the advanced long-context methods into CoLM-Air settings, but this task is not the main goal of our paper.

Overall, in our paper, we want to claim the main advantages of CoLM-Air is its prefilling efficiency, rather than its long-context capability. Our original experiments mainly use sliding window mask and then evaluate its prefilling speed under long-context data. We apologize for making this misunderstanding. We will add these discussions into our final manuscript to make a clear understanding and avoid any misunderstanding. In the future, we will also consider to conduct more experiments like Needle-In-a-Haystack to explore the potential setting of our CoLM architecture in long context tasks.

Thank you once again for your insightful engagement, which has helped us improve the clarity of our paper.

Dear Reviewer fgKQ,

Thanks for your comment and we hope our previous response can address your concerns. To further clarify your remaining concern regarding long-context performance, we would like to offer a concrete example to explain the mechanism of our model.

As mentioned in Section 4.4, our CoLM-Air can reduce the prefilling speed as we only need to calculate all keys and values within the first chain and then share it to the latter chains. Therefore, the capability of our model to handle long sequences is dependant on the design of the first chain. In the original setting, our model is trained with 4K, and thus we report the prefilling speed by using sliding window masks. Conversely, if the first chain is a model proficient in long-context tasks, our model can also inherit its advantages.

To make a better understanding, we perform the "Needle in a haystack" on the baseline LLaMA-3.2-1B and our expanded version. We can find their results are completely equivalent, because our expanded version also reuse the keys and values from the first chain and inherit its capability in processing long-sequence:

From these results, we can find that the long-context capability of our method comes from the first chain. Our first chain can also use some other methods to further improve its long-context capability (e.g., YaRN, StreamingLLM, Sparse Attention and so on). For model itself, our CoLM-Air design mainly reduces the prefilling speed by only using the first chain. Besides, just as mentioned in Line 124, the first chain could be any model, which proves the flexibility and generality of our model.

We hope our experimental explanation can resolve your concerns and we will add these discussions into our final version to make a clear understanding about our design. Thank you for your comments.