# Learning to Disentangle Latent Reasoning Rules with Language VAEs: A Systematic Study

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

## 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 21 of its 47 references that are in the map.
- First posted on arXiv 24 Jun 2025.

## 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, 21 of 21)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [MMaDA: Multimodal Large Diffusion Language Models](https://mlsanity.com/paper/2505.15809/) (2025)
- [In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering](https://mlsanity.com/paper/2311.06668/) (2023)
- [Explaining Answers with Entailment Trees](https://mlsanity.com/paper/2104.08661/) (2021)
- [Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI](https://mlsanity.com/paper/2402.00809/) (2024)
- [How much do language models memorize?](https://mlsanity.com/paper/2505.24832/) (2025)
- [WorldTree: A Corpus of Explanation Graphs for Elementary Science Questions supporting Multi-Hop Inference](https://mlsanity.com/paper/1802.03052/) (2018)
- [Transformer-based Conditional Variational Autoencoder for Controllable Story Generation](https://mlsanity.com/paper/2101.00828/) (2021)
- [Understanding Transformer from the Perspective of Associative Memory](https://mlsanity.com/paper/2505.19488/) (2025)
- [Searching Latent Program Spaces](https://mlsanity.com/paper/2411.08706/) (2024)
- [Mitigating Content Effects on Reasoning in Language Models through Fine-Grained Activation Steering](https://mlsanity.com/paper/2505.12189/) (2025)
- [A Symbolic Framework for Evaluating Mathematical Reasoning and Generalisation with Transformers](https://mlsanity.com/paper/2305.12563/) (2023)
- [Neurosymbolic Diffusion Models](https://mlsanity.com/paper/2505.13138/) (2025)
- [Do Large Language Models Truly Grasp Addition? A Rule-Focused Diagnostic Using Two-Integer Arithmetic](https://mlsanity.com/paper/2504.05262/) (2025)
- [Multi-Operational Mathematical Derivations in Latent Space](https://mlsanity.com/paper/2311.01230/) (2023)
- [Grammar-based Ordinary Differential Equation Discovery](https://mlsanity.com/paper/2504.02630/) (2025)
- [Towards Controllable Natural Language Inference through Lexical Inference Types](https://mlsanity.com/paper/2308.03581/) (2023)

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

- [The Knowledge-Reasoning Dissociation: Fundamental Limitations of LLMs in Clinical Natural Language Inference](https://mlsanity.com/paper/2508.10777/) (2025)
- [Bridging Compositional and Distributional Semantics: A Survey on Latent Semantic Geometry via AutoEncoder](https://mlsanity.com/paper/2506.20083/) (2025)

## Abstract (arXiv)

Incorporating explicit reasoning rules within the latent space of language models (LMs) offers a promising pathway to enhance generalisation, interpretability, and controllability. While current Transformer-based language models have shown strong performance on Natural Language Inference (NLI) tasks, they often rely on memorisation rather than rule-based inference. This work investigates how reasoning rules can be explicitly embedded and memorised within the LMs through Language Variational Autoencoders (VAEs). We propose a complete pipeline for learning reasoning rules within Transformer-based language VAEs. This pipeline encompasses three rule-based reasoning tasks, a supporting theoretical framework, and a practical end-to-end architecture. The experiment illustrates the following findings: Disentangled reasoning: Under explicit signal supervision, reasoning rules - viewed as functional mappings - can be disentangled within the encoder's parametric space. This separation results in distinct clustering of rules in the output feature space. Prior knowledge injection: injecting reasoning information into the Query enables the model to more effectively retrieve the stored value Value from memory based on Key. This approach offers a simple method for integrating prior knowledge into decoder-only language models. Performance bottleneck: In mathematical reasoning tasks using Qwen2.5(0.5B), increasing sample count doesn't improve performance beyond a point. Moreover, ffn layers are better than attention layers at preserving the separation of reasoning rules in the model's parameters.