Transformer 架构：LLM 如何理解 Token 之间的关系？

在前面的文章中，我们已经了解：

Tokenizer 把文本转换成 Token ID；
Chat Template 把对话组织成模型约定的输入格式；
上下文窗口限制模型一次能够处理多少 Token；
Transformer 最终输出 Logits，用来预测下一个 Token；
KV Cache 保存历史 Token 的 Key 和 Value，加速逐 Token 生成。

但中间还缺少最重要的一块：

Token ID 进入模型后，究竟经历了什么，模型又是怎样结合上下文理解每个 Token 的？

这个问题的答案就是 Transformer。

先看完整数据流

现代文本 LLM 通常采用 Decoder-only Transformer。一次前向计算可以简化为：

文本
  ↓
Tokenizer
  ↓
Token IDs
  ↓
Token Embedding + 位置信息
  ↓
Transformer Block × N
  ├─ Causal Self-Attention
  ├─ Residual Connection
  ├─ Normalization
  └─ Feed Forward Network
  ↓
Contextual Hidden States
  ↓
LM Head
  ↓
完整词表的 Logits
  ↓
选择下一个 Token

每个 Transformer Block 都会进一步更新 Token 的表示。经过多层之后，一个 Token 的向量不再只表示它自身，还融合了上下文中其他可见 Token 的信息。

Transformer 解决了什么问题？

语言是有上下文的。同一个词在不同句子中的含义可能完全不同：

1 2	苹果很好吃。苹果发布了新系统。

两个“苹果”拥有相同的 Token Embedding，但一个表示水果，另一个表示公司。

模型需要根据周围内容，动态更新每个 Token 的表示：

初始 Embedding
    “苹果”的基础表示
            ↓
Self-Attention 读取上下文
            ↓
上下文化 Hidden State
    当前句子中“苹果”的具体含义

Transformer 的核心价值不是简单记住单词，而是让序列中的每个位置根据其他相关位置动态计算自己的表示。

Transformer 出现之前

早期序列模型经常使用 RNN、LSTM 或 GRU，按照顺序逐步处理 Token：

1	Token 1 → Token 2 → Token 3 → ... → Token N

这种结构能够处理序列，但训练时存在一些限制：

时间步之间依赖较强，不容易充分并行；
很长距离的信息需要经过多步传递；
长序列上的梯度和记忆较难处理；
训练大规模模型的效率受到限制。

2017 年的论文《Attention Is All You Need》提出 Transformer，使用 Attention 为主的结构替代循环结构。训练时，同一层中的多个 Token 可以并行计算，这为扩大数据、参数和计算规模提供了重要基础。

需要注意，论文标题虽然叫“Attention Is All You Need”，但完整 Transformer 并不只有 Attention。它还包括：

Embedding；
位置信息；
Feed Forward Network；
Residual Connection；
Normalization；
输出层。

从 Token ID 到 Token Embedding

Tokenizer 输出的是整数 ID：

1 2	文本：我喜欢人工智能 Token IDs：[35946, 105205, 1773, 64312]

Token ID 本身只是词表索引，不表达“两个数字越接近，语义就越接近”。模型需要通过 Embedding Matrix 查表，把每个 ID 转换成连续向量：

1 2	Embedding Matrix Shape = [vocab_size, hidden_size]

查表后：

input_ids shape
= [batch_size, sequence_length]

hidden_states shape
= [batch_size, sequence_length, hidden_size]

例如词表大小为 151,936、隐藏维度为 896，那么 Embedding Matrix 可以理解为：

1	[151,936, 896]

每个 Token ID 会取出其中一行，得到长度为 896 的初始表示。

Embedding 是上下文无关的起点

同一个 Token ID 在查表时总会得到同一个基础向量。真正与上下文相关的变化发生在 Transformer Block 中：

1 2	Token Embedding：上下文无关的初始表示 Hidden State：经过上下文计算后的动态表示

因此，不能把 Embedding 和最终 Hidden State 当成完全相同的概念。

为什么需要位置信息？

Self-Attention 本身主要根据向量内容计算关系。如果不加入位置信息，模型很难区分：

1 2	小明批评了小红小红批评了小明

两句话包含相似的 Token，但顺序决定了不同含义。

原始 Transformer 使用正弦和余弦 Positional Encoding，并把它加入 Token Embedding。现代 LLM 还常见：

Learned Position Embedding；
RoPE，Rotary Position Embedding；
ALiBi；
相对位置偏置。

其中 RoPE 被许多 Decoder-only LLM 使用。它通常不是简单向输入向量加一个位置向量，而是对 Attention 中的 Query 和 Key 应用与位置相关的旋转，让注意力分数包含相对位置信息。

可以先记住：

Embedding 告诉模型“这是什么 Token”，位置信息告诉模型“它出现在什么位置以及与其他 Token 相隔多远”。

一个 Transformer Block 里有什么？

现代 Decoder-only 模型常见的 Pre-Norm Block 可以简化为：

x
│
├─────────────── Residual ───────────────┐
│                                        │
└→ Norm → Causal Self-Attention ─────────┼→ 相加
                                         ↓
                                        h
│
├─────────────── Residual ───────────────┐
│                                        │
└→ Norm → FFN / MLP ─────────────────────┼→ 相加
                                         ↓
                                      output

公式可以写成：

1 2	h = x + Attention(Norm(x)) output = h + MLP(Norm(h))

不同模型可能使用不同设计：

Pre-Norm 或 Post-Norm；
LayerNorm 或 RMSNorm；
GELU、SiLU、SwiGLU 等激活与门控结构；
标准 Attention、GQA、MQA 或局部 Attention；
是否使用 Bias；
Residual 的具体缩放方式。

因此，“Transformer Block”代表一类基本结构，不表示所有模型源码完全相同。

Self-Attention 是什么？

Self-Attention 让序列中的每个位置，根据当前任务动态关注同一序列中的其他相关位置。

考虑句子：

1	小明把书交给小红，因为她明天要考试。

当模型处理“她”时，需要结合上下文判断“她”更可能指向“小红”。Attention 提供了一种可学习的关系计算方式。

“Self”表示 Query、Key 和 Value 都来自同一个输入序列，而不是来自另一个编码器或外部数据库。

Q、K、V 分别是什么？

输入 Hidden States 记为 X。模型使用三个可训练的线性投影得到：

1
2
3

Q = XW_Q
K = XW_K
V = XW_V

其中 W_Q、W_K 和 W_V 是训练得到的参数。

可以用一个不完全等价但容易理解的比喻：

名称	直观问题	作用
Query	我当前要寻找什么？	表示当前位置的查询需求
Key	我可以通过什么特征被匹配？	表示每个位置用于匹配的索引
Value	匹配后实际取回什么？	表示每个位置提供的信息

假设当前正在更新“她”的表示：

“她”的 Query
    ↓
与“小明”“书”“小红”“因为”等位置的 Key 比较
    ↓
得到每个位置的注意力权重
    ↓
按照权重聚合各位置的 Value
    ↓
形成包含上下文信息的新表示

Q、K、V 不是人工编写的规则，也不直接等于自然语言中的问题、关键词和答案。它们是模型训练得到的高维向量。

Scaled Dot-Product Attention

Attention 的核心公式是：

1 2	Attention(Q, K, V) = softmax(QKᵀ / √d_head + Mask) V

可以分成五步理解。

第一步：计算相似度

1	scores = QKᵀ

每个 Query 与所有 Key 做点积，得到位置之间的匹配分数。

如果序列长度为 T，注意力分数矩阵通常具有：

[T, T]

第 i 行表示位置 i 对所有位置的关注分数。

第二步：除以缩放因子

1	scaled_scores = scores / √d_head

维度较大时，点积的绝对值可能变大，使 Softmax 进入非常陡峭的区域。除以 √d_head 可以帮助控制数值尺度，使训练更稳定。

第三步：加入 Mask

不允许关注的位置会被加上一个极大的负数，理论上可以理解为负无穷：

1	masked_score = -∞

经过 Softmax 后，对应概率接近 0。

第四步：转换为权重

1	weights = softmax(scaled_scores)

每一行的注意力权重之和为 1。

第五步：聚合 Value

1	output = weights × V

当前 Token 根据注意力权重，从其他可见位置的 Value 中聚合信息。

什么是 Causal Mask？

Decoder-only LLM 使用 Causal Attention，也叫 Masked Self-Attention。

训练模型预测下一个 Token 时，当前位置不能偷看未来答案。例如：

输入：我 喜欢 人工 智能

位置“我”只能看“我”
位置“喜欢”可以看“我 喜欢”
位置“人工”可以看“我 喜欢 人工”
位置“智能”可以看“我 喜欢 人工 智能”

对应的可见性矩阵为：

           Key 位置
           1   2   3   4
Query 1    ✓   ×   ×   ×
Query 2    ✓   ✓   ×   ×
Query 3    ✓   ✓   ✓   ×
Query 4    ✓   ✓   ✓   ✓

Mask 可以表示成：

0   -∞  -∞  -∞
0    0  -∞  -∞
0    0   0  -∞
0    0   0   0

它保证了训练阶段每个位置只能利用当前及之前的 Token，与真实生成时“未来 Token 尚不存在”的条件保持一致。

Padding Mask 与 Causal Mask 不相同

批量处理不同长度句子时，较短序列可能需要 Padding：

1 2	句子 A：[Token, Token, Token, Token] 句子 B：[Token, Token, PAD, PAD]

Padding Mask 用来阻止模型关注补齐位置；Causal Mask 用来阻止模型关注未来位置。

Mask	屏蔽什么	主要目的
Causal Mask	未来 Token	保持自回归因果性
Padding Mask	PAD Token	忽略批处理补齐内容

推理框架可能把多个 Mask 合并后交给 Attention Kernel。

为什么需要 Multi-Head Attention？

如果只有一组 Q、K、V，模型只能在一个表示子空间里计算关系。Multi-Head Attention 会并行使用多组投影：

Head 1 = Attention(Q₁, K₁, V₁)
Head 2 = Attention(Q₂, K₂, V₂)
...
Head H = Attention(Qₕ, Kₕ, Vₕ)

再把多个 Head 的结果拼接并投影：

1 2	MultiHead(Q, K, V) = Concat(head₁, head₂, ..., headₕ) W_O

不同 Head 可以学习不同的关系模式，例如：

临近 Token 关系；
指代关系；
语法结构；
实体与属性；
长距离依赖；
特定格式边界。

这些不是预先分配给每个 Head 的固定任务，而是训练过程中形成的表示模式。也不能简单把某个 Attention Head 的最高权重解释为模型完整的“思考原因”。

MHA、MQA 与 GQA

传统 Multi-Head Attention 中，每个 Query Head 通常拥有对应的 Key Head 和 Value Head：

1
2
3

Query Heads = 32
Key Heads   = 32
Value Heads = 32

在自回归生成中，所有历史 Key 和 Value 都需要写入 KV Cache。为了减少缓存和内存带宽，现代模型还会使用 MQA 或 GQA。

类型	Query Head 数	KV Head 数	特点
MHA	H	H	表达能力强，KV Cache 较大
MQA	H	1	多个 Query Head 共享一组 K/V，缓存最小
GQA	H	介于 1 和 H 之间	多组 Query Head 共享 K/V，兼顾质量和效率

例如：

1 2	28 个 Query Heads 4 个 KV Heads

意味着每 7 个 Query Heads 共享一组 K/V。GQA 论文提出的核心思想就是在 MHA 与 MQA 之间折中，使推理速度和质量取得更好的平衡。

Attention 之后为什么还需要 FFN？

Attention 主要负责不同 Token 位置之间的信息交换。每个位置聚合上下文后，还需要进行更复杂的非线性变换，这就是 Feed Forward Network，也经常叫 MLP。

原始 Transformer 的 FFN 可以简化为：

1	FFN(x) = W₂ · activation(W₁x + b₁) + b₂

典型的数据流是：

hidden_size
    ↓ 扩展
intermediate_size
    ↓ 激活与特征变换
hidden_size

例如：

1 2	hidden_size = 4096 intermediate_size = 11008

许多现代 LLM 使用带门控的结构，例如 SwiGLU：

1
2
3

MLP(x) = down_proj(
    SiLU(gate_proj(x)) * up_proj(x)
)

可以直观地理解为：

1 2	Attention：让 Token 之间交换信息 FFN / MLP：对每个位置的特征进行非线性加工

FFN 通常独立应用于每个 Token 位置，但所有位置共享同一组 FFN 参数。

Residual Connection 有什么作用？

Residual Connection，也叫残差连接或跳跃连接，会把子层输入直接加回输出：

1	output = x + sublayer(x)

它的作用包括：

保留原始信息通道；
让子层主要学习需要补充的变化；
改善深层网络的梯度传播；
使大量 Transformer Block 可以稳定堆叠。

不能把 Residual 理解为简单“复制一份文本”。它传递的是 Hidden State 张量。

LayerNorm 与 RMSNorm

深层神经网络中的数值分布会不断变化。Normalization 用来稳定各层输入尺度，改善训练过程。

LayerNorm

LayerNorm 会进行均值中心化和方差缩放，再应用可学习参数。

RMSNorm

RMSNorm 主要根据均方根缩放向量，不做与 LayerNorm 完全相同的均值中心化。许多现代 LLM 使用 RMSNorm，因为它结构更简单、计算高效，并且在实践中表现良好。

Pre-Norm 与 Post-Norm

原始 Transformer 常见 Post-Norm：

1	output = Norm(x + Sublayer(x))

现代 LLM 常见 Pre-Norm：

1	output = x + Sublayer(Norm(x))

Pre-Norm 往往更有利于训练很深的网络，但具体模型可能采用变体，阅读源码时应以实际结构为准。

为什么要堆叠很多层？

一个 Transformer Block 只能完成有限的表示变换。堆叠多层后，每层都在前一层 Hidden States 的基础上继续计算：

初始 Embedding
    ↓
Layer 1：建立基础局部和语法关系
    ↓
Layer 2：组合更复杂的上下文特征
    ↓
...
    ↓
Layer N：形成面向下一个 Token 预测的高层表示

这里的层级解释只是直观理解。真实模型不会严格规定“第几层只负责语法”或“第几层只负责事实”。能力通常分布在多个层、Head 和 MLP 参数中。

经过所有层后，Hidden States 的形状通常仍然是：

1	[batch_size, sequence_length, hidden_size]

序列长度和隐藏维度通常不变，但张量中的数值已经包含越来越丰富的上下文信息。

从 Hidden State 到 Logits

对于 Causal Language Model，最后还要把 Hidden State 映射回完整词表：

1	logits = hidden_states × W_vocab

LM Head 输出形状为：

1	[batch_size, sequence_length, vocab_size]

在生成下一个 Token 时，主要使用最后一个有效位置：

1	next_token_logits = logits[:, -1, :]

然后经过 Temperature、Top-k、Top-p 等处理，使用 Greedy 或 Sampling 选择下一个 Token。

一些模型会让 LM Head 与输入 Token Embedding 共享权重，称为 Weight Tying。它可以减少参数量，并让输入表示与输出词表映射共享结构；但并非所有模型都必须使用。

为什么现代文本 LLM 常用 Decoder-only？

原始 Transformer 为机器翻译设计，包含 Encoder 和 Decoder。后来形成了三种常见架构。

Encoder-only

代表模型：BERT 类模型。

输入 Token
   ↓
双向 Self-Attention
   ↓
每个位置融合左右两侧上下文
   ↓
分类、抽取、表示学习

它适合理解输入，但默认不是按顺序生成长文本的架构。

Decoder-only

代表模型：GPT、Llama、Qwen 等大量生成式 LLM。

已有 Token
   ↓
Causal Self-Attention
   ↓
预测下一个 Token
   ↓
循环生成

它把各种任务统一为文本续写：问答、翻译、摘要、代码生成都可以表示成“根据当前上下文继续生成”。

Encoder-Decoder

代表模型：原始 Transformer、T5 等。

Encoder 处理输入
        ↓
Decoder 通过 Cross-Attention 读取 Encoder 输出
        ↓
生成目标序列

它适合输入到输出的转换任务，例如翻译和摘要。

架构	Attention 特点	常见用途
Encoder-only	通常双向读取输入	分类、检索、Embedding、抽取
Decoder-only	Causal，只看当前及以前	对话、文本生成、代码生成
Encoder-Decoder	Encoder 双向，Decoder 因果并带 Cross-Attention	翻译、摘要、条件生成

Self-Attention 与 Cross-Attention

Self-Attention 的 Q、K、V 来自同一个序列：

1
2
3

Q = 输入序列
K = 输入序列
V = 输入序列

Cross-Attention 中，Query 通常来自 Decoder，Key 和 Value 来自 Encoder 输出：

1
2
3

Q = Decoder Hidden States
K = Encoder Hidden States
V = Encoder Hidden States

Decoder-only LLM 的标准 Block 主要使用 Causal Self-Attention，并不需要单独的 Encoder Cross-Attention 才能读取 Prompt，因为 Prompt 与输出前缀本来就在同一条 Token 序列中。

Transformer 与 KV Cache 的关系

在 Decoder 的每层 Self-Attention 中，历史 Token 都会产生 Key 和 Value：

Layer 1：K₁...Kₜ，V₁...Vₜ
Layer 2：K₁...Kₜ，V₁...Vₜ
...
Layer N：K₁...Kₜ，V₁...Vₜ

生成下一个 Token 时，过去位置的 K/V 不会因为未来 Token 出现而改变，所以可以缓存：

新 Token Hidden State
       ↓
只计算新 Token 在每层的 Q、K、V
       ↓
Q_new 与历史 K_cache 计算注意力
       ↓
聚合历史 V_cache
       ↓
把 K_new、V_new 追加到缓存

因此 KV Cache 不是 Transformer 之外的神秘记忆，而是对各层 Causal Self-Attention 中历史 K/V 张量的推理优化。

训练与推理时有什么不同？

训练阶段

训练数据中完整序列已知，可以一次并行计算所有位置：

1 2	输入：A B C D 标签：B C D E

Causal Mask 确保每个位置不能看到其标签对应的未来 Token，但 GPU 仍然可以并行计算多个位置的矩阵运算。

位置 A 预测 B
位置 B 预测 C
位置 C 预测 D
位置 D 预测 E

这种方式常称为 Teacher Forcing：训练时每个位置看到的历史通常来自真实训练序列，而不是模型刚刚生成的错误结果。

推理阶段

未来 Token 尚不存在，只能逐个生成：

1
2
3

A B C D → 生成 E
A B C D E → 生成 F
A B C D E F → 生成 G

KV Cache 用来避免每一步重复计算整个历史的 K/V。

Attention 的计算复杂度

标准全注意力需要构造 Token 两两之间的分数矩阵：

1	[sequence_length, sequence_length]

因此在序列长度维度上，Attention Score 的计算量通常近似：

O(n²)

当上下文从 4K 增加到 8K 时，仅从理论矩阵元素数量看：

1 2	4K × 4K → 约 1,678 万 8K × 8K → 约 6,711 万

序列长度翻倍，分数矩阵规模约变为四倍。这也是长上下文成本较高的原因之一。

需要区分：

FlashAttention 等算法可以减少中间显存读写并提高执行效率；
Sliding Window、Sparse Attention 等方法可以减少可见关系；
这些优化不等于所有场景都把标准全注意力数学关系自动变成 O(n)。

实战一：手写 Scaled Dot-Product Attention

下面使用 PyTorch 实现单头 Causal Self-Attention 的核心计算。

安装依赖

1	pip install -U torch

可执行代码

import math

import torch

torch.manual_seed(42)

batch_size = 1
sequence_length = 4
head_dim = 8

# 为了突出 Attention 核心，这里直接构造 Q、K、V。
query = torch.randn(batch_size, sequence_length, head_dim)
key = torch.randn(batch_size, sequence_length, head_dim)
value = torch.randn(batch_size, sequence_length, head_dim)

# [batch, query_position, key_position]
scores = query @ key.transpose(-2, -1)
scores = scores / math.sqrt(head_dim)

# 上三角表示未来位置，需要屏蔽。
future_mask = torch.triu(
    torch.ones(sequence_length, sequence_length, dtype=torch.bool),
    diagonal=1,
)

masked_scores = scores.masked_fill(future_mask, float("-inf"))
attention_weights = torch.softmax(masked_scores, dim=-1)
output = attention_weights @ value

print("Scores Shape：", scores.shape)
print("Attention Weights Shape：", attention_weights.shape)
print("Output Shape：", output.shape)
print("\nCausal Mask：")
print(future_mask)
print("\nAttention Weights：")
print(attention_weights.round(decimals=3))
print("\n每一行权重之和：")
print(attention_weights.sum(dim=-1))

观察输出时应该注意：

第一行只有第一个位置的权重非零；
第二行只能关注前两个位置；
被屏蔽的未来位置经过 Softmax 后权重为 0；
每一行允许位置的权重之和为 1；
最终 Output 是对应 Value 的加权和。

实战二：实现一个简化 Transformer Block

下面使用 PyTorch 内置 MultiheadAttention，构造一个 Pre-Norm Decoder Block。

import torch
from torch import nn


class CausalTransformerBlock(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        intermediate_size: int,
    ) -> None:
        super().__init__()

        self.attention_norm = nn.LayerNorm(hidden_size)
        self.attention = nn.MultiheadAttention(
            embed_dim=hidden_size,
            num_heads=num_heads,
            batch_first=True,
        )

        self.mlp_norm = nn.LayerNorm(hidden_size)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.GELU(),
            nn.Linear(intermediate_size, hidden_size),
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        sequence_length = hidden_states.shape[1]

        # True 表示这个 Query-Key 位置不允许参与 Attention。
        causal_mask = torch.triu(
            torch.ones(
                sequence_length,
                sequence_length,
                dtype=torch.bool,
                device=hidden_states.device,
            ),
            diagonal=1,
        )

        # 第一个 Pre-Norm + Attention + Residual。
        normed = self.attention_norm(hidden_states)
        attention_output, _ = self.attention(
            query=normed,
            key=normed,
            value=normed,
            attn_mask=causal_mask,
            need_weights=False,
        )
        hidden_states = hidden_states + attention_output

        # 第二个 Pre-Norm + MLP + Residual。
        hidden_states = hidden_states + self.mlp(
            self.mlp_norm(hidden_states)
        )

        return hidden_states


torch.manual_seed(42)

block = CausalTransformerBlock(
    hidden_size=64,
    num_heads=4,
    intermediate_size=256,
)

inputs = torch.randn(2, 10, 64)
outputs = block(inputs)

print("输入 Shape：", inputs.shape)
print("输出 Shape：", outputs.shape)
print("参数量：", sum(p.numel() for p in block.parameters()))

这个简化实现体现了核心结构，但不是生产级 LLM Block。真实模型还可能包含：

RMSNorm；
RoPE；
GQA；
SwiGLU；
Attention Dropout；
高性能 Attention Kernel；
KV Cache；
Tensor Parallel；
混合精度和量化。

实战三：查看真实模型的 Transformer 配置

下面读取 Qwen2.5-0.5B-Instruct 的配置，观察隐藏维度、层数、Attention Head 和 KV Head。

1	pip install -U torch transformers

from transformers import AutoConfig, AutoModelForCausalLM

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

config = AutoConfig.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

print("模型类型：", config.model_type)
print("词表大小：", config.vocab_size)
print("隐藏维度：", config.hidden_size)
print("Transformer 层数：", config.num_hidden_layers)
print("Attention Heads：", config.num_attention_heads)
print("KV Heads：", config.num_key_value_heads)
print("中间层维度：", config.intermediate_size)
print("最大位置配置：", config.max_position_embeddings)
print("总参数量：", f"{model.num_parameters():,}")

print("\n=== 第一层 Transformer Block ===")
print(model.model.layers[0])

输出的第一层结构中，可以重点寻找：

self_attn：Self-Attention；
q_proj、k_proj、v_proj：Q、K、V 投影；
o_proj：Attention 输出投影；
mlp：前馈网络；
gate_proj、up_proj、down_proj：门控 MLP；
input_layernorm：Attention 前的归一化；
post_attention_layernorm：MLP 前的归一化。

模型内部属性名称不是 Transformers 的跨模型统一接口，其他架构可能使用不同路径和命名。

实战四：观察 Hidden States、Attention 和 KV Cache

下面执行一次真实模型前向传播，并打印关键张量 Shape。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# eager 实现便于返回 Attention Weights，主要用于学习和观察。
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    attn_implementation="eager",
)
model.eval()

inputs = tokenizer(
    "Transformer 通过注意力机制理解上下文。",
    return_tensors="pt",
)

with torch.inference_mode():
    outputs = model(
        **inputs,
        output_hidden_states=True,
        output_attentions=True,
        use_cache=True,
    )

print("Input IDs：", inputs["input_ids"].shape)
print("Logits：", outputs.logits.shape)
print("Hidden States 数量：", len(outputs.hidden_states))
print("Embedding 输出：", outputs.hidden_states[0].shape)
print("最后一层 Hidden States：", outputs.hidden_states[-1].shape)
print("Attention 层数：", len(outputs.attentions))
print("第一层 Attention：", outputs.attentions[0].shape)

cache = outputs.past_key_values
legacy_cache = (
    cache.to_legacy_cache()
    if hasattr(cache, "to_legacy_cache")
    else cache
)

first_layer_key, first_layer_value = legacy_cache[0][:2]
print("第一层 Key Cache：", first_layer_key.shape)
print("第一层 Value Cache：", first_layer_value.shape)

# 下一个 Token 只使用最后一个有效位置的 Logits。
next_token_id = torch.argmax(outputs.logits[:, -1, :], dim=-1)
print("Greedy 下一个 Token：", tokenizer.decode(next_token_id))

常见的 Shape 关系为：

Input IDs
[batch, sequence]

Hidden States
[batch, sequence, hidden_size]

Attention Weights
[batch, attention_heads, query_sequence, key_sequence]

KV Cache
[batch, kv_heads, sequence, head_dim]

Logits
[batch, sequence, vocab_size]

这里可以把前几篇文章串起来：

Tokenizer 产生 Input IDs
        ↓
Transformer 产生 Hidden States
        ↓
每层 Attention 产生可缓存的 K/V
        ↓
LM Head 产生 Logits
        ↓
解码器选择下一个 Token

Attention 权重能解释模型吗？

Attention Weights 可以帮助观察某一层、某个 Head 在一次计算中如何分配权重，但不能直接等同于模型完整的推理解释。

原因包括：

模型有很多层和很多 Head；
Value 和输出投影会继续变换信息；
FFN 保存和处理大量特征；
Residual 会混合多条信息路径；
相同输出可能由多个内部路径共同形成；
最高 Attention 权重不等于对最终 Logit 的最大因果贡献。

因此，可视化 Attention 是分析工具之一，而不是“模型为什么这样回答”的完整证明。

Transformer 的知识存在哪里？

不能简单回答“都存在 Attention 里”。模型能力分布在：

Token Embedding；
Q/K/V/O 投影；
FFN / MLP 参数；
Normalization 参数；
LM Head；
多层组件之间的组合关系。

研究通常认为 FFN 与事实和特征表示有密切关系，但知识不是一个可以直接读取的普通数据库表。模型回答问题时，会通过多层 Attention 与 MLP 的共同计算激活相关表示。

常见误区

误区一：Transformer 就是 Attention

Attention 是核心组件，但完整 Transformer 还包括 Embedding、位置信息、FFN、Residual、Normalization 和输出映射。

误区二：Attention 是在搜索模型训练数据

Self-Attention 处理的是当前序列中的 Hidden States，不会自动搜索原始训练语料或互联网。

误区三：Attention Weight 就是概率正确率

它表示某个 Head 在特定层中的信息聚合权重，不是事实正确概率。

误区四：每个 Attention Head 都有固定的人类可读职责

Head 的功能通过训练形成，可能分布、重叠或依赖其他层，不能机械标记成“语法 Head”“事实 Head”。

误区五：Decoder-only 模型看不到 Prompt

它不能看未来 Token，但可以通过 Causal Attention 读取当前位置之前的 System、历史消息、用户问题和已生成内容。

误区六：训练时模型逐 Token 串行计算所有位置

训练中已知完整序列，可以利用 Causal Mask 并行计算多个位置。推理时未来 Token 未知，才必须逐步生成。

误区七：位置编码只是给 Token 编号

位置机制会进入 Attention 关系计算。RoPE 等方法还会显式影响 Query 与 Key 的相对位置关系。

误区八：层数越多，模型一定越好

模型能力还取决于隐藏维度、数据、训练方法、参数分配、计算预算和优化稳定性。单独增加层数不保证质量提升。

误区九：KV Cache 保存了完整 Hidden States

标准 KV Cache 主要保存每一层历史位置的 Key 和 Value，不等同于保存全部中间状态。

误区十：FlashAttention 改变了模型含义

在等价配置和允许的数值误差下，FlashAttention 主要优化 Attention 的计算与内存访问方式，不是把模型换成另一套语义规则。

如何阅读一个真实 LLM 架构？

拿到模型配置时，可以按下面顺序检查：

vocab_size：词表有多大；
hidden_size：每个 Token 的主表示维度；
num_hidden_layers：堆叠多少个 Block；
num_attention_heads：Query Head 数；
num_key_value_heads：是 MHA、GQA 还是 MQA；
head_dim：每个 Head 的维度；
intermediate_size：MLP 扩展维度；
hidden_act：MLP 激活函数；
max_position_embeddings 与 RoPE 配置：位置与上下文设置；
Norm 类型和位置：Pre-Norm、Post-Norm、RMSNorm 或 LayerNorm；
是否使用 Sliding Window、MoE 等特殊结构；
LM Head 是否与输入 Embedding 共享权重。

这些配置可以帮助我们判断参数主要花在哪里、KV Cache 大小如何估算，以及模型推理时可能有哪些性能特点。

小结

本文的核心内容可以总结为：

Transformer 让每个 Token 根据上下文动态更新自己的 Hidden State；
Token Embedding 提供初始表示，位置信息提供顺序和相对距离；
Self-Attention 的 Q、K、V 来自同一序列；
QKᵀ 计算匹配分数，Softmax 产生权重，再对 V 加权聚合；
Causal Mask 保证当前位置不能看到未来 Token；
Multi-Head Attention 允许模型在多个表示子空间中建模关系；
MQA 和 GQA 通过减少 KV Head 降低 KV Cache 与推理带宽；
FFN 对每个位置进行非线性特征变换；
Residual 和 Normalization 让深层网络更容易稳定训练；
多个 Block 堆叠后，LM Head 将最终 Hidden States 映射成词表 Logits；
Decoder-only、Encoder-only 和 Encoder-Decoder 适合不同任务；
KV Cache 本质上缓存的是各层 Causal Self-Attention 的历史 Key 和 Value。

现在，我们已经能够把一次 LLM 推理的核心路径完整串起来：

文本
→ Token
→ Embedding 与位置
→ Transformer Blocks
→ Hidden States
→ Logits
→ 解码
→ 下一个 Token

下一篇将进入训练阶段，介绍模型如何通过海量文本完成预训练，再通过 SFT 和偏好对齐逐步成为能够遵循指令的聊天助手。