# Rethinking Weight Tying: Pseudo-Inverse Tying for LM Stable Training and Updates

arXiv: https://arxiv.org/abs/2602.04556 · HTML: https://mlsanity.com/paper/2602.04556/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 16 of its 25 references that are in the map.
- First posted on arXiv 4 Feb 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 16 of 16)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Mass-Editing Memory in a Transformer](https://mlsanity.com/paper/2210.07229/) (2022)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster](https://mlsanity.com/paper/2304.03208/) (2023)
- [Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling](https://mlsanity.com/paper/1611.01462/) (2016)
- [Vocabulary-Defined Semantics: Latent Space Clustering for Improving In-Context Learning](https://mlsanity.com/paper/2401.16184/) (2024)
- [A Semantic-Aware Layer-Freezing Approach to Computation-Efficient Fine-Tuning of Language Models](https://mlsanity.com/paper/2406.11753/) (2024)

## Built on by (most cited first, 2 of 2)

- [SemRF: A Semantic Reference Frame for Residual-Stream Dynamics in Language Models](https://mlsanity.com/paper/2606.32022/) (2026)
- [Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State](https://mlsanity.com/paper/2609.33497/) (2026)

## Abstract (arXiv)

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.