Energy-Based Reasoners: Your Recurrent Reasoner Should Be an Energy-Based Model
Organizations: Duke University · UIUC · Harvard University
Abstract
Looping is a powerful inductive bias for reasoning. By repeatedly iterating a recurrent block, looped models can scale their effective depth and enable multi-step reasoning in latent space. However, existing looped models are not inherently guided by an explicit measure of progress, making it difficult to determine when to stop looping and how to improve upon the current state. To enable this, we introduce Energy-Based Reasoners (EBRs), which learn an energy function over latent reasoning states and formulate looped updates as gradient descent on this energy. The energy scores latent states, while its gradient guides refinement and provides a signal for stopping. To train EBRs effectively, we further introduce exploration-based training to discover additional modes and to better shape the energy landscape over a broader set of states. Across four algorithmic reasoning benchmarks, EBRs match or outperform prior recurrent reasoning models, achieving 99.9% accuracy on Sudoku-Extreme while using 345 less inference compute than the next-best model and maintaining comparable training compute. On all multi-solution benchmarks, EBRs achieve state-of-the-art coverage of distinct valid solutions while maintaining near-100% single-sample accuracy. These results suggest that energy-based models are a promising approach for strong reasoning performance.
Figures & tables
| Method | # Params | Sudoku-Extreme (↑) | Maze-Hard (↑) |
| Single sample (%) | |||
| TRM [ 12 ] | 5–7M | 87.4 | 85.3 |
| EqR [ 18 ] | 5–7M | 93.0 | – |
| FPRM [ 17 ] | 7M | 94.2 | 87.0 |
| GRAM [ 19 ] | 10M | 86.6 2.2 | – |
| Looped Flows [ 26 ] | 5–7M | 97.9 0.4 | 86.7 1.1 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| TRM | EBR | |
| Block output | the next latent state | a scalar energy ( eq. 7 ) |
| State update | replaced by the block output | a step along |
| Latent state | hierarchical, and | a single |
| Initial state | one fixed vector for every input | fresh draw from a normal |
| Halting | learned Q head | residual of the descent ( eq. 4 ) |
| N-Queens | Graph Coloring | |||||
| Sudoku | Maze | 8-vertex | 10-vertex | |||
| Model | ||||||
| Hidden size | 256 | 256 | 512 | 512 | 512 | 512 |
| Parameters | 1.8M | 1.8M | 7.1M | 7.1M | 7.1M | 7.1M |
| Per-cell energy | linear | softplus | linear | linear | linear | linear |
| Descent | ||||||
| N-Queens | Graph Coloring | |||||
| Sudoku | Maze | 8-vertex | 10-vertex | |||
| Training steps | 60k | 50k | 30k | 45k | 70k | 70k |
| Wall time (hours) | 3 | 3 | 4 | 8 | 5 | 8 |
| Training steps | Total EFLOPs | result | ||||
| Task | EBR | GRAM | EBR | GRAM | EBR | GRAM |
| Sudoku-Extreme | 60k | 65.1k | 1.92 | 2.26 | 99.41 0.12 | 86.6 2.2 |
| Maze-Hard | 50k | – | 1.16 | – | 87.37 1.1 | – |
| N-Queens | 30k | 47.0k | 2.45 | 1.08 | 100 0 | 99.5 0.3 |
| N-Queens | 45k | 140.8k | 5.82 | 4.74 | 99.51 0.45 | 89.7 2.7 |
| Graph Coloring, 8-vertex | 70k | 45.5k | 2.46 | 0.67 | 0 0 | 8.3 2.5 |
| Avg. iterations | Pass@1 | Pass@8 | TFLOPs / puzzle | Saving | |
| No stopping | 64 | 99.41 0.12 | 99.92 0.05 | 6.60 | 1.0 |
| 0.1 | 7.6 1.4 | 99.41 0.12 | 99.92 0.05 | 0.82 | 8.0 |
| 0.2 | 5.6 0.6 | 99.41 0.12 | 99.92 0.05 | 0.61 | 10.8 |
| 0.5 | 3.6 0.1 | 99.42 0.12 | 99.86 0.07 | 0.39 | 17.1 |
| 1 | 3.5 0.1 | 99.41 0.12 | 99.88 0.07 | 0.38 | 17.4 |
| 2 | 3.5 0.1 | 98.81 0.17 | 99.89 0.06 | 0.38 | 17.6 |
| Inference | Labeled-path exact match | Any shortest path |
| EBR, Pass@1 | 87.37 1.10 | 94.07 0.91 |
| EBR, Pass@5 | 88.27 0.32 | 94.77 0.59 |