# Learning Iterative Reasoning through Energy Diffusion

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

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 55 references that are in the map.
- First posted on arXiv 17 Jun 2024.

## Builds on (most cited first, 30 of 44)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Spatial Temporal Graph Convolutional Networks for Skeleton-Based Action Recognition](https://mlsanity.com/paper/1801.07455/) (2018)
- [MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms](https://mlsanity.com/paper/1905.13319/) (2019)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Neural Probabilistic Logic Programming in DeepProbLog](https://mlsanity.com/paper/1907.08194/) (2019)
- [Inter-GPS: Interpretable Geometry Problem Solving with Formal Language and Symbolic Reasoning](https://mlsanity.com/paper/2105.04165/) (2021)
- [Integrated Task and Motion Planning](https://mlsanity.com/paper/2010.01083/) (2020)
- [Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization](https://mlsanity.com/paper/1809.05504/) (2018)
- [DIFUSCO: Graph-based Diffusion Solvers for Combinatorial Optimization](https://mlsanity.com/paper/2302.08224/) (2023)
- [Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC](https://mlsanity.com/paper/2302.11552/) (2023)
- [A Theory of Generative ConvNet](https://mlsanity.com/paper/1602.03264/) (2016)
- [SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver](https://mlsanity.com/paper/1905.12149/) (2019)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [Adaptive Neural Networks for Efficient Inference](https://mlsanity.com/paper/1702.07811/) (2017)
- [End-to-End Differentiable Proving](https://mlsanity.com/paper/1705.11040/) (2017)
- [What Can Neural Networks Reason About?](https://mlsanity.com/paper/1905.13211/) (2019)
- [Neural Logic Machines](https://mlsanity.com/paper/1904.11694/) (2019)
- [Improved Contrastive Divergence Training of Energy Based Models](https://mlsanity.com/paper/2012.01316/) (2020)
- [Neural GPUs Learn Algorithms](https://mlsanity.com/paper/1511.08228/) (2015)
- [Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks](https://mlsanity.com/paper/2106.04537/) (2021)
- [Hierarchical Multiscale Recurrent Neural Networks](https://mlsanity.com/paper/1609.01704/) (2016)
- [VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models](https://mlsanity.com/paper/2010.00654/) (2020)
- [Neural Programmer-Interpreters](https://mlsanity.com/paper/1511.06279/) (2015)
- [NeurASP: Embracing Neural Networks into Answer Set Programming](https://mlsanity.com/paper/2307.07700/) (2023)

## Built on by (most cited first, 30 of 36)

- [Hierarchical Reasoning Model](https://mlsanity.com/paper/2506.21734/) (2025)
- [Inference-time Scaling of Diffusion Models through Classical Search](https://mlsanity.com/paper/2505.23614/) (2025)
- [Artificial Kuramoto Oscillatory Neurons](https://mlsanity.com/paper/2410.13821/) (2024)
- [Coevolutionary Continuous Discrete Diffusion: Make Your Diffusion Language Model a Latent Reasoner](https://mlsanity.com/paper/2510.03206/) (2025)
- [Energy-Based Transformers are Scalable Learners and Thinkers](https://mlsanity.com/paper/2507.02092/) (2025)
- [Fixed-Point Reasoners: Stable and Adaptive Deep Looped Transformers](https://mlsanity.com/paper/2606.18206/) (2026)
- [Non-myopic Generation of Language Models for Reasoning and Planning](https://mlsanity.com/paper/2410.17195/) (2024)
- [A Formal Comparison Between Chain of Thought and Latent Thought](https://mlsanity.com/paper/2509.25239/) (2025)
- [Thinking with Looped Flows](https://mlsanity.com/paper/2609.11801/) (2026)
- [VFScale: Intrinsic Reasoning through Verifier-Free Test-time Scalable Diffusion Model](https://mlsanity.com/paper/2502.01989/) (2025)
- [Self-Supervised Transformers as Iterative Solution Improvers for Constraint Satisfaction](https://mlsanity.com/paper/2502.15794/) (2025)
- [Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning](https://mlsanity.com/paper/2505.16950/) (2025)
- [On the Role of Computation in Reinforcement Learning](https://mlsanity.com/paper/2602.05999/) (2026)
- [A Kinetic Energy Perspective of Flow Matching](https://mlsanity.com/paper/2602.07928/) (2026)
- [Hyper-SET: Designing Transformers via Hyperspherical Energy Minimization](https://mlsanity.com/paper/2502.11646/) (2025)
- [Selection, Reflection and Self-Refinement: Revisit Reasoning Tasks via a Causal Lens](https://mlsanity.com/paper/2510.08222/) (2025)
- [Generalizable Reasoning through Compositional Energy Minimization](https://mlsanity.com/paper/2510.20607/) (2025)
- [EBT-Policy: Energy Unlocks Emergent Physical Reasoning Capabilities](https://mlsanity.com/paper/2510.27545/) (2025)
- [GenPlan: Generative Sequence Models as Adaptive Planners](https://mlsanity.com/paper/2412.08565/) (2024)
- [Thinking into the Future: Latent Lookahead Training for Transformers](https://mlsanity.com/paper/2603.20219/) (2026)
- [Large Neighborhood Search meets Iterative Neural Constraint Heuristics](https://mlsanity.com/paper/2603.20801/) (2026)
- [Learning 3D Reconstruction with Priors in Test Time](https://mlsanity.com/paper/2604.03878/) (2026)
- [Constraints-Guided Diffusion Reasoner for Neuro-Symbolic Learning](https://mlsanity.com/paper/2508.16524/) (2025)
- [PID-controlled Langevin Dynamics for Faster Sampling of Generative Models](https://mlsanity.com/paper/2511.12603/) (2025)
- [Can Continuous-Time Diffusion Models Generate and Solve Globally Constrained Discrete Problems? A Study on Sudoku](https://mlsanity.com/paper/2601.20363/) (2026)
- [Tutorial on Reasoning for IR & IR for Reasoning](https://mlsanity.com/paper/2602.03640/) (2026)
- [Blocked Gibbs meets Diffusion Transformers: Unsupervised Learning for Constraint Optimization](https://mlsanity.com/paper/2605.25129/) (2026)
- [Text Dictates, Music Decorates: Energy-based Attention for Editable Dance Motion Generation](https://mlsanity.com/paper/2606.22726/) (2026)
- [Retrieval-Warmed Energy-Based Reasoning: A Five-Arm Ablation Methodology for Diffusion-as-Inference on Structured Reasoning Tasks](https://mlsanity.com/paper/2606.26476/) (2026)
- [Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions](https://mlsanity.com/paper/2607.11555/) (2026)

## Abstract (arXiv)

We introduce iterative reasoning through energy diffusion (IRED), a novel framework for learning to reason for a variety of tasks by formulating reasoning and decision-making problems with energy-based optimization. IRED learns energy functions to represent the constraints between input conditions and desired outputs. After training, IRED adapts the number of optimization steps during inference based on problem difficulty, enabling it to solve problems outside its training distribution -- such as more complex Sudoku puzzles, matrix completion with large value magnitudes, and pathfinding in larger graphs. Key to our method's success is two novel techniques: learning a sequence of annealed energy landscapes for easier inference and a combination of score function and energy landscape supervision for faster and more stable training. Our experiments show that IRED outperforms existing methods in continuous-space reasoning, discrete-space reasoning, and planning tasks, particularly in more challenging scenarios. Code and visualizations at https://energy-based-model.github.io/ired/