Disentangling Reasoning Capabilities from Language Models with Compositional Reasoning Transformers

TL;DR

ReasonFormer is a unified, modular reasoning framework inspired by cognitive science, which dynamically composes specialized reasoning modules to improve performance and generalization across diverse complex decision-making tasks.

Contribution

This work introduces ReasonFormer, a novel modular reasoning framework that decouples representation and reasoning modules, enabling dynamic composition and improved task generalization.

Findings

Substantial performance improvements on 11 reasoning datasets.

Enhanced few-shot learning by composing pre-trained reasoning skills.

Modularity allows activation of distinct reasoning skills at different depths.

Abstract

This paper presents ReasonFormer, a unified reasoning framework for mirroring the modular and compositional reasoning process of humans in complex decision-making. Inspired by dual-process theory in cognitive science, the representation module (automatic thinking) and reasoning modules (controlled thinking) are decoupled to capture different levels of cognition. Upon the top of the representation module, the pre-trained reasoning modules are modular and professional in specific and fundamental reasoning skills (e.g., logic, simple QA, etc). To mimic the controlled compositional thinking process, different reasoning modules are dynamically activated and composed in both parallel and cascaded manners to control what reasoning skills are activated and how deep the reasoning process will be reached to solve the current problems. The unified reasoning framework solves multiple tasks with a single model, and is trained and inferred in an end-to-end manner. Evaluated on 11 datasets requiring different reasoning skills and complexity, ReasonFormer demonstrates substantial performance boosts, revealing the compositional reasoning ability. Few-shot experiments exhibit better generalization ability by learning to compose pre-trained skills for new tasks with limited data, and decoupling the representation module and the reasoning modules. Further analysis shows the modularity of reasoning modules as different tasks activate distinct reasoning skills at different reasoning depths.