为什么语言模型会共享输入 Embedding 与输出权重？

简化版

Weight Tying 让输入 token Embedding 矩阵 E∈R^{V×H} 与输出 LM Head 使用同一组参数，输出 logits 常写成 hEᵀ。这样节省约 V×H 参数，并让“读入 token 的语义空间”与“预测 token 的分类空间”对齐。它要求词表和隐藏维兼容；在输入输出词表不同、多模态投影或量化/并行约束下未必适合，是否使用要比较困惑度、容量和真实部署收益。

详细版

未共享时，输入查表 E_in 和输出分类器 W_out∈R^{V×H} 各占一份大矩阵；共享后令 W_out=E_in。若 V=50,000、H=4096、bf16，单矩阵约 50,000×4096×2≈391 MiB，共享可少一份权重及其训练状态。

input token id -> E[token] -> Transformer -> h
logits = h @ E.T (+ output bias)

共享产生双向梯度：E 同时被输入表示与输出预测更新，属于正则化也可能限制自由度。实现要确认参数对象真正 alias，而非初始化时复制；保存、优化器参数组、分片与量化也要避免重复。扩词表时同一矩阵的新增行同时影响输入与输出。

完整版教学

一、输入与输出为何都有词表矩阵

输入 Embedding 把离散 token ID 映射为 H 维向量；输出头把 H 维隐藏状态映射回 V 个词表 logits。两者形状转置相容。

E_in:  V × H
W_out: V × H, logits = h W_out^T

既然都描述 token 与隐藏空间的关系，共享参数是自然假设，但不是数学必需。

记忆钩子：输入矩阵负责“这个词怎样读进来”，输出矩阵负责“当前状态像哪个词”；Weight Tying 让两者用同一本词典。

二、参数节省有多大

V=128K、H=4096 时矩阵含约 524M 参数。bf16 权重约 1 GiB；训练还可能有梯度与 Adam fp32 状态，未共享的额外成本远大于 1 GiB。

状态	每参数典型字节	额外矩阵约占用
bf16 权重	2	1.0 GiB
bf16 梯度	2	1.0 GiB
Adam m/v fp32	8	4.0 GiB

实际取决于混合精度、ZeRO 分片和优化器。词表越大，共享价值越明显。

三、共享如何改变梯度

未共享时输入矩阵只从出现于上下文的位置收梯度，输出矩阵从每个预测位置的 softmax 收梯度。共享后两个来源累加到同一 E。

这使 token 表示同时适合编码与分类，常有正则化效果；也减少模型为输入和输出分别学习完全不同空间的自由度。某些任务或架构中，解除共享可能换取少量质量。

需要检查梯度不是被错误加两次，而是自动微分对同一 Parameter 正常累积。

四、为什么可能改善泛化

输出分类器要判断隐藏状态与各 token 的兼容性，共享相当于用输入语义向量作为分类原型，减少参数并施加结构先验。

低频 token 的输入与输出统计可以互相帮助，但也可能把两种不对称需求绑在一起。收益是经验性的，应通过等训练条件的困惑度和下游评测确认。

不能把参数减少后的收益简单归因于“语义更统一”，还可能来自正则化。

五、哪些场景无法直接共享

编码器—解码器可能有不同源/目标词表；多模态输入 token 与文本输出词表不同；若输出前投影维度不等于 H，也需额外变换。

场景	主要障碍	替代方案
不同语言独立词表	V 不同	部分共享或不共享
输出低维投影	维度不同	加适配投影
多模态输入	token 语义不同	各模态独立输入投影
分类任务	输出非词表标签	独立分类头

共享要从任务语义与张量形状同时判断。

六、实现中最容易出现“假共享”

W_out.data.copy_(E.data) 只复制初值，后续仍是两个参数；真正共享需让模块引用同一 Parameter。可比较对象 identity、存储指针和修改后是否同步。

优化器参数列表里同一参数不应重复登记，否则可能被更新两次。序列化加载时，有的格式会打破 alias，需要模型初始化后重新 tie。

lm_head.weight = embedding.weight  # alias，而不是 copy

具体框架和编译器可能有额外要求，测试比代码外观更可靠。

七、扩词表与微调的影响

新增 token 时扩展 E 的行，输入查表和输出分类器同步拥有新 token。新增行需合理初始化，并用包含新 token 的输入与目标训练。

若只想训练新增 token，冻结共享矩阵无法只冻结旧行而放开新行，通常需要梯度 mask、参数拆分或适配器。LoRA 对输入 Embedding 与输出头如何挂载也因共享而互相影响。

删除或重排 tokenizer ID 则会同时破坏读入和输出映射，必须迁移权重行。

八、量化、并行与部署

输入查表和输出 GEMM 的访问模式不同，共享一个量化权重未必满足两个 kernel 的布局。推理框架可能为性能生成两份打包副本，逻辑共享但显存未完全节省。

Tensor Parallel 常按词表维切 LM Head，Embedding 也要兼容同一分片与通信。导出后应检查文件大小、运行时显存和 logits 一致性。

理论少一个参数矩阵不保证加载后只有一份物理存储。

九、常见误区与追问

误区：初始化相同就是 Weight Tying。 必须持续引用同一参数存储。
误区：共享一定提高精度。 它是结构先验与正则化，效果需实验。
误区：只节省推理权重。 训练还可节省梯度和优化器状态。
误区：所有 Encoder-Decoder 都能完全共享。 源目标词表或维度可能不同。
误区：导出后一定保持物理共享。 量化打包和序列化可能复制。
追问：输出 bias 是否也共享？ 输入 Embedding 没有对应 bias，通常独立或省略。
追问：扩词表有什么额外影响？ 同一新增行同时承担输入和输出，需要双向训练覆盖。
追问：如何验证实现真的 tie？ 查 Parameter identity/存储指针、优化器去重和保存加载后 alias。

十、加强记忆

Weight Tying 记住“V×H 同形、输入输出共一本词典”：令 LM Head 权重等于 Embedding，可省一整块词表矩阵及训练状态，并让编码与预测空间相互约束。它要求词表与维度兼容，实现必须是参数 alias 而非复制；扩词表、LoRA、分片、量化和导出都要重新验证逻辑与物理共享，最终用困惑度和真实显存判断收益。

简化版

详细版