CausalMamba: Scalable Conditional State Space Models for Neural Causal Inference

TL;DR

CausalMamba tackles the challenge of inferring directed neural connectivity from fMRI by addressing the ill-posed BOLD-to-neural inverse problem through a differentiable BOLD deconvolution stage followed by ROI-conditioned causal inference, powered by a state-space backbone. The method employs a three-stage curriculum to pre-train deconvolution (Stage 1) and causality mapping (Stage 2) before end-to-end fine-tuning (Stage 3), anchored by the Conditional Mamba architecture that injects ROI-specific modulation into universal temporal dynamics. On synthetic data, it achieves a substantial accuracy improvement over DCM (approximately a 38% gain) and demonstrates robust scalability with ROI count, while on real HCP task fMRI data it recovers canonical pathways with high system-level fidelity (KPRR around 0.88) and reveals stimulus- and task-driven hub reconfigurations that conventional methods miss. The approach provides a practical, scalable tool for large-scale causal inference in neuroscience, enabling both canonical circuit discovery and flexible network dynamics understanding, with potential for broader adoption and cross-task generalization.

Abstract

We introduce CausalMamba, a scalable framework that addresses fundamental limitations in fMRI-based causal inference: the ill-posed nature of inferring neural causality from hemodynamically distorted BOLD signals and the computational intractability of existing methods like Dynamic Causal Modeling (DCM). Our approach decomposes this complex inverse problem into two tractable stages: BOLD deconvolution to recover latent neural activity, followed by causal graph inference using a novel Conditional Mamba architecture. On simulated data, CausalMamba achieves 37% higher accuracy than DCM. Critically, when applied to real task fMRI data, our method recovers well-established neural pathways with 88% fidelity, whereas conventional approaches fail to identify these canonical circuits in over 99% of subjects. Furthermore, our network analysis of working memory data reveals that the brain strategically shifts its primary causal hub-recruiting executive or salience networks depending on the stimulus-a sophisticated reconfiguration that remains undetected by traditional methods. This work provides neuroscientists with a practical tool for large-scale causal inference that captures both fundamental circuit motifs and flexible network dynamics underlying cognitive function.

Figures (11)

• Figure 1: Framework of CausalMamba.(a) Stage 1: Hemodynamic Deconvolution. The model takes a BOLD signal and jointly learns a region-specific Hemodynamic Response Function (HRF) and its underlying latent neural activity. These two components are then convolved to reconstruct the original BOLD signal, guiding training. (b) Stage 2: Causal Inference. A Conditional Mamba encoder maps the deconvolved neural activity into feature representations. A causality mapper then uses these features to estimate a directed causal graph, defined by coupling strengths and delay timings between regions. ROI-index conditioning is controlled by ablation (shuffling, shared adapters) to mitigate memorization. (c) Integrated Pipeline. The full model is fine-tuned end-to-end using a three-stage curriculum. All training occurs on simulated data; real HCP data is used only for evaluation.
• Figure 2: Architecture of Stage 1 BOLD Deconvolution Modules. The two parallel modules that map the observed BOLD signal to its underlying neural activity are illustrated. (a) The HRF Generator uses a Conditional Mamba encoder followed by a linear layer to regress the six parameters defining a unique, subject- and ROI-specific HRF.(b) Concurrently, the BOLD Deconvolver uses a Conditional Mamba encoder and a cross-attention mechanism with learnable queries to infer the parameters (amplitude, timing, width) of the latent LFP-like neural signal.
• Figure 3: Architecture of Stage 2 Causality Mapper. This module infers the final causal graph from the feature representation generated by the Conditional Mamba backbone in Stage 2. To preserve temporal structure, the time-series features are first averaged to produce node-level feature vectors for each ROI. These vectors are then expanded and concatenated to form pairwise features, which are passed to two lightweight MLPs—rather than deeper graph modules—to directly predict the final Coupling Strengths and Delay Timings between all pairs of ROIs, thereby maintaining model simplicity and preserving temporal fidelity.
• Figure 4: Architecture of Conditional Mamba. This module constitutes the central architectural element of our approach. The main Mamba pathway learns universal, ROI-agnostic temporal dynamics. In parallel, an adaptive stream uses the ROI index to select dedicated weights (W[i]) and biases (b[i]) that apply a region-specific affine modulation to the main pathway’s output. The two streams are then integrated via a residual connection, enabling the model to combine global dynamics with region-specific characteristics while mitigating the risk of the model memorizing ROI identifiers rather than learning true neurovascular heterogeneity.
• Figure 5: Signed, Directed Coupling on the 6-ROI Simulation. Each panel shows a signed, directed coupling matrix with rows = sources (from) and columns = targets (to); colors encode normalized weights in $[-1,1]$ with a zero-centered colormap identical across panels. Diagonal entries denote self-inhibition when present. a ground-truth coupling; b CausalMamba; c Granger Causality; d Neural Granger Causality; e rDCM; f DCM. All methods were evaluated under the same simulated data and preprocessing; additional metrics are summarized in the main results table.