Learning Iterative Reasoning through Energy Minimization

TL;DR

This paper introduces a neural network framework that models iterative reasoning as an energy minimization process, enabling adaptive computational effort and improved performance on complex algorithmic tasks.

Contribution

It proposes a novel energy-based iterative reasoning framework allowing dynamic computational effort and recursive problem solving in neural networks.

Findings

Successfully solves complex algorithmic reasoning tasks

Achieves better generalization on graph and continuous domains

Enables recursive reasoning for nested problems

Abstract

Deep learning has excelled on complex pattern recognition tasks such as image classification and object recognition. However, it struggles with tasks requiring nontrivial reasoning, such as algorithmic computation. Humans are able to solve such tasks through iterative reasoning -- spending more time thinking about harder tasks. Most existing neural networks, however, exhibit a fixed computational budget controlled by the neural network architecture, preventing additional computational processing on harder tasks. In this work, we present a new framework for iterative reasoning with neural networks. We train a neural network to parameterize an energy landscape over all outputs, and implement each step of the iterative reasoning as an energy minimization step to find a minimal energy solution. By formulating reasoning as an energy minimization problem, for harder problems that lead to more complex energy landscapes, we may then adjust our underlying computational budget by running a more complex optimization procedure. We empirically illustrate that our iterative reasoning approach can solve more accurate and generalizable algorithmic reasoning tasks in both graph and continuous domains. Finally, we illustrate that our approach can recursively solve algorithmic problems requiring nested reasoning