Elamur Research⌕
☰
home / papers / 2024 / Learning Iterative Reasoning through Energy Diffusion

Learning Iterative Reasoning through Energy Diffusion

Paper · 2024 · Yilun Du, Jiayuan Mao and Joshua B. Tenenbaum

Summary

Many reasoning and decision problems, such as puzzles, constraint satisfaction and planning, can be framed as finding an output that satisfies constraints tied to an input. The MIT authors (ICML 2024) present IRED. It learns an energy function E_θ(x, y) from paired examples so that valid answers have low energy, and it solves new problems by gradient-based minimization over y. Earlier energy-based reasoners trained slowly and unstably because learning passed gradients through the whole optimization loop, and their landscapes contained many local minima. IRED draws on diffusion models in two ways. It trains a series of annealed landscapes that are optimized from smooth to sharp. It also supervises them with a denoising (score) loss plus a contrastive loss on mined negatives, so it never backpropagates through the solver. Because the number of optimization steps can be scaled to problem difficulty, IRED handles harder inputs than it saw in training. Across continuous, discrete and planning tasks it outperforms recurrent, energy-based and diffusion baselines, most clearly on difficult instances. This points to one domain-agnostic recipe for learned iterative reasoning.

Key points

  • Tasks:

    • continuous: matrix completion and inversion;
    • discrete: Sudoku and graph connectivity;
    • planning: path finding on graphs.

    Harder test settings include larger value magnitudes, Sudoku with fewer given digits, and larger graphs.

  • Baselines: recurrent adaptive computation (Palm et al., 2018), the energy-based IREM (Du et al., 2022) and a diffusion model (Ho et al., 2020). IRED wins overall, and the plain diffusion model trails by a wide margin in planning.

  • Visual Sudoku: the board is made of MNIST digit images (Visual Sudoku dataset, Wang et al., 2019), encoded with a CNN. IRED keeps its advantage over the baseline.

  • Ablations and encoders: removing gradient-based optimization, multi-step optimization per landscape, or contrastive energy shaping hurts results. Planning models use a shared STGCN encoder.

  • Limitations: inference needs many gradient steps, so it is slower than dedicated algorithms on known problems. The noise-based landscape schedule is fixed rather than learned. The method has no explicit memory for chain-of-thought-style intermediate results.

  • Code: energy-based-model.github.io/ired

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 3The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksComputer Networks and Communications · 115 works

🧭 Topics 1Research subjects and languages the work is about.Reasoning · 1,948 works

Cite

Cite as Yilun Du, Jiayuan Mao and Joshua B. Tenenbaum. 2024. Learning Iterative Reasoning through Energy Diffusion. arXiv:2406.11179.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2406.11179 · arXiv 2406.11179

@misc{du2024learning,
  title = {{Learning Iterative Reasoning through Energy Diffusion}},
  author = {Yilun Du, Jiayuan Mao and Joshua B. Tenenbaum},
  year = {2024},
  doi = {10.48550/arxiv.2406.11179},
  eprint = {2406.11179},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2406.11179}
}

References

55 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Djolonga, J. and Krause, A. Differentiable Learning of Submodular Models. In NeurIPS, 2017. find
  2. Amini, A., Gabriel, S., Lin, S., Koncel-Kedziorski, R., Choi, Y., and Hajishirzi, H. MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms. In NAACL, 2019.
  3. Dong, H., Mao, J., Lin, T., Wang, C., Li, L., and Zhou, D. Neural Logic Machines. In ICLR, 2019. find
  4. Amos, B. and Kolter, J. Z. Optnet: Differentiable Optimization as a Layer in Neural Networks. In ICML, 2017. find
  5. Donti, P. L., Amos, B., and Kolter, J. Z. Task-based End-to-End Model Learning in Stochastic Optimization. In NeurIPS, 2017. find
  6. Anil, C., Pokle, A., Liang, K., Treutlein, J., Wu, Y., Bai, S., Kolter, J. Z., and Grosse, R. B. Path Independent Equilibrium Models Can Better Exploit Test-Time Computation. In NeurIPS, 2022. find
  7. Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019. find
  8. Arbel, M., Zhou, L., and Gretton, A. Generalized Energy Based Models. In ICLR, 2021. find
  9. Du, Y., Li, S., Tenenbaum, B. J., and Mordatch, I. Improved Contrastive Divergence Training of Energy Based Models. In ICML, 2021. find
  10. Bai, S., Kolter, J. Z., and Koltun, V. Deep Equilibrium Models. In NeurIPS, 2019. find
  11. Du, Y., Li, S., Tenenbaum, J., and Mordatch, I. Learning Iterative Reasoning through Energy Minimization. In ICML, 2022. find
  12. Banino, A., Balaguer, J., and Blundell, C. Pondernet: Learning to Ponder. In ICML Workshop on Automated Machine Learning, 2021.
  13. Du, Y., Durkan, C., Strudel, R., Tenenbaum, J. B., Dieleman, S., Fergus, R., Sohl-Dickstein, J., Doucet, A., and Grathwohl, W. S. Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC. In ICML, 2023. find
  14. Bolukbasi, T., Wang, J., Dekel, O., and Saligrama, V. Adaptive Neural Networks for Efficient Inference. In ICML, 2017. find
  15. Brockett, R. W. Dynamical Systems That Sort Lists, Diagonalize Matrices, and Solve Linear Programming Problems. Linear Algebra and Its Applications, 146:79–91, 1991. find
  16. Fang, X., Garrett, C. R., Eppner, C., Lozano-Pérez, T., Kaelbling, L. P., and Fox, D. DiMSam: Diffusion Models as Samplers for Task and Motion Planning under Partial Observability. arXiv:2306.13196, 2023. — [arXiv:2306.13196](https://arxiv.org/abs/2306.13196) find
  17. Bryson, A. E. Applied Optimal Control: Optimization, Estimation and Control. Routledge, 2018. find
  18. Garrett, C. R., Chitnis, R., Holladay, R., Kim, B., Silver, T., Kaelbling, L. P., and Lozano-Pérez, T. Integrated Task and Motion Planning. Annual Review of Control, Robotics, and Autonomous Systems, 4(1):265–293, 2021. find
  19. Cai, J., Shin, R., and Song, D. Making Neural Programming Architectures Generalize via Recursion. In ICLR, 2017. find
  20. Chen, X., Dai, H., Li, Y., Gao, X., and Song, L. Learning to Stop while Learning to Predict. In ICML, 2020a. find
  21. Grathwohl, W., Wang, K.-C., Jacobsen, J.-H., Duvenaud, D., and Zemel, R. Cutting Out the Middle-Man: Training and Evaluating Energy-Based Models Without Sampling. In ICML, 2020. find
  22. Chen, X., Liang, C., Yu, A. W., Song, D., and Zhou, D. Compositional Generalization via Neural-Symbolic Stack Machines. In NeurIPS, 2020b. find
  23. Graves, A. Adaptive Computation Time for Recurrent Neural Networks. arXiv:1603.08983, 2016. find
  24. Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchfiel, B., and Song, S. Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. In RSS, 2023.
  25. Graves, A., Wayne, G., and Danihelka, I. Neural Turing Machines. arXiv:1410.5401, 2014. find
  26. Chung, J., Ahn, S., and Bengio, Y. Hierarchical Multiscale Recurrent Neural Networks. In ICLR, 2017.
  27. Comas, A., Du, Y., Lopez, C. F., Ghimire, S., Sznaier, M., Tenenbaum, J. B., and Camps, O. Inferring Relational Potentials in Interacting Systems. In ICML, 2023. find
  28. Graves, A., Wayne, G., Reynolds, M., Harley, T., Danihelka, I., Grabska-Barwińska, A., Colmenarejo, S. G., Grefenstette, E., Ramalho, T., Agapiou, J., et al. Hybrid Computing Using a Neural Network with Dynamic External Memory. Nature, 538(7626):471–476, 2016. find
  29. Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, Ł. Universal Transformers. In ICLR, 2019. find
  30. He, K., Zhang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. In CVPR, 2016.
  31. Ho, J., Jain, A., and Abbeel, P. Denoising Diffusion Probabilistic Models. In NeurIPS, 2020. find
  32. Xiao, Z., Kreis, K., Kautz, J., and Vahdat, A. VAEBM: A Symbiosis Between Variational Autoencoders and Energy-based Models. In ICLR, 2020. find
  33. Kaiser, Ł. and Sutskever, I. Neural GPUs Learn Algorithms. In ICLR, 2016. find
  34. Xie, J., Lu, Y., Zhu, S.-C., and Wu, Y. A Theory of Generative Convnet. In ICML. PMLR, 2016. find
  35. Kautz, H., Selman, B., and Hoffmann, J. SATPlan: Planning as Satisfiability. In IPC, 2006. find
  36. Xie, J., Lu, Y., Gao, R., Zhu, S.-C., and Wu, Y. N. Cooperative training of descriptor and generator networks. T-PAMI, 42(1):27–45, 2018. find
  37. LeCun, Y., Chopra, S., and Hadsell, R. A Tutorial on Energy-based Learning, 2006. find
  38. Xu, K., Li, J., Zhang, M., Du, S. S., Kawarabayashi, K.-i., and Jegelka, S. What Can Neural Networks Reason About? In ICLR, 2019. find
  39. Lu, P., Gong, R., Jiang, S., Qiu, L., Huang, S., Liang, X., and Zhu, S.-C. Inter-GPS: Interpretable Geometry Problem Solving with Formal Language and Symbolic Reasoning. In ACL, 2021.
  40. Yan, S., Xiong, Y., and Lin, D. Spatial Temporal Graph Convolutional Networks for Skeleton-Based Action Recognition. In AAAI, 2018. find
  41. Manhaeve, R., Dumancic, S., Kimmig, A., Demeester, T., and De Raedt, L. Deepproblog: Neural Probabilistic Logic Programming. In NeurIPS, 2018.
  42. Yang, F., Yang, Z., and Cohen, W. W. Differentiable Learning of Logical Rules for Knowledge Base Reasoning. In NeurIPS, 2017. find
  43. Neelakantan, A., Le, Q. V., and Sutskever, I. Neural Programmer: Inducing Latent Programs with Gradient Descent. In ICLR, 2015. find
  44. Yang, Z., Ishay, A., and Lee, J. NeurASP: Embracing Neural Networks into Answer Set Programming. In IJCAI, 2020. find
  45. Palm, R., Paquet, U., and Winther, O. Recurrent Relational Networks. In NeurIPS, 2018. find
  46. Yang, Z., Ishay, A., and Lee, J. Learning to Solve Constraint Satisfaction Problems with Recurrent Transformer. In ICLR, 2023a. find
  47. Reed, S. and De Freitas, N. Neural Programmer-Interpreters. In ICLR, 2016. find
  48. Yang, Z., Mao, J., Du, Y., Wu, J., Tenenbaum, J. B., Lozano-Pérez, T., and Kaelbling, L. P. Compositional Diffusion-Based Continuous Constraint Solvers. In CoRL, 2023b. find
  49. Rocktäschel, T. and Riedel, S. End-to-End Differentiable Proving. In NeurIPS, 2017.
  50. Rubanova, Y., Sanchez-Gonzalez, A., Pfaff, T., and Battaglia, P. Constraint-Based Graph Network Simulator. In ICML, 2022. find
  51. Schwarzschild, A., Borgnia, E., Gupta, A., Huang, F., Vishkin, U., Goldblum, M., and Goldstein, T. Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks. In NeurIPS, 2021. find
  52. Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep Unsupervised Learning Using Nonequilibrium Thermodynamics. In ICML, 2015.
  53. Sun, Z. and Yang, Y. Difusco: Graph-Based Diffusion Solvers for Combinatorial Optimization. In NeurIPS, 2023. find
  54. Wang, P.-W., Donti, P., Wilder, B., and Kolter, Z. SATNet: Bridging Deep Learning and Logical Reasoning using a Differentiable Satisfiability Solver. In ICML, 2019. find
  55. Wilder, B., Dilkina, B., and Tambe, M. Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization. In AAAI, 2019. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/du-2024-learning-iterative-reasoning-through-energy/