Research map / Large language models / Reasoning and chain of thought

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

AAAI 2026 · Large language models › Reasoning and chain of thought · approach: Search and self-correction in reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes a Transformer language VAE (BERT prior/posterior encoders, Qwen2.5-0.5B decoder) that injects latent reasoning-rule vectors into decoder attention, and studies rule separation on math derivation, syllogistic and explanatory NLI-style tasks.

Problem

Explicitly encoding human-interpretable reasoning rules in the latent space of language models so that generation is rule-controlled rather than memorisation-based, especially for NLI-style tasks.

Approach

Uses a VAE whose posterior encoder averages latents of BM25-retrieved same-rule examples and whose prior encoder predicts the rule via a classifier; the latent is injected into the decoder's attention (kv_add, kv_mem, or query_add), with query injection giving the best retrieval of stored values, and supports the claim with an NTK-based argument about rule-specific gradient subspaces.

Key results

• On math reasoning, query_add with cls_weight=0.5 reaches 0.8130 BLEU and 0.6642 accuracy vs 0.6260 BLEU and 0.3800 accuracy for zero-shot Qwen2.5-0.5B.
• On syllogistic reasoning, query_add without prior reaches 0.9681 BLEU and 0.9200 accuracy, above the 0.9014/0.7000 zero-shot base.
• Few-shot prompting of the base Qwen2.5-0.5B drops sharply (e.g., math accuracy 0.0614), with the model repeating generated examples.

Surprising findings

• Few-shot prompting and few-shot fine-tuning markedly degrade performance because the model repeats generated examples.
• Performance on math reasoning declines when training data per operation exceeds about 2,000 samples.

Limitations

• Only lightweight decoders under 1B parameters are studied.
• Out-of-distribution math performance remains very low (best accuracy about 0.02).
• Results rely on single configurations and the NTK rule-separation claim is supported only by gradient heatmaps rather than formal proof.

Experimental setup

Datasets: Math derivation dataset (seven symbolic rules) Syllogistic reasoning dataset (24 valid patterns) EntailmentBank-based explanatory reasoning dataset

Benchmarks: Math Reasoning (in-distribution and OOD) Syllogistic Reasoning Explanatory Reasoning

Metrics: BLEU accuracy gradient cosine similarity heatmaps PCA/t-SNE visualisation

Baselines: GPT2-medium Llama3-1B Qwen2-0.5B Qwen2.5-0.5B few-shot

Builds on: Optimus Neural Tangent Kernel BERT Qwen2.5

Closest papers by meaning

• Efficient Post-Training Refinement of Latent Reasoning in Large Language Models (AAAI 2026 · similarity 0.89)
• Latent Concept Disentanglement in Transformer-based Language Models (ICLR 2026 · similarity 0.89)
• Reasoning Palette: Modulating Reasoning via Latent Contextualization for Controllable Exploration for (V)LMs (CVPR 2026 · similarity 0.88)
• I Have Covered All the Bases Here: Interpreting Reasoning Features in Large Language Models via Sparse Autoencoders (AAAI 2026 · similarity 0.88)
• Think in Latent, Explain in Language: Self-Explainable Latent Reasoning (ICML 2026 · similarity 0.88)
• Reasoning via Test-Time Instance-Level Policy Gradient in Latent Space (NeurIPS 2026 · similarity 0.88)