AbstRaL: Augmenting LLMs' Reasoning by Reinforcing Abstract Thinking

TL;DR

This paper introduces AbstRaL, a reinforcement learning-based method to enhance LLMs' abstract reasoning skills, improving robustness in grade school math tasks and general reasoning under distribution shifts.

Contribution

Proposes a novel RL-based approach to teach LLMs abstract reasoning, outperforming supervised fine-tuning in robustness and generalization for math and reasoning tasks.

Findings

AbstRaL significantly reduces performance drops on GSM benchmarks.

Abstract reasoning via RL improves LLMs' out-of-distribution mathematical reasoning.

Enhanced abstract thinking benefits general reasoning capabilities.

Abstract

Recent studies have shown that large language models (LLMs), especially smaller ones, often lack robustness in grade school math (GSM) reasoning. In particular, they tend to experience performance drops when faced with distribution shifts, such as changes to numerical or nominal variables, or insertions of distracting clauses. A possible strategy to address this involves generating synthetic data to further "instantiate" reasoning problems on potential variations. In this work, we instead focus on the strategy of "abstracting" reasoning problems. This not only helps counteract distribution shifts but also facilitates the connection to symbolic tools for deriving solutions. Focusing on GSM, we find that this abstraction process is better acquired through reinforcement learning (RL) than just supervised fine-tuning, which often fails to produce faithful abstractions. Our method, AbstRaL -- which promotes abstract reasoning in LLMs using RL on granular abstraction data -- significantly mitigates performance degradation on recent GSM perturbation benchmarks. Besides, improving GSM robustness via AbstRaL is shown to also implicitly benefit LLMs' capabilities on OOD mathematical and general reasoning tasks, indicating that abstract thinking broadly enables better generalizability.