LLaMA是目前很多SOTA開源大模型的基礎(chǔ)，包括DeepSeek、千問等在內(nèi)的很多大模型的模型機構(gòu)大體上都沿用了LLaMA。因此，LLaMA在模型結(jié)構(gòu)的設(shè)計上可以說是目前大模型的一種最佳實踐。這篇文章就給大家匯總一下LLaMA的核心模型結(jié)構(gòu)。

LLaMA的主體結(jié)構(gòu)仍然基于Transformer，本文主要介紹LLaMA各個版本相比于Transformer的改進部分，包括Pre-Normalization、RMSNorm、SwiGLU激活函數(shù)、Rotray Embedding等部分。

1.Pre-Normalization

基礎(chǔ)的Transformer每一層的計算邏輯是attention->add->norm->ffn->add->norm。其中norm為LayerNormalization。這種在attention和add之后的norm稱為Post-Normalization。而LLaMA中采用了Pre-Normalization，主要源自于ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE（ICLR 2020）這篇文章。其核心是將LayerNormalization放在每層Transformer的輸入位置。兩者的差異如下圖和表所示。

圖片

圖片

文中通過分析實驗和分析發(fā)現(xiàn)，Post-Normalization方法在訓(xùn)練開始階段模型的輸出層的梯度范數(shù)比較大，模型層數(shù)越靠后梯度越大，這給訓(xùn)練的初始階段帶來了不穩(wěn)定性。而使用Pre-Normalization，各層的梯度范數(shù)基本想同，因次使用Pre-Normalization可以提升訓(xùn)練的穩(wěn)定性。此外，通過warm-up等策略，讓初始的學(xué)習(xí)率小一些，可以緩解Post-Normalization的這種初始階段梯度范數(shù)較大的問題。

圖片

2.RMSNorm

基礎(chǔ)的Transformer在norm部分采用的是LayerNormalization，在LLaMA中則使用了RMSNorm，是一種針對LayerNormalization的改進，出自論文Root Mean Square Layer Normalization（NeuIPS 2019）中。LayerNorm每一層需要計算輸入的每條樣本向量各個元素的均值和方差，然后對輸入進行歸一化處理。這種方法雖然可以提升訓(xùn)練穩(wěn)定性，但是大幅增加了模型中的計算開銷。如下圖，相同步數(shù)下使用LayerNorm可以降低1.6的loss，而相同時間下只能降低1.1的loss，說明LayerNorm的計算開銷較大。

圖片

為了解決這個問題，文中提出的RMSNorm將LayerNorm進行了簡化，原來的LayerNorm是計算均值和方差，而RMSNorm改為計算元素的均方根，用均方根進行歸一化。這樣做雖然犧牲了LayerNorm對輸入數(shù)據(jù)的re-centering能力，但是最終效果和LayerNorm差不多，說明LayerNorm的有效并不來源于re-centering。

圖片

從實驗效果看，使用RMSNorm模型收斂更快，也取得了更好的效果。

圖片

3.SwiGLU

基礎(chǔ)的Transformer結(jié)構(gòu)中，激活函數(shù)使用的是ReLU。而LLaMA中將所有ReLU更換為SwiGLU激活函數(shù)，以此來提升模型的表現(xiàn)。SwiGLU發(fā)表于文章GLU Variants Improve Transformer（2020）中，SwiGLU是Swish激活函數(shù)和GLU激活函數(shù)的結(jié)合，Swish、GLU、SwiGLU激活函數(shù)的形式分別如下：

圖片

圖片

圖片

這里面的核心是利用了門控的思路，原始的輸入過一個sigmoid，得到一個0~1的和輸入向量維度相同的gate，以此對原始輸入各個維度進行縮放。論文中在基于Transformer Encoder-Decoder的語言模型中，分別進行了預(yù)訓(xùn)練、finetune等不同激活函數(shù)的效果測試，都驗證了SwiGLU相比ReLU等其他激活函數(shù)可以取得更好的效果。

圖片

4.Rotary Position Embedding

基礎(chǔ)的Transformer中采用絕對位置編碼，即每個位置映射成一個embedding，或者用三角函數(shù)定義位置編碼。但是絕對位置編碼的核心問題在于，無法支持超出訓(xùn)練長度的建模，因為這些embedding在訓(xùn)練過程中沒見過。在文本生成中，需要靈活支持很長的文本生成場景，因此絕對位置編碼的弊端就顯露出來。

旋轉(zhuǎn)位置編碼發(fā)表于論文ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING（2023）中，是一種相對位置編碼方法，能夠靈活支持不同長度的文本輸入和輸出。其核心出發(fā)點是將位置編碼信息放到attention的計算中，通過query和key的內(nèi)積引入相對編碼，目標是尋找一個函數(shù)f，其內(nèi)積能夠?qū)崿F(xiàn)只和query和key的向量輸入，以及兩個元素的相對位置相關(guān)：

圖片

對應(yīng)的函數(shù)在2維的情況下可以表示為如下形式，滿足上述需求：

圖片

這里利用了復(fù)數(shù)的計算實現(xiàn)相對位置編碼，整體的計算流程未，對于 token 序列中的每個詞嵌入向量，首先計算其對應(yīng)的 query 和 key 向量，然后對每個 token 位置都計算對應(yīng)的旋轉(zhuǎn)位置編碼，接著對每個 token 位置的 query 和 key 向量的元素按照 兩兩一組應(yīng)用旋轉(zhuǎn)變換，最后再計算 query 和 key 之間的內(nèi)積得到 self-attention 的計算結(jié)果。下圖是論文中的示意圖計算流程示意圖。

圖片

5.Grouped-Query Attention

為了提升infer的運算效率，LLaMA將基礎(chǔ)Transformer中的self-attention改成了GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（2023）中提出的grouped-query attention。在多頭self-attention中，每和head內(nèi)都會進行一次QKV的映射，再進行attention計算。而Grouped-query會將多個head分成多個組，每個組內(nèi)的query使用相同的K和V，而不再每個head都進行一次映射，以此節(jié)省計算開銷。

圖片