Research map / Probabilistic and scientific modeling / Molecules and proteins
Learning Iterative Reasoning through Energy Diffusion
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 IRED, which learns a sequence of annealed energy landscapes with denoising score and contrastive energy supervision so reasoning becomes iterative energy minimization that scales inference steps with difficulty.
Problem
Learning-to-reason models often fail to generalize to harder instances; existing energy-based reasoning models are slow and unstable to train because they backpropagate through the full optimization, and their energy landscapes have many local optima.
Approach
Learn K annealed energy functions trained by denoising score matching (gradient of energy supervised toward injected noise) plus a contrastive loss against noise-corrupted negatives, then solve by gradient descent through landscapes from high to low noise; no backprop through the optimization is needed and inference steps can be increased for harder problems.
Key results
- On Sudoku harder dataset, IRED reaches 62.1% accuracy vs 28.6% for RRN and 3.2% for SAT-Net.
- On connectivity harder test (18 nodes), IRED achieves 93.8% vs 89.8% for IREM and 61.3% for diffusion.
- On shortest path planning harder set (graphs of size 25), IRED achieves 91.9% vs 88.4% for IREM and 46.9% for diffusion.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase optimization steps from 10 to 40 (ablation) | IRED · Matrix Inverse, harder difficulty | MSE: 0.2110 -> 0.2063 | Helped |
| increase optimization steps from 10 to 40 (ablation) | IRED · Matrix Inverse, same difficulty | MSE: 0.0096 -> 0.0095 | Helped |
| use gradient descent instead of the noisy diffusion reverse process (ablation) | IRED · Matrix Inverse, same difficulty | MSE: 0.0158 -> 0.0097 | Helped |
| use gradient descent instead of the noisy diffusion reverse process (ablation) | IRED · Matrix Inverse, harder difficulty | MSE: 0.2223 -> 0.2135 | Helped |
| add multiple optimization steps at each energy landscape (ablation) | IRED · Matrix Inverse, harder difficulty | MSE: 0.2135 -> 0.2113 | Helped |
| add contrastive energy shaping (ablation) | IRED · Matrix Inverse, harder difficulty | MSE: 0.2113 -> 0.2063 | Helped |
| add gradient descent optimization (ablation) | IRED · Sudoku harder dataset | accuracy: 35.0% -> 45.1% | Helped |
| add multiple optimization steps at each energy landscape (ablation) | IRED · Sudoku harder dataset | accuracy: 45.1% -> 59.7% | Helped |
| add contrastive energy shaping (ablation) | IRED · Sudoku harder dataset | accuracy: 59.7% -> 62.1% | Helped |
| add multiple optimization steps at each energy landscape (ablation) | IRED · Shortest Path test dataset | success rate: 80.7% -> 88.6% | Helped |
| add contrastive energy shaping (ablation) | IRED · Shortest Path test dataset | success rate: 88.6% -> 92.6% | Helped |
| use IRED instead of RRN | Sudoku solving · Sudoku harder dataset | accuracy: 28.6% -> 62.1% | Helped |
Surprising findings
- Plain diffusion baselines perform much worse than energy-based approaches on several tasks, e.g., 45.2% vs 92.6% on planning.
- Without gradient descent refinement and contrastive shaping, Sudoku harder-set accuracy drops from 62.1% to 35.0%.
Limitations
- Inference requires many gradient descent steps, making it slower than specialized algorithms for known tasks.
- No additional memory mechanism, so it may be less effective for tasks needing intermediate storage such as chain-of-thought reasoning.
- Annealed landscape sequence is fixed by a predefined Gaussian noise schedule rather than learned.
Experimental setup
Datasets: SAT-Net Sudoku dataset RRN Sudoku dataset Visual Sudoku randomly generated connectivity graphs randomly generated planning graphs continuous algorithmic reasoning matrix tasks
Benchmarks: Continuous algorithmic reasoning (addition, matrix completion, matrix inverse) Sudoku (SAT-Net test and harder RRN dataset) Visual Sudoku graph connectivity shortest path planning
Metrics: MSE accuracy
Baselines: Feedforward Recurrent Programmatic Diffusion IREM SAT-Net RRN
Builds on: IREM denoising diffusion models Energy-based models
Closest papers by meaning
- Generalizable Reasoning through Compositional Energy Minimization (NeurIPS 2025 · similarity 0.89)
- VFScale: Intrinsic Reasoning through Verifier-Free Test-time Scalable Diffusion Model (ICLR 2026 · similarity 0.87)
- PEAR: Phase Entropy Aware Reward for Efficient Reasoning (ICLR 2026 · similarity 0.86)
- Constraints-Guided Diffusion Reasoner for Neuro-Symbolic Learning (AAAI 2026 · similarity 0.86)
- Boosting Inference with Guided Reasoning: Stochastic Exploration for Recursive Models (NeurIPS 2026 · similarity 0.86)
- Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.