Chain-of-Model Learning for Language Model

本文提出 Chain-of-Model (CoM) 学习范式，通过在 Transformer 架构中引入因果依赖的多尺度表示（Chain-of-Representation），实现高效模型扩展和弹性推理，实验表明 CoLM 系列在性能上与标准 Transformer 相当，同时在预填充速度和灵活性上具有优势。

Large Language Model, Transformer, Pre-training, Efficiency, Multimodality

Kaitao Song, Xiaohua Wang, 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

Microsoft Research, Fudan University, Zhejiang University, ShanghaiTech University

Generated by grok-3

Background Problem

随着大型语言模型（LLMs）的快速发展，模型参数规模呈指数级增长，从十亿级到万亿级，导致训练成本极高且推理时无法灵活适应不同部署环境。现有扩展策略（如 Dense 或 MoE 架构）存在两个主要问题：一是无法基于已有模型增量学习，总是需要从头训练；二是缺乏动态调整模型规模以适应不同任务需求的能力。论文提出了一种新的学习范式，旨在解决这些问题，通过引入多尺度表示和因果依赖机制，实现模型的高效扩展和弹性推理。

Method

论文提出了 Chain-of-Model (CoM) 学习范式，其核心思想是通过 Chain-of-Representation (CoR) 将隐藏状态在维度上分割为多个子表示（称为链，chains），并在每一层引入因果依赖，确保每个输出链只能依赖于输入中的前序链。具体步骤如下：

Chain-of-Representation (CoR)：将隐藏表示 $x \in \mathbb{R}^D$ 分割为 $n$ 个子表示 $\xi(x, n) = \{x_1, \dots, x_n\}$ ，每个子表示对应一个链，代表不同尺度信息。
Chain-of-Layer (CoL)：设计网络层，使得输出链 $y_i$ 仅依赖于输入链 $x_{\leq i}$ ，确保因果性。这种设计应用于线性层、注意力模块和前馈网络等组件。
Chain-of-Language-Model (CoLM)：将 CoL 思想集成到 Transformer 架构的每一层，支持单次前向传播中多尺度训练，并通过不同链数量提供多个子模型用于弹性推理。
CoLM-Air：引入 KV 共享机制，仅在第一个链计算键（Key）和值（Value），并共享给其他链，以提升推理灵活性（如无缝切换模型规模、加速预填充）。 批判性思考：虽然 CoM 的因果设计理论上支持增量扩展和弹性推理，但其复杂性可能增加实现难度和计算开销。此外，KV 共享机制虽然提升了灵活性，但可能限制模型对不同链的独立表达能力，导致性能下降。

Experiment

实验基于 SlimPajama 数据集（600B tokens），预训练了约 200B tokens，使用 32 个 NVIDIA A100 GPU，序列长度为 4096，批大小为 1024。模型配置包括不同链数量（如 C={16,16} 和 C={8,8,8,8}）和维度设置，与 LLaMA-3.2-1B 作为基线对比。评估涵盖常识推理任务（如 HellaSwag、Obqa）和其他场景（如链扩展、弹性推理、预填充速度）。

性能：CoLM 在性能上与基线相当，C={16,16} 配置略优于 C={8,8,8,8}，但 CoLM-Air 由于 KV 共享导致性能略有下降（平均准确率下降约 1-2%）。
链扩展：通过将预训练模型作为第一链并扩展新链，性能略有提升（如 Tiny-LLaMA-v1.1 提升 0.92%），验证了增量扩展的可行性。
弹性推理：CoLM-Air 能通过不同链数量提供多个子模型，参数从 0.33B 到 0.86B，性能随规模增加而提升。
预填充速度：CoLM-Air 在预填充阶段显著加速，例如处理 1M tokens 时比 LLaMA 快 1.6-3.0 倍，结合 MInference 可达 27 倍加速。 批判性思考：实验设置受资源限制，预训练规模较小，可能无法反映方法在大规模场景下的真实表现。此外，虽然预填充速度提升显著，但性能下降是否为可接受的权衡未充分讨论。与现有方法的直接对比缺失，难以判断 CoM 的相对优势。

Further Thoughts

尽管 CoM 提供了一种有趣的模型扩展思路，但其实际应用价值仍需进一步探索。例如，KV 共享机制虽然提升了推理灵活性，但可能限制模型在多尺度任务中的表达能力，是否可以通过动态调整共享策略来缓解这一问题？此外，CoM 的因果设计与现有的混合专家（MoE）架构有一定相似性，未来可以尝试将 CoM 与 MoE 结合，利用专家机制进一步提升链间的信息隔离和特化能力。在更广泛的领域中，CoM 的多尺度思想或许可以应用于多模态模型，通过为不同模态分配不同链条，实现模态间的渐进式信息融合，这可能为构建更高效的多模态基础模型提供新思路。