언어 모델을 위한 체인-오브-모델 학습

Chain-of-Model Learning for Language Model

May 17, 2025

저자: 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

cs.AI

초록

본 논문에서는 각 계층의 은닉 상태에 인과 관계를 체인 형태로 통합하여 모델 학습의 확장 효율성과 배치 시의 추론 유연성을 크게 향상시키는 새로운 학습 패러다임인 Chain-of-Model(CoM)을 제안한다. 여기서는 Chain-of-Representation(CoR) 개념을 소개하며, 이는 각 계층의 은닉 상태를 은닉 차원 수준에서 다수의 하위 표현(즉, 체인)의 조합으로 형식화한다. 각 계층에서 출력 표현의 각 체인은 입력 표현의 선행 체인들만을 볼 수 있다. 결과적으로, CoM 프레임워크에 기반한 모델은 이전 모델(즉, 체인)을 기반으로 체인을 증가시켜 모델 크기를 점진적으로 확장할 수 있으며, 다양한 체인 수를 사용하여 탄력적 추론을 위한 다수의 하위 모델을 제공할 수 있다. 이 원리를 바탕으로, Transformer 아키텍처의 각 계층에 CoM 아이디어를 통합한 Chain-of-Language-Model(CoLM)을 설계한다. CoLM을 기반으로, 첫 번째 체인 내에서 모든 키와 값을 계산한 후 이를 모든 체인에 공유하는 KV 공유 메커니즘을 도입하여 CoLM-Air를 추가로 소개한다. 이 설계는 원활한 언어 모델 전환, 프리필링 가속화 등과 같은 추가적인 확장성을 보여준다. 실험 결과는 CoLM 계열이 표준 Transformer와 비슷한 성능을 달성하면서도, 학습 효율성을 개선하기 위한 점진적 확장 및 탄력적 추론을 위한 다양한 모델 크기 제공과 같은 더 큰 유연성을 동시에 가능하게 함을 보여준다. 이를 통해 언어 모델 구축을 위한 새로운 방향을 제시한다. 본 연구의 코드는 향후 https://github.com/microsoft/CoLM에서 공개될 예정이다.

English

In this paper, we propose a novel learning paradigm, termed Chain-of-Model (CoM), which incorporates the causal relationship into the hidden states of each layer as a chain style, thereby introducing great scaling efficiency in model training and inference flexibility in deployment. We introduce the concept of Chain-of-Representation (CoR), which formulates the hidden states at each layer as a combination of multiple sub-representations (i.e., chains) at the hidden dimension level. In each layer, each chain from the output representations can only view all of its preceding chains in the input representations. Consequently, the model built upon CoM framework can progressively scale up the model size by increasing the chains based on the previous models (i.e., chains), and offer multiple sub-models at varying sizes for elastic inference by using different chain numbers. Based on this principle, we devise Chain-of-Language-Model (CoLM), which incorporates the idea of CoM into each layer of Transformer architecture. Based on CoLM, we further introduce CoLM-Air by introducing a KV sharing mechanism, that computes all keys and values within the first chain and then shares across all chains. This design demonstrates additional extensibility, such as enabling seamless LM switching, prefilling acceleration and so on. Experimental results demonstrate our CoLM family can achieve comparable performance to the standard Transformer, while simultaneously enabling greater flexiblity, such as progressive scaling to improve training efficiency and offer multiple varying model sizes for elastic inference, paving a a new way toward building language models. Our code will be released in the future at: https://github.com/microsoft/CoLM.