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

Citation

• Authors: Yingji Zhang, Marco Valentino, Danilo S. Carvalho, and Andre Freitas
• Year: 2025
• Venue: arXiv preprint
• arXiv: arXiv:2506.19418
• URL: https://arxiv.org/abs/2506.19418
• Citation count: 0 (deeplearn.org paper detail page, retrieved 2026-04-20)

One-Line Takeaway

This paper studies whether explicitly labeled reasoning rules can be encoded and separated inside a language-VAE latent space, and argues that with direct supervision they can.

Technical Core

• The paper treats a reasoning rule as a functional input-output mapping and asks whether those mappings can be separated in a VAE latent space.
• The architecture distinguishes:
  • a prior-side representation that tracks fine-grained input semantics;
  • a posterior-side representation that captures input-output reasoning patterns.
• The supporting theory uses a Neural Tangent Kernel argument: if rule identity is supervised during training, then different rules can become separable in the encoder's parameter space and therefore cluster in the learned latent feature space.
• The decoder receives the latent through several alternative attention-side injection strategies, so the latent can influence generation or inference behavior explicitly.

Method

• Study three rule-based reasoning task families: mathematical reasoning, syllogistic reasoning, and explanatory reasoning.
• Train lightweight Transformer-based language VAEs with full decoder retraining, so the experiments can inspect parameter-space and latent-space geometry directly.
• Compare three latent injection strategies, including key/value-style injection and query-side injection.
• Measure rule separation in two places:
  • parameter space, using gradient-based similarity analyses;
  • latent feature space, using clustering / visualization diagnostics and downstream reasoning performance.

Empirical Evaluation

• With explicit rule supervision, the model learns latent spaces where different reasoning rules cluster distinctly.
• Query-side latent injection is the strongest decoding configuration in the paper's experiments.
• In the mathematical-reasoning case study, increasing samples per operation eventually stops helping, suggesting a performance bottleneck unrelated to brute-force sampling alone.
• The paper also reports that FFN layers preserve rule separation more effectively than attention layers in the studied setup.

Assumptions And Limits

The paper is about supervised rule encoding, not about discovering latent strategies from an already trained model. That distinction matters. The model is taught which rule family a sample belongs to, and the paper mostly studies memorized or explicitly specified input-output rule patterns rather than open-ended reasoning generalization. It also retrains small decoders instead of factorizing a frozen pretrained generator.

Relevance To This Project

This is the closest paper in the batch to the repo's central question, because it directly asks whether reasoning structure can be disentangled in a language-VAE latent space.

The overlap is real but not complete:

• both the paper and the repo care about latent reasoning structure, controllability, and interpretability;
• this paper uses explicit rule supervision and small controlled NLI-style tasks;
• the repo's harder target is to recover useful latent strategies from a pretrained model distribution without assuming clean rule labels.

Comparison Notes

LlaMaVAE is the direct nearby latent-language-model ancestor, but it focuses more on semantic control than on explicit reasoning-rule separation. Latent Thought Models also place continuous latent states inside generation, but they aim at efficiency and inference-time computation rather than supervised rule disentanglement. This paper is also a useful external comparison point for the repo's note Abstract Strategy vs Solution, because it operationalizes something closer to labeled rule classes than to reusable abstract strategy.

Related Notes

• LlaMaVAE: Guiding Large Language Model Generation via Continuous Latent Sentence Spaces
• Latent Thought Models with Variational Bayes Inference-Time Computation
• CVAE problem setup
• Abstract Strategy vs Solution