Beyond Interpolation: Extrapolative Reasoning with Reinforcement Learning and Graph Neural Networks

2025-02-06

AI Summary

This work addresses the out-of-distribution (OOD) generalization challenge for neural models on logic puzzles—specifically, correct reasoning on unseen puzzles significantly larger and more complex than those in training. Method: We propose a reinforcement learning framework integrating graph neural networks (GNNs) with proximal policy optimization (PPO), explicitly encoding structured logical relations in puzzles. Our approach incorporates recursive state updates and a structured sparse reward mechanism to guide learning. Contribution/Results: We present the first systematic disentanglement of the effects of inductive bias, reward design, and recurrent modeling on sequential logical extrapolation, establishing an interpretable OOD generalization analysis framework. Experiments demonstrate substantial improvements in extrapolation performance on logic puzzles of unseen scales and difficulty levels. Results validate that synergistic optimization of architectural priors and training mechanisms is critical for robust generalization in logical reasoning tasks.

Technology Category

• Machine Learning: Statistical Relational/Logic Learning
• Knowledge Representation and Reasoning: Logic Programming
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Despite incredible progress, many neural architectures fail to properly generalize beyond their training distribution. As such, learning to reason in a correct and generalizable way is one of the current fundamental challenges in machine learning. In this respect, logic puzzles provide a great testbed, as we can fully understand and control the learning environment. Thus, they allow to evaluate performance on previously unseen, larger and more difficult puzzles that follow the same underlying rules. Since traditional approaches often struggle to represent such scalable logical structures, we propose to model these puzzles using a graph-based approach. Then, we investigate the key factors enabling the proposed models to learn generalizable solutions in a reinforcement learning setting. Our study focuses on the impact of the inductive bias of the architecture, different reward systems and the role of recurrent modeling in enabling sequential reasoning. Through extensive experiments, we demonstrate how these elements contribute to successful extrapolation on increasingly complex puzzles.These insights and frameworks offer a systematic way to design learning-based systems capable of generalizable reasoning beyond interpolation.

Problem

• Generalizing beyond training distribution
• Solving scalable logical puzzles
• Enhancing extrapolative reasoning with RL

Innovation

• Graph Neural Networks
• Reinforcement Learning
• Extrapolative Reasoning