Attention 与 Transformer：让上下文彼此说话

Attention 解决什么问题

Embedding 只表示当前 token。理解“猫坐在窗边，因为它喜欢阳光”，模型还要让“它”读取“猫”的信息。Attention 为每个位置动态计算：我现在应该从哪些上下文位置取多少信息？

INTERACTIVE FIGURE谁在关注谁？Attention 热力图

100%

选择查询 token，观察它对之前 token 的示意权重。因果遮罩意味着它不能看见未来；数值为教学构造，不是模型内部日志。

Query、Key、Value 的直觉

把每个 token 想成一张资料卡：

Query（Q）：当前任务在寻找什么。
Key（K）：这张卡适合被怎样检索。
Value（V）：如果命中，真正取走的内容。

模型对隐藏状态分别做线性投影得到 $Q,K,V$ ，用 Query 与所有可见 Key 的相似度决定权重，再对 Value 加权：

\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V

$d_k$ 是每个 head 的 key 维度；除以 $\sqrt{d_k}$ 避免内积随维度过大； $M$ 是 mask，未来位置被加上极小值，从而得到接近 0 的权重。直观上：相关性负责选，Value 负责搬。

多头为什么有用

一个Attention head 只能形成一种检索投影。多头让模型同时追踪局部搭配、长距离依赖、代码括号、实体指代等关系。不要把 head 说成固定的“语法专家”；这种解释最多是分析结果，不是结构承诺。

Grok-1 的 Grouped-Query Attention

标准多头 Attention 可以给每个 Query head 独立 K/V head。Grok-1 配置 48 个 Query heads、8 个 Key/Value heads，每个 K/V 组被 6 个 Query heads 共享。这是Grouped-Query Attention（GQA）。

为什么要共享 K/V？自回归生成会缓存历史 K/V。48 组改为 8 组，Grok-1 每层 KV Cache 的 head 维度减少到六分之一；代价是不同 Q heads 不再拥有完全独立的 K/V 表示。

Transformer 层怎样组合

Attention 负责跨 token 通信，前馈网络负责每个 token 独立加工。残差连接把加工结果加回主干，RMSNorm 稳定数值。Grok-1 的前馈位置不是一个稠密 MLP，而是 Router 选择两个 Expert 的 MoE。

INTERACTIVE FIGURE一次前向传播：从 token 到下一个 token

100%

INPUTtoken IDs[71, 2034, …]

→

LOOKUPEmbedding离散 → 连续

→

× 64 layers

RMSNormGrouped-query AttentionResidual

RMSNormRouter → Top-2 ExpertsResidual

→

DECODELM head131,072 logits

→

SAMPLEnext token一次只生成一个

每一层都先让 token 交换上下文信息（Attention），再各自加工（Grok-1 中是 MoE）。64 层之后投影成词表分数并采样。

在 DecoderLayer.__call__() 中可以追踪两段主干：

# 简化代码：保留真实结构，省略分片、dropout 和类型细节
h = h + attention(rms_norm(h), mask, kv_memory)
h = h + moe(rms_norm(h), padding_mask)

这段是简化代码，不是仓库原文。真实实现位于 model.py 的 MHABlock、DecoderLayer 与 MoELayer；源码还处理 bfloat16、JAX mesh、sharding constraint、attention logit soft-cap 与 KVMemory。

一个容易忽略的 Grok-1 细节

公开实现对 Attention logits 使用 tanh soft-cap：先乘 attn_output_multiplier，再把极端值平滑压到约 ±30，最后 softmax。它控制过大的注意力分数，改善数值行为。源码还明确让 softmax 在 fp32 执行，再转回输入 dtype。

优势、代价与限制

全局 Attention 能直接连接远距离 token，但随序列长度增长，预填充阶段的注意力矩阵代价近似 $O(n^2)$ 。
GQA 显著减小 KV Cache，却可能牺牲部分 K/V 表达自由度。
残差和归一化帮助 64 层信号稳定流动，但深网络依然依赖精细初始化、训练和并行策略。
看到 GQA、RoPE、RMSNorm 等 Grok-1 设计，不能推出 Grok 4.5 使用同一配置；它们只是一扇理解现代解码器模型的公开窗口。

本章检查点

Attention 不是“找最相似的词”这么简单，而是让当前表示生成 Query，在可见历史中对 Key 打分，再把 Value 混合回来。Transformer 把这种通信与每 token 加工重复很多层。