提出PIT方法,让语言模型训练更稳定,词表映射始终一致。
Rethinking Weight Tying: Pseudo-Inverse Tying for LM Stable Training and Updates
- 用共享潜空间和正交初始化,实现编码解码同步
- 在256M-1.3B模型上提升持续预训练稳定性
- 适合追求训练稳定性和可解释性的轻量级模型开发者
权重共享广泛用于紧凑型语言模型以减少参数,但仅靠共享无法保证稳定的词元接口:训练中编码与解码的对应关系可能漂移,加剧优化敏感性并削弱依赖有意义词表空间解码的可解释性探测。本文提出伪逆绑定(Pseudo-Inverse Tying, PIT),将嵌入与反嵌入作为共享潜空间的耦合投影,确保训练全程伪逆一致性。PIT维护一个正交共享记忆,通过源检查点极化初始化(用于持续预训练)或随机正交初始化(从头预训练),引入由Cholesky分解参数化的对称正定隐空间变换。输出头在词汇投影前对隐藏状态应用该变换,嵌入层则使用稳定三角求解计算逆变换,避免显式伪逆重算及词表规模的辅助参数。实验表明,PIT在256M至1.3B参数的设备端模型上提升了持续预训练稳定性,实现了跨场景近似精确的词元接口一致性,并带来更可预测的轻量化微调表现;从头训练时则揭示了严格接口一致性与自由优化之间的权衡。
原文摘要 · Abstract (English)
Weight tying is widely used in compact language models to reduce parameters by sharing the token table between the input embedding and the output projection. However, parameter sharing alone does not guarantee a stable token interface: during training, the correspondence between encoding tokens into hidden states and decoding hidden states into logits can drift, worsening optimization sensitivity and weakening explainability probes that rely on a meaningful vocabulary-space decoder. We propose Pseudo-Inverse Tying (PIT), which synchronizes embedding and unembedding as coupled projections of a shared latent token memory, guaranteeing a pseudo-inverse-consistent interface throughout training. 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. We evaluate PIT on on-device models spanning 256M-1.3B parameters. The results show that PIT improves continued-pretraining stability, enforces near-exact token-interface consistency across settings, and yields more predictable lightweight adaptation after continued pretraining, while from-scratch pretraining reveals a trade-off between strict interface consistency and unconstrained optimization.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。