Generalizable Reasoning through Compositional Energy Minimization

TL;DR

This paper introduces a compositional energy minimization framework for reasoning tasks, enabling models to generalize better to complex problems by constructing and optimizing over energy landscapes of subproblems.

Contribution

It proposes a novel approach that learns energy landscapes over subproblems and combines them for reasoning, improving generalization to larger and more complex problems.

Findings

Outperforms state-of-the-art methods on reasoning benchmarks

Enables reasoning over larger, more complex problems

Incorporates constraints during inference for better solutions

Abstract

Generalization is a key challenge in machine learning, specifically in reasoning tasks, where models are expected to solve problems more complex than those encountered during training. Existing approaches typically train reasoning models in an end-to-end fashion, directly mapping input instances to solutions. While this allows models to learn useful heuristics from data, it often results in limited generalization beyond the training distribution. In this work, we propose a novel approach to reasoning generalization by learning energy landscapes over the solution spaces of smaller, more tractable subproblems. At test time, we construct a global energy landscape for a given problem by combining the energy functions of multiple subproblems. This compositional approach enables the incorporation of additional constraints during inference, allowing the construction of energy landscapes for problems of increasing difficulty. To improve the sample quality from this newly constructed energy landscape, we introduce Parallel Energy Minimization (PEM). We evaluate our approach on a wide set of reasoning problems. Our method outperforms existing state-of-the-art methods, demonstrating its ability to generalize to larger and more complex problems. Project website can be found at: https://alexoarga.github.io/compositional_reasoning/