终章：从向量到语言——Transformer 输出生成层（Output Head）深度解析

在 Transformer 模型的流水线中，输入层负责将离散的符号映射为连续的几何向量，中间的编码器与解码器层在多维流形中完成了极其繁复的特征提炼与因果关联推演。然而，神经网络内部流动的所有特征向量，本质上都是高维连续空间中的浮点数阵列。

计算机需要向人类输出可读的自然语言，就必须将解码器最终吐出的连续隐层状态，重新映射回离散的字符或子词符号。承担这“最后一公里”跃迁任务的组件，正是 输出生成层（Output Head / Language Model Head）。

输出生成层不仅包含将隐藏维度映射到词表维度的线性投影，还融合了概率归一化、参数权重共享、训练损失优化以及决定生成内容多样性的采样解码策略。

一、 线性投影层（Linear Projection / LM Head）：维度的飞跃

解码器最后一层经过层归一化后，输出张量的形状为 $(B, L, d_{\text{model}})$，其中 $B$ 为批大小，$L$ 为序列长度，$d_{\text{model}}$ 为模型隐藏层维度。

线性投影层通过一个全连接矩阵，将特征维度从抽象的语义空间强制膨胀到物理词表空间：

解码器最终隐状态: H ∈ R^(B × L × d_model)
              │
              ▼
    [ 线性映射 W_v ∈ R^(|V| × d_model) + Bias ]
              │
              ▼
   未归一化对数概率: Logits ∈ R^(B × L × |V|)

1. 数学表述与 Logits 的物理意义

设词表大小为 $\vert{}V\vert{}$（例如现代大模型中通常在 $32,000 \sim 128,000$ 之间），线性投影层的权重矩阵为 $W_{\text{vocab}} \in \mathbb{R}^{\vert{}V\vert{} \times d_{\text{model}}}$，偏置项为 $b \in \mathbb{R}^{\vert{}V\vert{}}$。对于任意时刻的隐状态向量 $h \in \mathbb{R}^{d_{\text{model}}}$：

$$z = h W_{\text{vocab}}^T + b$$

输出的向量 $z \in \mathbb{R}^{\vert{}V\vert{}}$ 被称为 Logits（未归一化的对数几率）。$z$ 中的每一个分量 $z_i$，在几何上代表当前上下文表征 $h$ 与词表中第 $i$ 个词的特征原型向量之间的点积相似度：分量值越大，模型认为当前位置填入该词的合理性越高。

2. 计算与显存开销

随着分词粒度的演进，词表大小 $\vert{}V\vert{}$ 越来越大。计算 Logits 矩阵的矩阵乘法操作是长文本训练与推理阶段最消耗显存和计算带宽的算子之一。例如，当 Batch Size 为 4、序列长度为 4096、词表为 128,000 时，单一浮点批次的 Logits 矩阵即占用：

$$4 \times 4096 \times 128,000 \times 2 \text{ Bytes (FP16)} \approx 4.19 \text{ GB}$$

因此在工程落地中，主流框架通常采用融合交叉熵（Fused Cross-Entropy）算子，直接在 GPU 显存内将点积与损失函数一步算出，避免物化存储巨大的 Logits 中间矩阵。

二、 权重绑定机制（Weight Tying）：优雅的对称与参数瘦身

在原始论文《Attention Is All You Need》中，作者遵循了 Press & Wolf (2017) 等人提出的 Weight Tying（权重绑定 / 权重共享） 技术：将输入词嵌入层（Input Embedding）与最终输出线性层（LM Head）共享同一个权重矩阵。

$$W_{\text{vocab}} \equiv W_{\text{embed}}$$

1. 为什么两处权重可以完全共享？

词嵌入层与输出投影层在语义几何上互为对偶：

• 输入 Embedding： $x = \text{Lookup}(W_e, \text{ID})$，通过索引提取词向量，本质是单热编码向量（One-Hot）与词表矩阵相乘。
• 输出 Head： $z = h W_v^T$，通过隐状态与词表各向量做点积求相似度。

这两个过程针对的都是词表中静态的“概念词向量”。如果同一个词在输入端和输出端拥有完全不同的几何空间，不仅浪费参数容量，还会迫使网络耗费深度层去学习两套词表语义之间的对齐。强制二者共享同一套参数，能够促使模型学到更紧凑、泛化能力更强的双向词空间。

2. 参数量维度的巨大收益

假设模型 $d_{\text{model}} = 4096$，词表大小 $\vert{}V\vert{} = 128,000$：

• 单一词表矩阵参数量为：$128,000 \times 4096 \approx 5.24 \text{ 亿参数}$。
• 如果独立设置，两个矩阵将占用超过 10.5 亿参数；使用 Weight Tying，直接节省了约 5 亿显存常驻参数。

现代大模型的工程取舍：
在较小的模型（如 GPT-2、PaLM、Gemma）中，Weight Tying 被广泛使用；而在极大规模参数模型（如 LLaMA-2/3、DeepSeek）中，由于总参数量已达千亿级，5 亿参数占比极低，为了给予输出层更独立的判别拟合自由度，部分模型选择解除绑定，采用独立的 Head 矩阵。

三、 Softmax 归一化与温度系数（Temperature）

将实数范围的 Logits 转换为严格合法的概率分布（和为 1 且分量均在 $[0, 1]$ 之间），必须经过 Softmax 函数。为了在推理阶段动态调节生成的随机性，通常引入温度系数（Temperature, $T$）：

$$P(w_i) = \frac{\exp(z_i / T)}{\sum_{j=1}^{\vert{}V\vert{}} \exp(z_j / T)}$$

          高温 (T > 1)                     标准 (T = 1)                     低温 (T < 1)
   Logits 差异被缩小，分布平坦        维持原始相对概率差异          极化分布，高分更高，低分更低
       ┌───┬───┬───┐                     ┌─┐                           █
       │   │   │   │                     │ │ ┌─┐                       █  
       └───┴───┴───┘                     │ │ │ │ ┌─┐                   █  
     多样、富有创意但易胡说                平衡输出与连贯性              确定、精准但易陷入死板循环

• 标准温度（$T = 1$）： 保持模型原始训练得到的打分梯度。
• 降低温度（$T < 1$）： 将 Logits 整体放大，数值间的差距被指数级拉大。Softmax 输出呈现极度“尖锐（Sharp）”的分布，模型会极度聚焦于概率最高的前几个候选词。当 $T \to 0$ 时，输出退化为完全确定的贪婪搜索（Argmax）。常用于代码生成、数学推理等严谨场景。
• 升高温度（$T > 1$）： 将 Logits 整体缩小，数值间的差距被抹平。Softmax 输出趋向“平坦（Flat / Uniform）”，原本概率较低的词被选中的几率显著提升，模型更具想象力与创意，但过高会导致逻辑混乱、文本胡话。常用于文学创作、头脑风暴场景。

四、 训练目标：交叉熵损失与标签平滑

在有监督训练阶段，Transformer 的输出层使用标准的多分类交叉熵损失（Cross-Entropy Loss）指导优化：

$$\mathcal{L}_{\text{CE}} = - \sum_{t=1}^{L} \log P(y_t^* \mid y_{<t}^*, X_{\text{src}})$$

其中 $y_t^*$ 为当前位置的真实标注词（Ground Truth）。

标签平滑（Label Smoothing）

原生 Transformer 引入了正则化手段——标签平滑（通常设平滑系数 $\epsilon = 0.1$）：

• 传统的 One-Hot 标签：目标词概率为 1，其余全为 0。这会迫使模型无限拉大真实词与竞品词之间的 Logits 差值，容易导致模型在深层陷入极端过自信（Overconfidence），损害泛化能力。
• 软化后的目标概率：真实词的标签调整为 $1 - \epsilon + \frac{\epsilon}{\vert{}V\vert{}}$，其余词均分 $\frac{\epsilon}{\vert{}V\vert{}}$。
• 效果：抑制了模型过拟合，虽然在训练初期使困惑度（Perplexity）略微升高，但最终在 BLEU 分数等端到端生成指标上获得了稳定的提升。

五、 推理解码策略：如何从连续分布中采样文本？

在推理自回归阶段，Softmax 输出的是一个长度为 $\vert{}V\vert{}$ 的离散概率向量。如何将概率分布具象化为一个确切的 Token，业界演变出了一系列解码控制技术：

在实际工程使用中，通常以 Temperature + Top-P + Repetition Penalty 的组合形式协同工作：

• 重复惩罚（Repetition Penalty）： 对前文中已经出现过的 Token，人为对其 Logits 进行除法或减法折损，强力压制自回归模型常见的“死循环车轱辘话”现象。

六、 PyTorch 完整实现：LM Head 与动态采样器

以下代码展示了带权重绑定的输出生成层封装，以及融合了 Temperature、Top-K、Top-P 联合控制的端到端单步采样实现：

import torch
import torch.nn as nn
import torch.nn.functional as F

class TransformerLMHead(nn.Module):
    """Transformer 输出生成层模块"""
    def __init__(self, d_model: int, vocab_size: int, shared_embedding: nn.Embedding = None):
        super().__init__()
        self.vocab_size = vocab_size
        self.linear = nn.Linear(d_model, vocab_size, bias=False)

        # 实施 Weight Tying
        if shared_embedding is not None:
            self.linear.weight = shared_embedding.weight

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        hidden_states: (batch_size, seq_len, d_model)
        return: logits (batch_size, seq_len, vocab_size)
        """
        return self.linear(hidden_states)


def sample_next_token(
    logits: torch.Tensor,
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 0.0,
    repetition_penalty: float = 1.0,
    generated_tokens: torch.Tensor = None
) -> torch.Tensor:
    """
    单步采样器：支持 Temperature, Top-K, Top-P, Repetition Penalty
    logits: (batch_size, vocab_size) -> 仅取最后一个时间步的输出
    """
    batch_size, vocab_size = logits.shape

    # 1. 重复惩罚处理
    if repetition_penalty != 1.0 and generated_tokens is not None:
        for b in range(batch_size):
            unique_tokens = torch.unique(generated_tokens[b])
            for token_id in unique_tokens:
                # 若 logit > 0 则除以 penalty 降低它；若 logit < 0 则乘以 penalty 进一步压低它
                if logits[b, token_id] > 0:
                    logits[b, token_id] /= repetition_penalty
                else:
                    logits[b, token_id] *= repetition_penalty

    # 2. 温度平滑
    if temperature > 0:
        logits = logits / temperature
    else:
        # 当 temperature 趋近于 0，退化为贪婪选择
        return torch.argmax(logits, dim=-1, keepdim=True)

    # 3. Top-K 过滤
    if top_k > 0:
        top_k = min(top_k, vocab_size)
        # 寻找第 K 个最大值，低于该值的候选全部置为 -inf
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = -float("Inf")

    # 4. Top-P (核采样) 过滤
    if 0.0 < top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # 移除累积概率超过 top_p 的多余元素
        sorted_indices_to_remove = cumulative_probs > top_p
        # 保证至少保留第一个概率最高的候选词 (向右平移一位)
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = False

        # 映射回原始 index 维度并将对应位置置为 -inf
        for b in range(batch_size):
            indices_to_remove = sorted_indices[b][sorted_indices_to_remove[b]]
            logits[b, indices_to_remove] = -float("Inf")

    # 5. 归一化并进行多项式随机抽样
    probs = F.softmax(logits, dim=-1)
    next_token = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)

    return next_token


# ---------------- 验证前向传播与采样流水线 ----------------
if __name__ == "__main__":
    batch_size = 2
    d_model = 512
    vocab_size = 1000

    # 1. 构建词嵌入与输出头 (启用权重绑定)
    token_embedding = nn.Embedding(vocab_size, d_model)
    lm_head = TransformerLMHead(d_model=d_model, vocab_size=vocab_size, shared_embedding=token_embedding)

    # 2. 模拟解码器末层的输出隐状态
    mock_hidden_states = torch.randn(batch_size, 1, d_model)

    # 3. 计算未归一化对数概率 Logits
    logits = lm_head(mock_hidden_states).squeeze(1)  # (batch_size, vocab_size)
    print("LM Head 产出 Logits 尺寸:", logits.shape)

    # 4. 模拟已生成 token 并进行多策略采样
    mock_history = torch.tensor([[10, 24, 88], [42, 99, 10]])
    next_token = sample_next_token(
        logits=logits,
        temperature=0.8,
        top_k=50,
        top_p=0.9,
        repetition_penalty=1.1,
        generated_tokens=mock_history
    )

    print("最终预测生成的下一个 Token ID:", next_token.squeeze(-1).tolist())

输出生成层处于整个 Transformer 架构收敛信息与吐出预测的最末端。它既是数学优化落地的交汇点，也是连续多维世界通往人类离散词句的桥梁。从底层的线性映射和权重绑定，到损失函数中的标签平滑，再到推理时基于温度与核采样的动态裁决，输出生成层的每一处设计，都深刻决定着最终生成文本的质量、准确度与自然灵性。