Last Updated on 2024-08-19 by Clay
Cross-lingual Modular (X-Mod) 是一個有趣的語言模型架構，透過模組化不同語言的參數作為模組單元（Module Unit）組裝在模型中，好讓語言模型在新增全新語言時，可以使用獨立的參數進行微調，而（比較）不會出現災難性遺忘的窘境。
會來看這篇論文，最大的原因是因為最近看到 ColBERT-XM: A Modular Multi-Vector Representation Model for Zero-Shot Multilingual Information Retrieval，然後同事拿這個 ColBERT 的模型在我們自己內部的中文資料上測試，驚訝地發現它遠比其他模型高了 10% 左右的分數。
快速讀了這篇論文後，發現它最想解決的核心痛點，是被稱為多語言詛咒（Curse of Multilinguality）的現象 —— 意即當語言模型想要覆蓋、學習更多的語言時，其原本學會的語言性能會開始衰退。於是才會提出這樣的新架構去解決這個問題。
以下就來探究它和其他語言模型的架構差異吧！
模型架構
模型的基礎架構仍然是 Transformer，該有的嵌入層、注意力機制、前向傳播層通通一應俱全；而最大的不同就是它在那之後加入了模組化層（Modular Layer），輸入前與輸出後也都會再額外接 LayerNorm，並且留了一個殘差通道把進入模組化層之前的資訊導向輸出後，可以想像這讓梯度更好流通、訓練更穩健，算是一個常規的添加配置，想必也做過實驗測試後才保留的。
我們可以注意到多添加的 modular layers 是從 Language 1 到 Language n，代表有多少種語言就會有多少個不同的模組存在於這裡。不過在推理時，只會同時啟用一個語言模組。
從 HuggingFace 的模型卡中，我們可以看到其需要顯式地指定啟用的語言：
from transformers import XmodModel
model = XmodModel.from_pretrained("facebook/xmod-base")
model.set_default_language("en_XX")並且在原始碼中可以看到在接受 hidden_states 傳入時，只有對應啟用語言的適配器會處理 hidden_states：
...
def lang_adapter(self, lang_ids: torch.Tensor, hidden_states: torch.Tensor):
# Process subsequent samples with the same lang_id in parallel
lang_ids, lang_lengths = torch.unique_consecutive(lang_ids, return_counts=True)
if not self.ln_before_adapter:
residual = hidden_states
if self.adapter_layer_norm is not None:
hidden_states = self.adapter_layer_norm(hidden_states)
elif self.adapter_reuse_layer_norm:
hidden_states = self.LayerNorm(hidden_states)
if self.ln_before_adapter:
residual = hidden_states
split_hidden_states = torch.split(hidden_states, lang_lengths.tolist(), 0)
lang_wise_outputs = []
for i, (lang_id, split_hidden_state) in enumerate(zip(lang_ids, split_hidden_states)):
lang = list(self.adapter_modules.keys())[int(lang_id.item())]
lang_wise_outputs.append(self.adapter_modules[lang](split_hidden_state))
hidden_states = torch.cat(lang_wise_outputs, 0)
hidden_states = self.dropout(hidden_states)
hidden_states += residual
return hidden_states不過正是因為如此，所以不同語言的適配器是分開的。如下圖 (b) X-MOD 的部份，當我們需要添加新的語言時，可以只添加新的適配器（以及 Embedding Layer 的詞彙量，下一節會提到）。
模型訓練方式
X-Mod 的模型的訓練可以分成三個不同的階段來理解：
- 預訓練（Pre-training procedure）
- 擴增新語言（Extending to new languages）
- 下游任務微調（Fine-tuning on downstream tasks）
1. 預訓練（Pre-training procedure）
X-Mod 模型的預訓練是採用經典的掩碼語言建模（Masked Language Modeling, MLM），按照一定破壞比例隨機把 token 遮蔽，再訓練時透過上下文解碼還原出正確的 token，增加模型對文本的理解。
在這個任務中，會採用各種不同語言的資料，各種語言的資料會訓練到共享的嵌入層、注意力機制和前饋層；在不同語言的模組層時，則會由資料的設定導向各自對應的語言模組。
2. 擴增新語言（Extending to new languages）
由於模組化的設計，所以在添加新的語言模組層時並不會太多地影響到原本學習的語言（這也是本篇研究想要解決的核心問題）。除了添加專屬的語言模組層外，也得根據語言多出的詞彙量，更新嵌入層的尺寸和斷詞器（Tokenizer）。
在這一階段同樣是做 MLM 的訓練，不過只訓練嵌入層和語言模組層，其餘部份的參數是凍結（freeze）的。
這一步驟結束後，我們便得到了一個擁有新語言理解能力的語言模型；但直覺上來說，雖然新語言的理解能力是有了，但是為什麼不直接訓練一個專門學習新語言的語言模型呢？我的理解是，在下一步驟的下游任務微調中，是可以進一步讓模型擁有跨語言的理解能力的。
3. 下游任務微調（Fine-tuning on downstream tasks）
這一步驟主要是讓語言模型對跨語言的文本（比方說中英混雜的資料）擁有理解能力的，其作法也非常單純直接。
我們剛剛得到了新的語言模組層和嵌入層，但現在反而是凍結這些參數不讓其加入訓練；這次我們要訓練的是不同語言模組層共享權重的部份，也就是剛剛提到注意力機制和前饋層。不過路由的分配上是要導向剛剛新添加的語言模組層。
（2024/08/19 更新：我的同事提醒我說第三階段從論文中並不一定是採用多語言的資料集，而是可能只使用單一語言資料也可以，原文寫為 source language data）
結語
其他的細項評估就略過不看了，畢竟對我而言已經在我工作上的實際資料證實過後續的 ColBERT-XM 性能較其他開源模型更佳。所以本篇的閱讀著重在理解其內部的架構、與其想要解決的痛點。
若要更發散一點地去思考，這種模組化的架構是否不只侷限於語言呢？畢竟適配器本來就是為了下游任務微調而產生的，是否也可以讓模組層負責不同的下游任務來避免開頭提到的災難性遺忘呢？
另外，現在的許多 LLM 可以透過思維鏈（Chain-of-Thought, CoT）去產生更好的效果，是否也可以讓模型透過模組化層一邊做思考、一邊生成回覆呢？或許找一找已經有相關的論文了也說不定。
References
- arXiv - Lifting the Curse of Multilinguality by Pre-training Modular Transformers
- arXiv - ColBERT-XM: A Modular Multi-Vector Representation Model for Zero-Shot Multilingual Information Retrieval