重新思考权重绑定:伪逆绑定用于 LM 稳定训练与更新
摘要
权重绑定广泛用于紧凑型语言模型中,通过在输入嵌入和输出投影之间共享词表来减少参数。然而,仅凭参数共享并不保证词汇界面稳定:在训练过程中,编码令牌到隐藏状态之间以及解码隐藏状态到逻辑 logits 之间的对应关系可能会漂移,恶化优化灵敏度并削弱依赖于有意义的词汇空间解码器的可解释性探针。我们提出了伪逆绑定 (PIT),通过同步嵌入和解嵌入作为共享潜在词汇记忆的耦合投影来实现这一点,确保在整个训练期间保持伪逆一致的界面。PIT maintains an orthonormal shared memory, obtained by polar initialization from a source checkpoint for continued pretraining or by random orthonormal initialization for from-scratch pretraining, and introduces a learned symmetric positive definite hidden-space transform parameterized via a Cholesky factor. The output head applies this transform to hidden states before the vocabulary projection, while the embedding applies the inverse transform to token vectors using stable triangular solves, avoiding explicit pseudo-inverse recomputation and vocabulary-sized auxiliary parameters. Beyond improving training stability, PIT provides a cleaner substrate for logit-lens-style and vocabulary-space explainability probes by keeping the input and output token geometries synchronized. 我们评估了 PIT 在 256M-1.3B 参数范围内的 on-device 模型上的效果。结果表明,PIT 改进了持续预训练的稳定性,在各种设置下实现了接近精确的词汇界面一致性,并在持续预训练后yield更可预测的轻量级适应;而从头开始的预训练则揭示了严格界面一致性与无约束优化之间的权衡。
引用
@article{arxiv.2602.04556,
title = {Rethinking Weight Tying: Pseudo-Inverse Tying for LM Stable Training and Updates},
author = {Jian Gu and Aldeida Aleti and Chunyang Chen and Hongyu Zhang},
journal= {arXiv preprint arXiv:2602.04556},
year = {2026}
}
备注
an early-stage version