簡介

上下文長度 (Context Length) 指的是一個模型輸入 + 輸出最多可以到多少個 Tokens。例如一開始的 LLaMA 只支援 2048 個 Tokens，後來 Llama 2 支援到 4096 個 Tokens，模型的可用性就大幅提昇了。

Context Length 對中文應用尤其重要，因為中文所佔用的 Token 數量通常都比英文來得多。若是進行 Retrieval-Based 的應用，那消耗的 Token 數是相當可觀。今天就來探討 LLM Context Length 的相關議題。

ChatGPT & Claude

目前 OpenAI 的 gpt-3.5-turbo 可以支援 4K Tokens，而 gpt-3.5-turbo-16k 則可以支援到 16K Tokens。在 GPT-4 這邊，最低則是 8K Tokens，而 gpt-4-32k 則可以支援到 32K Tokens。

其實對一般應用來說 32K Tokens 已經相當足夠，然而 GPT-4 的費用還是相對高一些，且 GPT-4 32K 收取的費用又更高了。目前 GPT-4 8K 的費用是 0.03/0.06 鎂，而 GPT-4 32K 則是 0.06/0.12 鎂，長度多了四倍，價格多了兩倍。

但是會特地用到 32K 版本的應用，消耗的 Token 數量鐵定也是特別高！長期用下來其實成本是滿大的，現階段要創造出與此成本匹配的收益，是不太容易的。

來自 Anthropic 的 Claude 最高可以支援到 100K Tokens，是個相當驚人的長度。筆者最近終於收到 Claude 的測試邀請，並實際輸入了幾篇非常長的文章，其效果是真的挺不錯的，而且處理速度也滿快的。

可惜的是，這些 Long Context LLM 並沒有發表他們的實際訓練手法。

也有可能他們純粹就是拉長訓練，反正他們機器也是夠多夠硬。但這對一般中小型研究單位而言不太實際，於是大家開始思考有沒有更低成本又能拉長長度的方法。

Positional Interpolation

Positional Interpolation 簡稱 PI，是由 Meta 發表的一篇論文。在 Llama 裡面，使用 Rotary Positional Embedding (RoPE) 來表達位置資訊。在 RoPE 裡面，主要使用 Cosine 函數來表達相對位置的資訊，因此可以將 RoPE 表達為一個餘弦線圖。

如果我們要將 RoPE 支援的長度擴大，一個最簡單的想法就是直接外插 (Extrapolation) 上去，把 Cosine 函數接續下去，如下圖所示：

但是實際使用時發現外插效果並不好，於是這篇論文提出了內插 (Interpolation) 的方法，直接將 Cosine 函數給「拉寬」到指定長度，像這張圖一樣：

因此這個 Cosine 波的週期還是一樣的，但是可以覆蓋的範圍變廣了。就像是幫相機裝上廣角鏡頭一樣，可以看到的視野變寬廣了，但是拍出來的照片解析度還是一樣的，因此在細節上會看的相對模糊，但是可以看到的風景樣貌整體上變多了。

在 CodeLlama 的論文中提到，對 RoPE 進行內插確實是重要的一件事情。但是 CodeLlama 並不是像 PI 論文所述，直接將 Cosine 函數拉寬到指定長度，而是用了一個特殊的線性變換來改造 RoPE。因此 CodeLlama 可以在訓練階段將原本 Context Length 4K 的 Llama 2 內插擴展到 16K 的長度，並在推論時可以外插到將近 100K 都能維持穩定的效果。

LongLoRA

LongLoRA 為 CUHK 與 MIT 在 2023 年 9 月發表的論文，他們提出 Shift Short Attention (S^2-Attn) 來減少 Attention 的計算量，並結合 LoRA 減少記憶體消耗量，使我們能夠使用更少的硬體成本來將 LLM 的 Context Length 擴大。

根據作者的實驗結果指出，如果單純只把 LLM 的 Context Length 提高而不去做額外的 Finetuning，那模型的效果會很差，評估起來的 PPL 如下：

但是對高 Context Length 的模型進行訓練，需要的硬體成本又太高了。因此作者提出結合 S^2-Attn 與 LoRA+ 的 LongLoRA 訓練方式，讓我們可以用相對低成本的機器來訓練高 Context Length 的模型。

S^2-Attn

S^2-Attn 的做法如下：

1. 將 Attention Heads 切成兩半。
2. 將其中一半的 Attention 位移 (Shift) 一點點。
3. 將所有 Attention 切成好幾個 Groups。
4. 進行 Group-Wise Self-Attention 運算。

位移的目的是要讓 Group 之間的資訊可以交流，能夠提昇模型訓練的效果。而切成 Group 做 Group-Wise Self-Attention 則可以降低 Attention Layer 的計算量。

從上圖可以看到，原本的 Self-Attention 要把整個大方格都計算一遍，但 Group-Wise Self-Attention 只需要計算對角線上的那幾格就好。

作者強調，這個 S^2-Attn 不僅能減少訓練階段的計算量，使用 S^2-Attn 訓練出來的模型，在推論階段也能夠直接使用原本的 Full Attention 進行運算，因此最後是不會改動到模型架構的。

LoRA+

原本的 LoRA 預設是不會訓練 Embedding & Normalization 這兩種 Layer 的，但是 LongLoRA 的作者發現，如果將這兩種 Layers 設定為可訓練，那 LoRA 的訓練效果就會追上 FFT 的做法，其效果遠比提高 LoRA Rank 來得好：

要啟用 Embedding & Normalization 的程式碼如下：

model = ModelCls.from_pretrained(...)
model = get_peft_model(model, ...)

for name, param in model.named_parameters():
    if "embed" in name or "norm" in name:
        param.requires_grad_()
        break

其中 embed 與 norm 需要根據模型實際 Layer 名稱調整，這裡是 Llama 的 Layer 名稱。筆者使用 7B 模型測量了一下，如果使用原本的 LoRA，小模型的參數量約 4M。而改用 LoRA+ 的話，小模型的參數量會變成 135M。雖然說參數量大了不少，但這也僅是 7B LLM 的 2% 參數量而已，所以可訓練的參數量依然維持的很少。

小結

從實驗數據來看，雖然 LongLoRA 依然追不上 FFT 的效果，但應該是目前最經濟的一種 Long Context 訓練方法，可以參考官方的 GitHub 使用 LongLoRA 進行訓練，作者也有釋出實驗用的 FFT & LongLoRA 的 Long Context Llama 模型權重。

Llama Long

Llama Long 也是由 Meta 發表的一篇論文，這篇論文提出 Adjusted Base Frequency (ABF) 的做法來改造 RoPE。論文中列出了數條極為凶悍的數學公式分析此做法，最後畫出以下概念圖：

從圖中可以看到，使用 ABF 的 RoPE 的「頻率」是比較高的，因此 ABF RoPE 的粒度相對較細，可以獲取相對細節的資訊，較細的粒度也能使模型更有辦法分辨不同位置的資訊。

StreamingLLM & Attention Sinks

StreamingLLM 是由 MIT, Meta, CMU 合作發表的論文，他們提出了 StreamingLLM 框架，試圖解決 LLM 的 Context Length 問題。

在 Long Context 的情況下，如果計算 Full Attention 會帶來相當巨大的計算量與記憶體消耗。因此有些人可能會採用 Window Attention 的做法，也就是只計算過去一段固定長度的 Attention 就好，超過的就通通丟掉不算。

但是這樣的做法最後會導致 PPL 炸開而無法使用。另外一種做法是結合 Sliding Window 的機制，將整個 Window 滑動一段距離，然後重新計算一次所有的 Attention 分數。這樣的做法雖然能有相對穩定的 PPL，但是遇到需要重新計算的地方時，會明顯感受到很長的停頓。

作者在探討 Window Attention 為何會爆掉的過程中，觀察到一個相當有趣的現象，叫做 Attention Sinks。Sink 是指沉陷、沒入的意思，在 LLM 計算 Attention Score 的時候，作者發現輸入模型的前幾個 Token 平均都能獲得相當高的分數。

即便把 Initial Tokens 換成換行符號這類不具備任何實質語意的 Tokens，依然會獲得相對高的分數。而 Window Attention 之所以會爆掉，就是因為 Initial Tokens 被截掉了。

之所以會有這個現象，作者提出了兩個猜測：

1. 因為 Attention 裡面的 Softmax 運算，無論何時都不會讓所有 Token 都是零分，即便下個預測的 Token 其實並不需要 Attend 在任何 Context 上，因此這類的分數就會被 Initial Tokens 給吸收。
2. 在 Autoregressive Decoding 的過程中，所有的 Token 都能看上 Initial Tokens 一眼，因此大家都認得，所以 Initial Token 能夠捕捉一些不重要的 Attention。

作者將這些 Initial Tokens 稱呼為 Attention Sinks，這些 Tokens 相當重要，就像是船錨一樣可以把整艘船穩住而不會亂飄。基於這個觀察，提出了 StreamingLLM 的做法，一樣是基於 Window Attention，但是把 Initial Tokens 給保留下來。

這樣的做法非常驚人的有效果，不僅能穩住 PPL 不會爆掉，而且推論上也相當有效率，計算量不會隨著長度而增加。因此作者建議以後大家訓練 LLM 的時候，一定要多放幾個 Attention Sinks，這樣 StreamingLLM 就可以利用 Attention Sinks 變相的讓 LLM Context Length 變得「無限長」。這樣的研究發現，再次重申了 BOS Token 的重要性。

StreamingLLM 並不是真的把 Context Length 變長，而是一種能讓模型穩定處理 Long Context 的方法。

結論

今天介紹了幾篇與擴展 LLM Context Length 有關的論文，各自的做法其實都還滿有效果的。有些論文也一同釋出了成果模型或程式碼，可以去驗證其效果如何。在 Long Context LLM 的議題裡面，無論是訓練手法還是推論加速，都逐漸有不同的做法被提出。再讓子彈飛一會，未來 LLM 的 Context Length 能發展到什麼樣的程度是相當令人期待的。

參考

• HF TGI Docs: RoPE Scaling
• arXiv Paper - Giraffe: Adventures in Expanding Context Lengths in LLMs
• arXiv Paper - Extending Context Window of Large Language Models via Positional Interpolation
• arXiv Paper - Effective Long-Context Scaling of Foundation Models
• arXiv Paper - LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models
• arXiv Paper - Efficient Streaming Language Models with Attention Sinks
• arXiv Paper - Code Llama: Open Foundation Models for Code
• iThome News: Llama Long