Last Updated on 2024-07-25 by Clay
介紹
Mistral 7B 是 2023 年 09 月 27 日所提出的大型語言模型（Large Language Model, LLM），由同名的 Mistral AI 團隊訓練而成並將其權重開源釋出。有趣的是，它所採用的授權是相當自由的 Apache 2.0，不像 Llama 2 採用了自己的 Llama 授權條款，所以是真正意義上的『開源』（Llama 授權需要在服務量達到七億時跟 Meta AI 討論授權問題）。
研發團隊很自豪地宣佈，Mistral 擁有以下突破性的表現：
- 在全部 benchmarks 上贏過 Llama 2 13B
- 在部分 benchmarks 上贏過 Llama 1 34B
- 接近 CodeLlama 7B 的程式能力，並且在英文任務上仍然優秀
- 使用 Grouped-query Attention (GQA) 加速推理
- 使用 Sliding Window Attention (SWA) 以更小的開銷換取更長的序列
而近期（2023/11）我所知的最好的一些 7B 模型，如 zephyr-7b-beta、OpenHermes-2.5-Mistral-7B，也都是基於 Mistral 去微調的。由此可見，Mistral 作為 7B 量級的模型其潛力相當不錯，已經可以適用於許多生成任務了。
架構
Mistral 本質上仍然是 Transformer 架構，並且同樣是 auto-regressive 的模型，使用之前的輸入來產生下一個 token —— 直到把序列補完為止。
而 Mistral 比較特別有提出來說明的架構分別是：滑動窗口注意力（Sliding Window Attention）、Rolling Buffer Cache 以及 Pre-fill Chunking。
以下分別簡單談談我對他們的理解。
Sliding Window Attention (SWA)
雖然 Sliding Window Attention 並不是由 Mistral 最早提出，但這在架構選擇上是一個相當有趣的選擇。比方說滑動窗口所關注的長度有有限的，所以在長上下文長度的任務情境中，可以理解為時間複雜度為 O(n*w) （w 為視窗尺寸）而非原先自注意力機制的 O(n^2)。
也即是說，在生成本來假設我們的 KV 維度為 (batch_size, num_head, seq_len, head_size)，其中我們的 seq_len 只會保持固定的長度，在超過過定長度的 token 資訊將不會參與注意力機制的計算。
對我來說，這種高效率的 Attention 機制仍然有令人憂慮的點，比如其擴展上下文的能力究竟如何？完成『大海撈針』測試的能力如何？（甚至是 2024/04 現在更加新穎的捕捉星星測試又如何？）
只要能證明其 Sliding Window Attention，我會毫不懷疑這是 Transformer 架構的穩定下一個大版本。
除此之外，由於 Mistral 是採用 Group-Query Attention（GQA）的機制，所以 kv 的 num_head 的數量也並非與 Q 完全一致，更進一步減少參數量。
另外，Mistral 的開發團隊想要強調的是，即便使用了固定的注意力視窗，也不代表計算的資訊僅止於當前 token 的視野範圍，因為模型在架構中會堆疊非常多層，所以在後頭的權重所計算的資訊，其包含一定數量的先前 token 資訊 —— 這就是 Effective Context Length 的圖示想要說明的。
Rolling Buffer Cache
另外，也正因為採用了 Sliding Window Attention，故其 KV Cache 所儲存的資料也可以只存固定尺寸，所以才會有 Rolling Buffer Cache 的出現。
簡單來說，跟傳統的 KV Cache 會無限期地成長不同，在 Sliding Window Attention 機制中，KV Cache 有其上限（畢竟 Window Size 固定），所以可以滾動式地更新 KV Cache。
Pre-fill and Chunking
Pre-fill 的意思基本上就是預先填充 KV Cache。
有趣的是，Mistral 的模型架構雖然都有其道理，但並不是多麼新穎神奇的方法，其模型效果如此之好，許多研究者不禁假設其背後所使用的訓練資料品質非常之高。
根據某份設定檔透漏，很有可能有 8T tokens 的訓練資料量級。這件事情告訴我們，在資料品質的面前，其模型的架構跟調校可能並沒有那麼高的權重。
效果
其這張技術報告中的實驗結果，可以窺見 Mistral 其性能之強。後續基於 Mistral 微調的 LLM 效果也都相當不錯，例如 HuggingFace 的 Zephyr 以及 NousResearch 的 OpenHermes，都曾是我愛用來測試研究的 LLM。
心得
這篇我早就開始寫了，大概從 2023 年 11 月開始動筆，一直斷斷續續寫到了 2024 年 2 月（現在更是 2024 年的 4 月了）。雖然 Mistral-7B 還不真正說已經過氣、被淘汰，但 MistralAI 團隊後續所做的 Mixtral-8x7b 的效果實在太驚人，在開源社群引起了一波 MoE 架構的風潮，也再度刷新了開源 benchmarks 的分數，追趕甚至超過了 GPT-3.5。
後來其 Mistral-large 的效果也同樣令人驚豔，可惜的是並沒有開源，讓人相當失望。
最近我研究了 HuggingFace transformers 套件中的原始碼，嘗試自己刻 Mistral 的模型架構，有了一些不錯的收穫，也對 Mistral 模型的設計有了更深入的了解，希望之後有機會能分享自己的程式碼。
在複現 Mistral 架構的途中，我也參考著實現了 RoPE 的實作，可以參考：[Machine Learning] 旋轉位置嵌入 （Rotary Position Embedding, RoPE）筆記