ReasonFormer: Modular Compositional Reasoning

• ReasonFormer is a unified compositional reasoning framework that mimics human cognitive processes by separating intuitive representation from controlled, stepwise reasoning.
• It employs a T5-based encoder-decoder architecture with distinct representation and modular reasoning modules, dynamically orchestrated via routing and gating mechanisms.
• Empirical evaluations demonstrate enhanced accuracy and robust few-shot generalization across diverse linguistic and logical benchmarks.

ReasonFormer is a unified compositional reasoning framework designed to mirror the modular and compositional reasoning mechanisms of humans in complex decision-making. It is explicitly inspired by dual-process theory from cognitive science, decoupling representation learning (automatic, intuitive cognition) from reasoning (deliberate, controlled cognition). The system is implemented as a T5-based encoder–decoder Transformer, where a dedicated representation module produces contextual features serving as the foundation for a set of reasoning modules, each responsible for specialized and fundamental reasoning skills. ReasonFormer applies these capabilities dynamically through a routing and gating mechanism, enabling complex, task-specific, and interpretable compositions of reasoning skills. The system demonstrates state-of-the-art compositional reasoning performance, robust few-shot generalization, and modularity across 11 diverse reasoning benchmarks (Zhong et al., 2022).

1. Theoretical Motivation

ReasonFormer operationalizes insights from dual-process theory (Daniel, 2017), which distinguishes between two modes of human cognition: System 1 (fast, automatic, intuitive responses) and System 2 (slow, controlled, stepwise reasoning). In ReasonFormer, the representation module functions analogously to System 1, producing intuitive contextualized representations $H^0$ from tokenized input $X$. The subsequent reasoning modules implement System 2, each corresponding to specialized "thinking skills" such as logic, question-answering (QA), factual recall, natural language inference (NLI), and named entity recognition (NER). These modules can be dynamically composed to reflect the compositional structure of complex human reasoning tasks.

2. Model Architecture

ReasonFormer is constructed as a T5-based encoder–decoder Transformer with two principal encoder components: the representation module and a set of modular reasoning modules.

• Representation Module: Accepts tokenized input $X = [x_1, ..., x_m]$ (prefixed by a [CLS] token) and parameterized by nine stacked Transformer layers initialized from T5-base encoder layers 1–9. Outputs a contextual representation $H^0 \in \mathbb{R}^{(m+1)\times d}$, denoted as $H^0 = f_{\mathrm{rep}}(X; \theta_{\mathrm{rep}})$.
• Reasoning Modules (RMs): A set of $K$ skill-specific modules $R_1, ..., R_K$ (e.g., logic, QA, factual recall, NLI, NER, general). Each module is implemented by $L$ shared Transformer layers (from T5-base encoder layers 10–12) augmented with step-specific, two-layer "bottleneck" adapters. Adapter parameters are unique per module and reasoning step but are parameter-efficient due to sharing across modules; RMs share their base parameters across reasoning steps.

The decoder is standard T5, consuming the final contextualized state to generate output text.

3. Compositional Reasoning Process

Reasoning in ReasonFormer unfolds over up to $N$ iterative steps, each involving parallel and cascaded skill activation, routing, module output fusion, and depth control.

1. Parallel Skill Activation: At each step $i$, a skill router computes scores $X$0 over the $X$1 reasoning modules:

$X$2

followed by top-$X$3 sparsification to activate only the $X$4 most relevant modules.

1. Module Outputs: Each selected module processes $X$5, producing $X$6 for $X$7.
2. Fusion of Skills: Activated module outputs are fused by a weighted sum:

$X$8

1. Depth Control: Depth is dynamically controlled via a gating network $X$9 that determines the extent of the residual update:

$X = [x_1, ..., x_m]$0

The representation is updated with:

$X = [x_1, ..., x_m]$1

If $X = [x_1, ..., x_m]$2 is near zero, the iterative reasoning process terminates.

Inference Pseudocode

$X = [x_1, ..., x_m]$8

4. Training and Adaptation

Training ReasonFormer comprises both pre-training and downstream adaptation phases:

• Pre-training: Supervised with a generative loss $X = [x_1, ..., x_m]$3 (teacher-forcing cross-entropy over text targets) across tasks corresponding to the reasoning skills. Additionally, the skill-routing loss $X = [x_1, ..., x_m]$4 is applied using known skill labels for each pre-training instance:

$X = [x_1, ..., x_m]$5

The total pre-training loss is $X = [x_1, ..., x_m]$6.

• Downstream Adaptation: Only the generative loss $X = [x_1, ..., x_m]$7 is applied. All parameters (representation, reasoning modules, adapters, router, stop-gate) are updated end-to-end. The architecture allows "decoupling": in few-shot or transfer settings, freezing either the representation module or the reasoning modules (plus adapters/router) yields comparable performance after minimal adaptation.

5. Empirical Evaluation

ReasonFormer was evaluated on 11 reasoning-focused benchmarks, encompassing a range of linguistic and logical inference challenges:

Dataset	Task Type	Example Metrics
ReClor	Logical Reasoning	Accuracy
CSQA/ARC/PIQA/HellaSwag	Commonsense Reasoning	Accuracy
aNLI	Natural Language Inference	Accuracy
HotpotQA/WikiHop	Multi-Hop QA	EM (HotpotQA), Accuracy
MuTual/DREAM	Dialogue Reasoning	Accuracy
RACE	General QA	Accuracy

Key empirical findings:

• Performance: On average, ReasonFormer improves absolute accuracy by +4.6 points over T5-base and outperforms T5 with reasoning-centric continual pre-training (RPT-T5). For instance, on CSQA: 68.2% (ReasonFormer) vs. 65.1% (vanilla); HotpotQA (EM): 65.2% vs. 63.3%; HellaSwag: 54.9% vs. 33.7%.
• Few-Shot Generalization: In 32-example fine-tuning, freezing reasoning modules or the representation module yields performance comparable to full model fine-tuning, indicating that skill modules and representations are independently reusable for new tasks.
• Routing and Depth Patterns: Routing scores (α) correspond to intuitive skill usage by task (e.g., CSQA relies on factual and QA modules; aNLI activates NER and NLI modules; HotpotQA chains QA skill over steps), and the gating mechanism reflects task complexity (easy tasks terminate after 1–2 steps; multi-hop tasks use all 3 steps).

6. Modularity, Interpretability, and Analysis

ReasonFormer’s modular structure and explicit routing confer interpretability and task-adaptive flexibility:

• Modularity: Reasoning skill modules are architecturally distinct and selectively composable. The routing scores provide insight into which skills are activated per instance and the dynamic “program” comprising the reasoning trace.
• Depth Control: The gating mechanism ensures adaptive depth, with reasoning steps corresponding to problem complexity; multi-hop tasks invoke more steps, aligning module invocation with reasoning demand.
• Interpretability: Analysis of per-instance router activations enables transparent inspection of the model’s compositional reasoning process and skill orchestration.

7. Significance and Implications

ReasonFormer establishes an architecture that separates representation ("seeing") from reasoning ("thinking"), decomposes complex tasks into modular, specialized skill components, and enables dynamic, instance-specific skill composition via parallel and cascaded module orchestration. The empirical results demonstrate substantial improvements across a diverse suite of reasoning tasks, effective generalization with minimal labeled examples, and interpretability of model decisions.

A plausible implication is that similar modular and compositional architectures could offer robust reasoning and transfer capabilities in broader natural language understanding and decision-making applications, while simultaneously providing transparent, task-adaptive reasoning "programs" that can facilitate analysis and diagnosis in high-stakes settings.