---
title: 'ReasonFormer: Modular Compositional Reasoning'
url: https://www.emergentmind.com/topics/reasonformer
type: topic
---

# ReasonFormer: Modular Compositional Reasoning

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 [2210.11265].

## 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 $\alpha_i$ over the $K$ reasoning modules:
   $$
   \alpha_i = \operatorname{Softmax}(\mathrm{FFN}(T(H^{i-1}))),
   $$
   followed by top-$k$ sparsification to activate only the $k$ most relevant modules.
2. **Module Outputs:** Each selected module processes $H^{i-1}$, producing $H_j^i = M_j(H^{i-1}; \theta_j)$ for $j\in \text{TopK}(\alpha_i)$.
3. **Fusion of Skills:** Activated module outputs are fused by a weighted sum:
   $$
   \widetilde H^i = \sum_{j=1}^K \alpha_{i,j} M_j(H^{i-1}; \theta_j)
   $$
4. **Depth Control:** Depth is dynamically controlled via a gating network $G_{\mathrm{stop}}$ that determines the extent of the residual update:
   $$
   G_{\mathrm{stop}}(H^i) = \sigma(W_2 \mathrm{ReLU}(W_1 \mathrm{Pool}(H^i) + b_1) + b_2)
   $$
   The representation is updated with:
   $$
   H^i = H^{i-1} + G_{\mathrm{stop}}(H^i)(\widetilde H^i - H^{i-1})
   $$
   If $G_{\mathrm{stop}}$ is near zero, the iterative reasoning process terminates.

### Inference Pseudocode

```
1. H⁰ ← f_rep(X;θ_rep)
2. for i←1 to N:
   αᵢ ← softmax_router(H^{i−1})
   select top-k modules J ← TopK(αᵢ)
   forall j∈J: H_j^i ← R_j(H^{i−1};θ_j)
   \widetilde H^i ← Σ_{j∈J} α_{i,j}·H_j^i
   g ← G_stop(\widetilde H^i)
   H^i ← H^{i−1} + g·(\widetilde H^i−H^{i−1})
   if g<ε: break
3. Decode answer Y ← Decoder(H^i;θ_dec)
```

## 4. Training and Adaptation

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

- **Pre-training:** Supervised with a generative loss $L_{\mathrm{gen}}$ (teacher-forcing cross-entropy over text targets) across tasks corresponding to the reasoning skills. Additionally, the skill-routing loss $L_{\mathrm{router}}$ is applied using known skill labels for each pre-training instance:
  $$
  L_{\mathrm{router}} = -\sum_{i=1}^N \log \alpha_{i,j^*}
  $$
  The total pre-training loss is $L_{\mathrm{pre}} = L_{\mathrm{gen}} + \lambda L_{\mathrm{router}}$.
- **Downstream Adaptation:** Only the generative loss $L_{\mathrm{gen}}$ 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.

Source: https://www.emergentmind.com/topics/reasonformer