Source-linked AI summary

Can Natural Image Autoencoders Compactly Tokenize fMRI Volumes for Long-Range Dynamics Modeling?

Peter Yongho Kim, Juhyeon Park, Jungwoo Park, Jubin Choi, Jungwoo Seo, Jiook Cha, Taesup Moon

arXiv:2604.03619v1cs.CV

TL;DR

Long-range fMRI dynamics are difficult to model because voxel-based representations are memory-intensive and restrict temporal context. TABLeT uses a pre-trained natural-image 2D autoencoder to tokenize volumes for lightweight Transformer modeling, achieving competitive or better performance across three datasets while reducing memory and computation; masked-token pre-training further improves downstream results.

  • Problem

    Voxel-based fMRI models face prohibitive memory demands that restrict the temporal windows available for modeling long-range spatiotemporal dynamics.

  • Method

    TABLeT uses a pre-trained natural-image 2D autoencoder to convert each fMRI volume into compact tokens for a lightweight Transformer, with masked token modeling for self-supervised pre-training.

  • Results

    Across three datasets, TABLeT achieves competitive or better performance while substantially reducing memory and computation relative to voxel-based baselines.

  • Takeaways & Limitations

    TABLeT provides a scalable approach for efficient spatiotemporal brain-activity modeling and enables much longer temporal sequences.

  • Takeaways & Limitations

    TABLeT tokenizes frames independently, which may disrupt subtle temporal dynamics, and its benefits from longer temporal modeling vary across tasks.

Abstract

from arXiv · show

Modeling long-range spatiotemporal dynamics in functional Magnetic Resonance Imaging (fMRI) remains a key challenge due to the high dimensionality of the four-dimensional signals. Prior voxel-based models, although demonstrating excellent performance and interpretation capabilities, are constrained by prohibitive memory demands and thus can only capture limited temporal windows. To address this, we propose TABLeT (Two-dimensionally Autoencoded Brain Latent Transformer), a novel approach that tokenizes fMRI volumes using a pre-trained 2D natural image autoencoder. Each 3D fMRI volume is compressed into a compact set of continuous tokens, enabling long-sequence modeling with a simple Transformer encoder with limited VRAM. Across large-scale benchmarks including the UK-Biobank (UKB), Human Connectome Project (HCP), and ADHD-200 datasets, TABLeT outperforms existing models in multiple tasks, while demonstrating substantial gains in computational and memory efficiency over the state-of-the-art voxel-based method given the same input. Furthermore, we develop a self-supervised masked token modeling approach to pre-train TABLeT, which improves the model's performance for various downstream tasks. Our findings suggest a promising approach for scalable and interpretable spatiotemporal modeling of brain activity. Our code is available at https://github.com/beotborry/TABLeT.

1. Introduction

fMRI modeling must capture long-range brain dynamics despite the high dimensionality of voxel data and severe memory constraints. TABLeT addresses this by compressing each 3D volume into 27 continuous tokens with a transferred 2D natural-image autoencoder, enabling longer Transformer sequences and efficient modeling.

  • Motivation: Voxel-based fMRI models are limited to short temporal windows by the massive scale of 4D volumes.Examples cited use only 20 timesteps simultaneously, potentially missing informative long-range dynamics.
  • Approach: TABLeT asks whether a natural-image 2D autoencoder can effectively tokenize 4D fMRI while preserving functional information.The approach is motivated by DCAE's compression of 256×256 images into 64 continuous tokens.
  • Approach: Each 3D fMRI volume is rearranged from slice-level representations into only 27 continuous tokens for Transformer processing.This compact representation reduces input size and supports substantially longer temporal sequences with limited VRAM.
  • Results: TABLeT achieves consistent yet modest improvements over voxel-based baselines on demographic attribute prediction and ADHD diagnosis.The contribution summary also reports performance improvements across multiple resting-state fMRI tasks.
  • Results: The model substantially reduces memory and computation costs while masked token modeling further improves downstream performance.Self-supervised pre-training is performed on tokenized fMRI data.

2. Related Work

Prior fMRI methods either use efficient but information-losing ROI preprocessing or retain voxel-level information at much higher computational cost. TABLeT instead tokenizes fMRI frames with a 2D autoencoder before Transformer processing.

  • ROI-Based Methods: ROI-based methods convert parcellated BOLD signals into functional-connectivity matrices for efficient modeling.Examples include BrainNetCNN, Brain Network Transformer, meanMLP, and Brain-JEPA.
  • ROI-Based Methods: ROI preprocessing depends strongly on region definitions and can discard structural information and other signals.This is the main limitation identified for ROI-based approaches.
  • Voxel-Based Methods: Voxel-based methods operate on high-dimensional brain volumes, creating substantial computational and memory demands.The supplied related-work passage is truncated after introducing voxel-based methods, so no further method-specific comparison is stated.
  • TABLeT: TABLeT processes each fMRI frame through a 2D autoencoder and then sends the resulting tokens to a Transformer.This pipeline is presented as an alternative to ROI-based functional-connectivity preprocessing.

3. Method

TABLeT uses a pre-trained 2D natural-image autoencoder to tokenize fMRI volumes slice-wise, then models the resulting sequences with a Transformer. The method also introduces masked token modeling for self-supervised pre-training.

  • 3. Method: Directly using a natural-image-trained autoencoder avoids computationally prohibitive, data-hungry fMRI-specific autoencoder training.The authors select DCAE because it maintains reconstruction fidelity at high compression rates.
  • 3. Method: 2D DCAE preserves fine-grained spatial details and global functional patterns comparably to a 3D DCAE trained directly on fMRI.This comparison motivates slice-wise tokenization with the off-the-shelf 2D autoencoder.
  • 3.1. Tokenization of fMRI with 2D Natural Image Autoencoder: Each fMRI volume is channel-duplicated, sliced along one spatial axis, and compressed independently into latent representations.The procedure is repeated across the three slicing axes before latent aggregation.
  • 3.1. Tokenization of fMRI with 2D Natural Image Autoencoder: Each 96×96×96 volume becomes 27 tokens with embedding dimension 3072 after aggregation across slicing axes.Tokenization is performed once and cached for later training.
  • 3.2. Transformer Architecture: TABLeT applies a Transformer encoder with grouped query attention to model spatiotemporal dynamics in the tokenized sequences.The architecture is designed to handle long sequences efficiently.
  • 3.3. Self-supervised Pre-training with Masked Token Modeling: Masked token modeling randomly replaces tokenized representations with [MASK] tokens and trains the Transformer to predict them.The default masking ratio is 0.5, and the same masking pattern is repeated across frames to prevent leakage from corresponding unmasked locations.

4. Experimental Results

Across four tasks and three datasets, TABLeT outperformed baseline methods while reducing memory and computation demands. Its 2D DCAE tokenizer preserved functional information, and longer temporal contexts and masked pre-training further supported performance, although gains over SwiFT were modest in resting-state settings.

  • TABLeT outperformed both ROI-based and voxel-based baselines across four tasks and three datasets, with marginal gains on HCP-Age and competitive UKB-Sex performance.
  • Performance showed a positive association with input time frames for HCP-Intelligence and ADHD diagnosis, suggesting benefits from modeling longer temporal variability.
  • The 2D DCAE moderately outperformed the 3D DCAE and better preserved functional connectivity patterns in reconstruction comparisons.The comparison assessed voxel-level fidelity and functional-connectivity preservation; downstream performance was also comparable, with 2D DCAE moderately ahead in most cases.
  • Fine-tuning the 2D DCAE on fMRI consistently harmed reconstruction quality, so the authors advocate using it as a training-free tokenizer.The authors attribute this result to the relatively small and homogeneous fMRI dataset potentially harming generic filters.
  • 7.33× lower memory use and 3.8× faster training than SwiFT were achieved at T = 50 on a single RTX A6000 with batch size 4.SwiFT could only run up to T = 50 under the stated memory limit.

5. Conclusion & Limitations

TABLeT uses a natural-image 2D autoencoder to tokenize fMRI volumes for scalable long-range modeling, achieving competitive or better performance with lower memory and computation. Its limitations include independently tokenized frames, absent explicit token structure, and task-dependent benefits from longer temporal modeling.

  • TABLeT uses a natural-image 2D autoencoder as a training-free tokenizer and models the resulting fMRI tokens with a lightweight Transformer.
  • Across three datasets, TABLeT achieved competitive or better performance while substantially reducing memory and computation relative to voxel-based baselines.
  • Masked token pre-training further enhanced downstream performance.
  • Limitations: TABLeT tokenizes each fMRI frame independently, which may disrupt subtle temporal dynamics.
  • Limitations: TABLeT jointly processes tokens without explicitly modeling their spatial or temporal structure, and longer temporal modeling benefits vary across tasks.

A. Implementation Details

The experiments use shared AdamW optimization, task-specific validation criteria, early stopping, mixed precision where stable, and task-appropriate losses. Baseline and TABLeT training use specified learning-rate, batch-size, epoch, and window-sampling procedures across GPUs.

  • Hardware and Precision: Training used fp16 mixed precision for all models except TFF, whose training produced NaN errors.
  • Losses: Classification used BCEWithLogitsLoss, ADHD used positive weighting for class imbalance, and regression used L1Loss.
  • Sampling and Evaluation: Voxel-based models trained on randomly sampled consecutive 3D-volume windows and averaged outputs over all evaluation windows.
  • Shared Settings: All experiments used AdamW with cosine decay and weight decay of 10^-2, with hyperparameters selected using task-specific validation metrics.
  • TABLeT: TABLeT training used learning rates from 3 × 10^-7 to 5 × 10^-5 and task-specific epoch schedules on NVIDIA A100-40GB and RTX A6000 GPUs.The cited setup lists separate epoch counts for HCP-Sex, HCP-Intelligence, ADHD, age regression, and UKB-Sex.

B. Training Details of 3D fMRI-trained DCAE

The 3D DCAE adapts the 2D DCAE architecture for volumetric fMRI inputs and is trained on UK-Biobank frames with reconstruction loss. Training used extensive convergence-oriented procedures for fair comparison.

  • 3D DCAE Architecture: The 3D DCAE replaces 2D operations with 3D counterparts to process volumetric fMRI inputs.Its configuration uses one input channel, 1024 latent channels, and specified encoder-decoder widths and depths.
  • Training: Training used 8,178 UK-Biobank subjects, a 9:1 stratified split, 100 epochs, and L2 reconstruction loss.One randomly selected fMRI frame per subject was used in each epoch.
  • Training: The training curve was used to verify efforts toward best performance and full convergence for fair comparison.

C.1. Experiments with Matched T

Matched-T experiments on HCP and ADHD-200 show that TABLeT remains competitive with SwiFT even when using fewer input frames. This indicates its gains are not solely attributable to a larger temporal input length.

  • Matched-T Results: TABLeT showed comparable performance to SwiFT on HCP and ADHD-200 with matched T values of 20 and 50.
  • Matched-T Results: The results indicate that TABLeT's performance gain is not solely due to using a larger number of input frames.

C.2. HBN-Movie Experiments

HBN-Movie tests whether TABLeT can model temporally dynamic fMRI through movie-viewing prediction. TABLeT matches SwiFT at the same temporal length and outperforms it with longer temporal windows, while token aggregation trades token count for dimensionality at fixed total values.

  • HBN-Movie Experiments: 680 subjects and 1,360 fMRI scans from HBN-Movie were used to predict which of two watched movies each subject viewed.Each subject watched two movies, providing a temporally dynamic classification task.
  • HBN-Movie Experiments: TABLeT achieves performance comparable to SwiFT with the same temporal length T on HBN-Movie.The comparison is reported in Table A6.
  • HBN-Movie Experiments: TABLeT outperforms SwiFT with longer T, demonstrating its ability to handle fMRI temporal dynamics.This conclusion is drawn from the HBN-Movie results.
  • HBN-Movie Experiments: At T = 256, token-count and latent-dimensionality variations keep the total representation size fixed at 27 × 3072 = 82,944 values.The aggregation scheme combines tokens while preserving the total number of values.
  • HBN-Movie Experiments: Combining more tokens increases token dimensionality while reducing the total number of resulting tokens.Token aggregation concatenates and rearranges tokens without changing the total number of values.

D. Detailed Experimental Results

The appendix reports experimental results with standard deviations across UKB, HCP sex classification and age regression, and HCP intelligence regression and ADHD diagnosis.

  • Detailed Experimental Results: Experimental results with standard deviation are reported for the UKB benchmark.These results are provided in Table A8.
  • Detailed Experimental Results: Experimental results with standard deviation cover HCP sex classification and age regression.These results are provided in Table A9.
  • Detailed Experimental Results: Experimental results with standard deviation cover HCP intelligence regression and ADHD diagnosis.These results are provided in Table A10.

E. Detailed Data Description

The appendix provides detailed dataset descriptions and demographic information for the datasets used in the study, including HCP and ADHD-200 experimental coverage.

  • Detailed Data Description: Detailed descriptions of each dataset used in the study are provided in Table A11.The table is identified as containing demographic information for the study datasets.
  • Detailed Data Description: Table A10 covers HCP intelligence regression and ADHD diagnosis among the reported experiments.This identifies two datasets or task settings represented in the detailed experimental results.
Loading 2604.03619v1…