Transformer架构基础：注意力机制、层归一化、前馈网络等核心组件

Transformer架构自2017年提出以来，已成为大语言模型（LLM）的绝对主流底座。无论是GPT系列、LLaMA还是ChatGLM，其核心计算单元都遵循同一套设计范式。理解Transformer的组件，不仅有助于读懂模型结构，更是后续掌握TensorRT-LLM中算子融合、KV Cache、量化策略的前提。本节将拆解其三大核心组件：多头注意力（MHA）、层归一化（LayerNorm）与前馈网络（FFN），并说明它们如何协同完成一次token的上下文感知变换。

3.1 核心组件一：多头注意力机制

注意力机制的本质是“按相关性加权聚合”。对于输入序列中的每个token，模型通过查询（Query）、键（Key）、值（Value）三个向量计算它与其他所有token的相关性分数，再用该分数加权求和所有Value，从而让每个位置的输出都携带全局上下文信息。

多头注意力则是将这一过程并行执行多次。假设模型隐藏维度为 \(d_{model}\)，将其切分为 \(h\) 个头，每个头独立在 \(d_k = d_{model}/h\) 维子空间上做注意力计算。这样做有两个好处：其一，不同头可以关注不同的语义关系（如语法结构、指代关系、局部共现）；其二，降低了单头计算的维度压力，使模型在合理算力下获得更丰富的表示。

其数学表达如下：

Attention(Q, K, V) = softmax(QK^T / √d_k) · V
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O
其中 head_i = Attention(Q·W_Q^i, K·W_K^i, V·W_V^i)

在TensorRT-LLM的推理优化视角下，注意力机制有两个关键观察点：

• QK^T矩阵的规模随序列长度平方增长，是长序列推理的主要瓶颈，因此实际部署中常使用Flash Attention减少显存读写。
• KV Cache存储历史token的K和V矩阵，避免每步重复计算，但显存占用随序列长度线性增长，需要配合分页管理（Paged KV Cache）使用。

3.2 核心组件二：层归一化与RMSNorm

层归一化（Layer Normalization）的作用是对每个token的隐藏向量在特征维度上做标准化，使其均值为0、方差为1，然后通过可学习的缩放参数 \(\gamma\) 和平移参数 \(\beta\) 恢复表达能力。其公式为：

LayerNorm(x) = γ · (x - μ) / √(σ² + ε) + β
其中 μ = mean(x), σ² = var(x)，ε 为防止除零的小常数

相比BatchNorm，LayerNorm不依赖batch内其他样本，因此非常适合变长序列和在线推理场景。它被放置在注意力层和前馈层之后，用于稳定深层网络的梯度传播，加速收敛。

现代LLM（如LLaMA、Mistral）普遍采用RMSNorm（Root Mean Square Normalization）替代标准LayerNorm。RMSNorm去掉了均值中心化步骤，仅对均方根进行归一化：

RMSNorm(x) = x / RMS(x) · γ，其中 RMS(x) = √(mean(x²) + ε)

RMSNorm计算量更小（少一次均值减法），且在实验中被证明与LayerNorm效果相当。在TensorRT-LLM中，RMSNorm通常与后续的线性层或激活函数融合为QKV融合或MLP融合的一部分，减少显存往返。

3.3 核心组件三：前馈网络（FFN）

每个Transformer块中的FFN是一个两层的全连接网络，中间使用非线性激活函数。标准FFN结构为：

FFN(x) = Act(x · W_1 + b_1) · W_2 + b_2

其中 \(W_1\) 将维度从 \(d_{model}\) 扩展到 \(d_{ff}\)（通常为4倍），\(W_2\) 再压缩回 \(d_{model}\)。这个“先升维再降维”的设计让模型有机会在更高维空间中进行非线性特征交互。常用的激活函数包括ReLU、GELU和SwiGLU。近年来的LLM（如LLaMA）采用SwiGLU变体，将FFN改为三个权重矩阵的门控结构：

FFN_SwiGLU(x) = (SwiGLU_gate(x) ⊙ SwiGLU_up(x)) · W_down
其中 SwiGLU_gate(x) = x·W_gate, SwiGLU_up(x) = x·W_up, ⊙ 为逐元素乘

SwiGLU引入了门控机制，让网络能更灵活地控制信息流动，但参数量比标准FFN多约50%。在TensorRT-LLM中，FFN的权重通常被量化到INT8或FP8，并通过Weight Only方式存储，以降低显存占用和带宽压力。

3.4 组件协同与整体数据流

一个标准的Transformer解码器层按以下顺序执行：输入token序列先经过自注意力子层，输出与残差连接相加后做第一次LayerNorm；随后进入FFN子层，输出再次与残差相加并做第二次LayerNorm。残差连接（Residual Connection）让梯度能直接回传，避免深层网络退化。整体数据流如下图所示：

上图展示了一个解码器层的完整数据流。值得注意的是，残差连接将每个子层的输入直接加到输出上，形成“高速公路”，使得梯度可以无损回传。在TensorRT-LLM中，残差连接通常与LayerNorm融合为AddBiasResidualLayerNorm算子，减少一次kernel调用。

3.5 关键参数与计算量对比

为了更直观地理解各组件在推理时的计算与显存开销，下表以LLaMA-7B为例（\(d_{model}=4096\)，层数32，头数32，\(d_{ff}=11008\)），对比单层各部分的参数量与单token推理时的FLOPs：

从上表可以看出，FFN的参数量和计算量都占据主导地位，这也是为什么TensorRT-LLM对FFN层的量化（INT8/FP8）收益最明显。而注意力部分的计算量虽然相对较小，但其KV Cache的显存占用随序列长度线性增长，是长上下文场景下的主要瓶颈。

3.6 常见误区与工程复盘