FlexiBrain: Resolution-Agnostic Voxel-Level Encoding for Native fMRI

Conference: ECCV2026
Paper: ECCV Paper
Code: https://github.com/OneMore1/FlexiBrain
Area: Medical Imaging
Keywords: native fMRI, physical-unit patches, dynamic kernels, Mamba-JEPA, disease classification

TL;DR

FlexiBrain defines fMRI patches in millimeters and seconds, combines dynamically resized shared kernels with physical positional encoding and Mamba-JEPA, and learns directly from heterogeneous native-space data, achieving the highest accuracy on four of five disease classification tasks while avoiding expensive template registration.

Background & Motivation

Functional magnetic resonance imaging (fMRI) consists of three-dimensional volumes evolving over time, but a voxel and a time point do not represent the same extent across hospitals and acquisition protocols. The paper gives typical spatial resolutions of approximately 2–4 mm and temporal sampling intervals from 720 ms to over 3000 ms; subject-specific anatomy further changes the volume shape. With a fixed voxel-count patch, the same kernel may cover local tissue in one dataset but a larger region or longer interval in another. Conventional pipelines nonlinearly register functional images to a common template before voxel-level modeling or extraction of average region-of-interest (ROI) time series. This simplifies batching but introduces interpolation, anatomical deformation, and substantial preprocessing time into the learning pipeline.

Retaining native space appears to avoid these costs, yet omitting registration does not itself solve parameter sharing across subjects. Fixed grid locations are no longer reliable, and voxel counts vary; subject-specific mappings used in visual decoding methods such as MindEye2 do not directly provide a unified encoder for large cohorts. Finer voxels also produce long sequences, while reconstructing every raw signal is expensive and may overemphasize high-frequency noise. The paper therefore addresses three connected questions: how to share patch-encoding parameters across resolutions, how to retain location information without a regular grid, and how to make voxel-level self-supervised training affordable.

FlexiBrain changes the shared unit from a fixed array block to measurements covering a comparable physical extent. This aligns the physical meaning of receptive fields, not the precise anatomical region at a coordinate across subjects, and it does not turn low-resolution measurements into high-resolution data. A linear-complexity sequence model and latent prediction then reduce computation enough to make preserving native geometry a practical training choice. Core Idea: instead of resampling every subject's fMRI into one input format, adapt shared encoding kernels using millimeters and seconds, then learn transferable representations from physically positioned voxel patches.

Method

Overall Architecture

Inputs are minimally processed four-dimensional fMRI retaining its spatial and temporal resolution, together with metadata describing voxel spacing, sampling interval, and spatial transforms. Dynamic Physical Patches aggregate short-term signals before spatial projection; Foreground and Physical Position removes background patches and restores real-space information to the remaining sequence. Mamba Long-Sequence Encoding extracts features, while Asymmetric JEPA Prediction supplies a latent learning objective only during pretraining. Downstream fine-tuning and inference use the pretrained backbone and a classification head, without reconstructing masked raw voxels or requiring the teacher branch at inference time.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Native fMRI<br/>Resolution and spatial metadata"] --> Patch["Dynamic Physical Patches"]
    Patch --> Position["Foreground and Physical Position"]
    Position --> Backbone["Mamba Long-Sequence Encoding"]
    Backbone -->|Pretraining| Predict["Asymmetric JEPA Prediction<br/>MoE, predictor, and EMA targets"]
    Backbone -->|Downstream fine-tuning and inference| Output["Classification head<br/>Disease prediction"]

Key Designs

1. Dynamic Physical Patches: resize discrete kernels instead of standardizing the data grid

The model specifies a temporal extent \(\tau\) and physical sizes \(\rho_d\) along the three spatial axes, then computes how many samples its kernels must cover using the subject's sampling interval \(tr\) and voxel spacing \(v_d\). Following Equation (1), discrete dimensions use nearest-integer rounding:

\[ k_t=\operatorname{round}(\tau/tr),\qquad k_d=\operatorname{round}(\rho_d/v_d),\quad d\in\{x,y,z\}. \]

What is shared is therefore the base kernel parameters and intended physical coverage, not an identical kernel-array shape. Integer rounding makes the realized physical extent only approximately consistent across protocols; resolution agnosticism should not be interpreted as exact continuous-scale invariance. The model maintains temporal and spatial base kernels and adapts the kernel-resizing idea from FlexiViT, using the Moore-Penrose pseudo-inverse of a forward interpolation matrix to map base weights to the required dimensions. This differs from resampling the input image into template space and from initializing independent parameters for every resolution. The aim is to preserve the functional mapping of the filters as far as possible while allowing different protocols to update the same base parameters.

Tokenization proceeds through time first and space second. A one-dimensional temporal convolution with stride \(k_t\) aggregates short-term dynamics at each voxel, producing features with a reduced temporal length and intermediate channels. The reduced time dimension is folded into channels, and the spatial kernel is repeated along its input-channel dimension to accommodate those temporal blocks. The paper explicitly connects this operation to the assumption that spatial structure does not depend on temporal indexing; a three-dimensional convolution with matching spatial strides then produces the embeddings. The four-dimensional input should not be mistaken for a complete four-dimensional token grid entering the backbone: the described implementation folds time into channels before flattening spatial patches into a sequence. This decomposition concentrates cross-resolution sharing in two adaptable encoding kernels rather than introducing numerous subject-specific adapters inside the backbone.

2. Foreground and Physical Position: remove uninformative regions without losing patch locations

Voxel space contains substantial non-brain background, which lengthens sequences and wastes prediction targets. FlexiBrain uses a subject-specific foreground mask to discard purely background patches and flattens only valid patches. The paper reports removing approximately 80% of purely background patches while retaining around 4,000 tokens per subject, so foreground filtering does not eliminate the long-sequence problem. After filtering, neighboring sequence indices no longer adequately describe distances within the brain, and discrete positional indices are particularly difficult to reuse across grids. The model therefore reads the affine matrix and translation from the NIfTI header and converts patch centers from voxel indices to physical coordinates.

\[ c_i=p_iA^{\top}+t. \]

Here, \(p_i\) is the discrete center of patch \(i\), \(A\) and \(t\) define the subject-specific spatial transform, and \(c_i\) is its continuous three-dimensional physical position. Coordinates are scaled, encoded with sinusoidal features at logarithmically spaced frequencies along each axis, concatenated, and linearly projected to the token dimension before being added to patch embeddings. The model thus receives locations measured in millimeters rather than merely resolution-dependent array indices. This does not perform nonlinear anatomical registration across subjects: comparable physical coordinates and precise correspondence between brain regions remain distinct requirements. That distinction explains why preserving native space may retain individual information while still leaving scanner, pose, and acquisition-protocol variation to be addressed.

3. Mamba Long-Sequence Encoding: avoid quadratic attention costs for valid voxel sequences

Retaining thousands of spatial patches is the cost of fine-grained modeling; returning to ROI averages would discard information the framework is designed to retain. The backbone therefore uses the selective state-space model Mamba rather than global self-attention over the entire long sequence. The complexity comparison concerns \(O(N)\) versus \(O(N^2)\) scaling in sequence length, not a claim that every component of the complete system has linear cost. During pretraining, random masking divides valid tokens into visible context and masked targets, and the stacked Mamba context encoder processes only the visible subset. The target branch uses the same backbone structure to encode the actual masked tokens; its weights follow an exponential moving average (EMA) of the context backbone.

This reduces invalid inputs and work on the context path, but spatial relationships must still be learned through positions and sequence organization. Under a roughly 50M-parameter comparison, the paper reports 28.24 GFLOPs and 8.29 GB memory for ViT, versus 18.59 GFLOPs and 1.97 GB for Mamba. These are measurements for the configuration on page 11, not fixed speedup factors that apply to every sequence length and hardware platform. The downstream model is not entirely Transformer-free: a learnable class token is prepended to backbone features, processed by a lightweight Transformer classification head, and normalized before an MLP produces logits. The evidence therefore supports assigning the main encoding workload to Mamba rather than suggesting that all attention is ineffective.

4. Asymmetric JEPA Prediction: predict latent targets while reducing representation-collapse risk

Reconstructing every masked voxel would give high-resolution samples more raw-signal terms and could allow noise to dominate the objective. The joint-embedding predictive architecture (JEPA) instead predicts latent features generated by a target encoder, avoiding reconstruction of raw voxel arrays on every acquisition grid. The target backbone follows the context backbone through EMA rather than receiving direct prediction-loss updates, providing slowly changing targets. Because this arrangement can still yield low-diversity representations, the authors append a lightweight mixture-of-experts (MoE) module to the context encoder. A token-level router computes soft mixture weights and sends each visible token to a subset of MLP experts, adding content-dependent transformations on the context side. MoE complements the existing context/target asymmetry and EMA; it is neither a replacement teacher nor a universal proof against collapse.

Before prediction, MoE-processed visible features are scattered back to their original sequence positions, with learnable mask tokens inserted at masked locations. A shallow Mamba predictor propagates information across this restored sequence and produces predicted features at the masked locations. The objective on page 9 is the mean squared distance between separately unit-normalized predictions and teacher features; the following uses clear notation for the same objective:

\[ \mathcal{L}=\frac{1}{|\mathcal{M}|}\sum_{i\in\mathcal{M}} \left\|\frac{\hat H_i}{\|\hat H_i\|_2}-\frac{H_i}{\|H_i\|_2}\right\|_2^2. \]

Here, \(\mathcal{M}\) is the masked-position set, \(\hat H_i\) is the predictor output, and \(H_i\) is the EMA target feature. Normalization makes feature direction central to the loss rather than allowing vector magnitude alone to match the targets. The training chain therefore learns to infer missing-region representations from visible brain signals; it does not complete images at inference time. The paper supports MoE's empirical benefit through downstream ablations and target-token similarity maps, which should not be promoted into a mathematical guarantee that collapse is impossible.

A Worked Example

Using the recommended \(\tau=6\) seconds and \(\rho_d=12\) mm along every axis, consider an input with isotropic 2 mm voxels and a 2-second sampling interval. Equation (1) gives a temporal kernel covering 3 samples and a spatial kernel covering 6 voxels per axis. For another input with isotropic 3 mm voxels and a 3-second sampling interval, the dimensions become 2 time points and 4 voxels per axis. This is an illustration derived from the formula, not an additional experiment reported by the authors. The kernel arrays have different dimensions but target the same 6-second and 12 mm extents, using resized versions of the same base weights. After encoding, background removal and physical positions produce the token inputs; visible tokens enter the context backbone, while masked tokens supply supervision through the target path. Inference no longer constructs this random masked-prediction task and instead extracts backbone features for classification.

Loss & Training

Self-supervised pretraining combines ABIDE, ADHD-200, and ADNI, followed by fine-tuning for disease classification; PPMI is excluded from pretraining. ABIDE and ADHD-200 use binary classification, ADNI uses separate MCI-versus-CN and AD-versus-CN binary tasks, and PPMI uses three-class classification. The external PPMI result evaluates leave-one-dataset-out pretraining within a downstream adaptation workflow, not fully zero-shot diagnosis. The main text states that baselines follow their required preprocessing and that the subject set is restricted to match baselines requiring structural scans. Detailed hyperparameters, dataset information, and additional experiments are delegated to Appendices A–C; the supplied full-text cache contains only the main paper and references, so learning rates, masking ratios, and scanner-holdout scores cannot be supplied from it.

Key Experimental Results

Main Results

The following values come from Table 1 on page 10, with all metrics in percent and results reported as means and standard deviations over three random seeds. The baseline column selects the existing method with the highest accuracy for each task; its F1 belongs to that same method and is not necessarily the highest baseline F1 for the task.

Task	FlexiBrain ACC	FlexiBrain F1	Highest-ACC baseline	Baseline ACC	Baseline F1
ADHD-200	69.76 ± 1.75	69.85 ± 1.69	SwiFT	62.50 ± 1.50	59.99 ± 4.71
ABIDE	65.45 ± 0.69	64.82 ± 0.52	BrainMass	63.33 ± 1.18	65.35 ± 2.00
ADNI (MCI)	79.16 ± 1.48	78.82 ± 1.51	NeuroSTORM	66.67 ± 1.06	62.17 ± 8.93
ADNI (AD)	82.34 ± 2.39	79.41 ± 4.42	NeuroSTORM	84.26 ± 0.80	83.91 ± 0.75
PPMI	74.31 ± 1.77	70.54 ± 1.19	NeuroSTORM	68.25 ± 1.18	65.18 ± 1.19

The largest accuracy gain is on ADNI (MCI): 79.16 versus 66.67, a difference of 12.49 percentage points; the PPMI gain is 6.06 percentage points. However, ADNI (AD) accuracy is lower by 1.92 percentage points, and ABIDE F1 is below BrainMass, so the abstract's consistent-outperformance wording must be narrowed by metric. Asterisks in Table 1 indicate a large effect size with Cohen's \(d\geq0.8\), not a p-value threshold for a significance test.

Ablation Study

The following values come from Table 2 on page 12, all in percent; the table does not provide per-entry standard deviations. The four task columns are ACC except for the final explicitly labeled F1 column; rows vary backbone, pretraining, and MoE.

Config	ABIDE ACC	ADHD-200 ACC	ADNI (MCI) ACC	ADNI (AD) ACC	ADNI (AD) F1
Mamba, pretrained, no MoE	52.55	61.86	71.88	71.43	70.40
Mamba, no pretraining, with MoE	59.85	65.98	68.75	80.95	80.59
ViT, pretrained, with MoE	61.31	63.92	66.66	76.85	72.92
Mamba, pretrained, with MoE	65.45	69.76	79.16	82.34	79.41

Removing MoE lowers ABIDE ACC from 65.45 to 52.55, a drop of 12.90 percentage points, the largest decrease among these four ACC columns. Pretraining improves all four ACC values, but ADNI (AD) F1 changes from 80.59 without pretraining to 79.41, so it does not improve every metric. Replacing Mamba with ViT changes ADNI (MCI) ACC from 79.16 to 66.66; this supports the chosen backbone under the tested conditions, not the exclusion of all Transformer variants.

Key Findings

• Page 11 and Figure 3 describe approximately 3 minutes of native-space preprocessing versus about 5 hours for template processing; these are preprocessing costs, not end-to-end training times.
• The ADHD patch-size ablation in Figure 3 trains from scratch and selects 6 seconds and 12 mm; the cache does not retain every reliably readable curve value, so individual accuracies are not reconstructed.
• The main text reports degradation under stricter ABIDE leave-one-scanner-center-out evaluation, showing that accepting unseen resolutions does not solve unseen-center domain shift.

Highlights & Insights

• Handling resolution changes through geometric transformations of shared kernels rather than input-image resampling is the most reusable design. It preserves measurement grids while giving model parameters cross-protocol meaning.
• Background removal, linear sequence modeling, and latent objectives address distinct costs. Their combination reduces invalid tokens, backbone computation, and heterogeneous voxel-reconstruction burden, respectively.
• Equal physical extent does not imply equal information content. The controlled downsampling degradation described in the paper reminds readers that accepting lower-resolution inputs cannot restore details absent from those measurements.

Limitations & Future Work

• Native space is still not raw acquisition data. Page 11 describes minimal processing as brain extraction with motion and slice-timing parameter estimation, whereas page 15 gives motion correction as an example; the precise processing strength requires unavailable implementation details.
• Scanner-center domain shift remains unresolved, and the authors propose site harmonization, nuisance auditing, or domain-adaptive pretraining. These are necessary deployment questions rather than issues that resolution adaptation alone can be assumed to solve.
• The cache lacks Appendices A–C and supplementary tables, preventing independent verification of complete settings and values for resolution stratification, downsampling, gender prediction, and strict scanner-center holdout.
• From a reader's perspective, physical positional encoding does not guarantee anatomical correspondence, and better classification does not directly establish preservation of causal disease signals. More explicit confound control and prospective validation remain necessary.

Related Work & Insights

• Compared with FlexiViT: the paper adapts resizable patch kernels to fMRI scales measured in millimeters and seconds, combined with time-first, space-second encoding for heterogeneous sampling.
• Compared with Brain-JEPA / BrainMass: these methods establish the brain-representation pretraining context; FlexiBrain's focus is not reinventing latent prediction but admitting native voxel data to a shared model without template standardization.
• Compared with MindEye2: subject-specific brain-to-image mappings address individual visual decoding, whereas this paper develops a unified input interface for cross-cohort disease classification; their tasks and scaling bottlenecks differ.
• Research direction: test whether physical-unit encoding complements explicit scanner-center harmonization, reporting within-center, held-out-center, and controlled-resolution results separately; this is a reader-proposed follow-up experiment.

Rating

• Novelty: 4/5. Physical-unit shared kernels address a specific native-fMRI bottleneck, while the components build on existing methods.
• Experimental Thoroughness: 4/5. Five classification tasks and architecture ablations are included, but external-center robustness remains limited and supplementary experiments cannot be fully checked from this cache.
• Writing Quality: 4/5. The main argument and architecture are clear, though some overall-superiority and collapse-prevention wording exceeds the evidence.
• Value: 4/5. A concrete interface design for reducing template dependence and scaling voxel-level brain models, not a ready-to-deploy diagnostic system.