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

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

1.Pre-Normalization

基礎(chǔ)的Transformer每一層的計(jì)算邏輯是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的輸入位置。兩者的差異如下圖和表所示。

圖片

圖片

文中通過分析實(shí)驗(yàn)和分析發(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的改進(jìn)，出自論文Root Mean Square Layer Normalization（NeuIPS 2019）中。LayerNorm每一層需要計(jì)算輸入的每條樣本向量各個(gè)元素的均值和方差，然后對輸入進(jìn)行歸一化處理。這種方法雖然可以提升訓(xùn)練穩(wěn)定性，但是大幅增加了模型中的計(jì)算開銷。如下圖，相同步數(shù)下使用LayerNorm可以降低1.6的loss，而相同時(shí)間下只能降低1.1的loss，說明LayerNorm的計(jì)算開銷較大。

圖片

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

圖片

從實(shí)驗(yàn)效果看，使用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ù)的形式分別如下：

圖片

圖片

圖片

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

圖片

4.Rotary Position Embedding

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

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

圖片

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

圖片

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

圖片

5.Grouped-Query Attention

為了提升infer的運(yùn)算效率，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)都會(huì)進(jìn)行一次QKV的映射，再進(jìn)行attention計(jì)算。而Grouped-query會(huì)將多個(gè)head分成多個(gè)組，每個(gè)組內(nèi)的query使用相同的K和V，而不再每個(gè)head都進(jìn)行一次映射，以此節(jié)省計(jì)算開銷。

圖片