[论文解读] Learning Iterative Reasoning through Energy Minimization

Yilun Du, Shuang Li+2 more|arXiv (Cornell University)|Jun 30, 2022

Advanced Graph Neural Networks被引用 5

一句话总结

该论文提出了迭代推理作为能量最小化（IREM）框架，将推理建模为在学习到的能量景观上的迭代能量最小化过程，以实现在神经网络中可扩展、可泛化的算法推理。通过将每个推理步骤视为一个优化过程，IREM能动态调整复杂问题的计算资源，并在连续和图结构的算法任务上取得最先进性能，包括操作的递归组合。

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

研究动机与目标

• 为解决深度神经网络在执行非平凡推理，尤其是需要迭代、受控处理的算法任务方面的局限性。
• 开发一种框架，使神经网络能够通过调整优化步数，动态增加对更复杂问题的计算投入。
• 设计一种可泛化到更复杂和分布外输入的推理系统，通过迭代优化实现。
• 实现对已学习算法的递归应用，以解决嵌套推理任务。
• 提供一种稳定、可微且可泛化的替代方案，以替代现有依赖强化学习或启发式终止策略的迭代推理方法。

提出的方法

• 使用神经网络参数化能量函数 $ E_{\theta}(\mathbf{z}|\mathbf{x},\mathbf{y}) $，以在给定输入 $ \mathbf{x}, \mathbf{y} $ 的情况下，定义可能输出 $ \mathbf{z} $ 的能量景观。
• 通过梯度下降在能量函数上进行迭代能量最小化，以在多步内逐步优化候选解。
• 将能量最小化过程的收敛性（即达到局部最小值）作为推理完成的信号。
• 通过优化步骤的反向传播端到端训练能量函数，实现可微推理。
• 将能量最小化过程整合到可微流水线中，支持分布内和分布外的泛化。
• 将该框架应用于连续向量输入和图结构数据，展示了在不同领域中的鲁棒性。

实验结果

研究问题

• RQ1 能否将神经网络中的迭代推理有效建模为能量最小化，以提升在算法任务上的性能？
• RQ2 所提出的能量最小化框架在更难、更复杂的算法问题实例上如何泛化？
• RQ3 该框架能否支持已学习算法的递归组合，即使中间输出处于分布外？
• RQ4 与依赖学习到的终止策略的现有迭代推理方法相比，该方法在稳定性和可扩展性方面表现如何？
• RQ5 能量值与解的质量相关程度如何？能否作为可靠的终止信号？

主要发现

• IREM在连续算法推理任务上达到最先进性能，在2个操作的加法组合任务上测试均方误差（MSE）为0.0014，显著优于次优方法（5个操作时为0.0078）。
• 在加法操作的递归组合中，IREM表现出缓慢而稳定的误差增长，10个组合操作的MSE仅为0.0422，而循环基线方法为4.8706。
• 候选解的能量值与其与真实值的距离高度相关，表明低能量解接近最优解。
• IREM在递归执行过程中对分布外的中间输出具有鲁棒性，即使组合多达10个操作，仍能保持低误差。
• 该方法对步长超参数的选择不敏感，固定步长为100时在各项实验中均表现出一致性能。
• IREM能够准确近似复杂操作，例如对四个输入矩阵求和，其预测输出在可视化中与真实值高度吻合。