flow-matching-1

Flow Matching fMRI Encoder

A two-stage architecture for predicting fMRI BOLD responses from naturalistic video stimuli (Friends TV show + Movie10 dataset) using Conditional Flow Matching.

---

Overview

The pipeline decodes brain activity in two sequential stages:

The design mirrors the MedARC approach: Stage 1 provides a stable conditional mean $\mu$; Stage 2 integrates a continuous normalizing flow $\phi_t$ conditioned on $\mu$ to sample from the true posterior over voxel activations.

---

Architecture

Stage 1 — Mean Anchor (medarc_architecture.py)

Features (N, T, D_i)  →  DepthConv1d + Linear  →  embed (N, T, d)
                                                         ↓
                                               [Shared Decoder]  +  [Subject Decoders]
                                                         ↓
                                               fMRI Prediction (N, S, T, V)

• MultiSubjectConvLinearEncoder: Projects each feature stream through a LinearConv (depthwise conv + linear) to a shared embedding dimension d = 192.
• Global average pooling across the feature stack combines multi-model features.
• A shared linear decoder and per-subject linear decoders combine additively to predict V voxels for each of S subjects.
• Trained with MSE loss against ground-truth BOLD.

Stage 2 — Neural Vector Field (matcha_architecture.py)

x1 (B, V, T)  →  proj_in (V → d)  →  latent x1 (B, d, T)
mu (B, V, T)  →  proj_in (V → d)  →  latent mu (B, d, T)
                                              ↓
                              OT-CFM loss (vector field u_t)
                              Matcha-TTS U-Net estimator
                              [Conformer / Transformer blocks]
                                              ↓
                              latent pred (B, d, T)  →  proj_out (d → V)  →  fMRI (B, V, T)

• Latent Bottleneck: fMRI voxels (V ≈ 1000) are projected down to a dense latent dimension (d = 128) before any convolution, reducing the first-layer parameter count from ~6M to ~98K and preventing gradient collapse.
• CFM wraps a Matcha-TTS style 1D U-Net (Decoder) with ResNet-1D blocks and Conformer/Transformer attention at each scale.
• At inference, noise $z \sim \mathcal{N}(0,I)$ is integrated from $t=0$ to $t=1$ over 25 Euler steps conditioned on $\mu$.
• An auxiliary reconstruction loss (weight 0.1) on proj_in → proj_out trains the projection pair jointly with the vector field.

---

Data

Subjects used: 1, 2, 3, 5.

Feature Models

The encoder can ingest intermediate activations from any combination of:

Active features are set in config.yml under include_features.

---

File Structure

flow_matching/
├── config.yml               # Main training config (GPU, full data)
├── debug_config.yml         # Fast local debug config (CPU, tiny data)
├── environment.yml          # Conda environment spec
│
├── src/
│   ├── training.py          # Two-stage training loop + evaluation
│   ├── matcha_architecture.py  # CFM + Matcha-TTS U-Net decoder
│   ├── medarc_architecture.py  # Stage 1 MultiSubjectConvLinearEncoder
│   ├── data.py              # Algonauts2025 dataset + loaders
│   ├── metric.py            # Pearson's r voxel-wise scoring
│   ├── visualize.py         # Loss curve plotting
│   └── inference.py         # Standalone inference helper
│
├── test/
│   ├── overfit_test.py      # Tiny-batch overfit sanity check for Stage 2
│   ├── check_pearson.py     # Load checkpoints and plot per-voxel Pearson's r heatmaps
│   └── debug_training.py    # End-to-end smoke test
│
├── experiments/
│   └── *.ipynb              # Analysis notebooks (RSA, OOD, brain region plots)
│
└── Matcha-TTS/              # Vendored Matcha-TTS source (U-Net + solver)

---

Training

Full training (server)

cd flow_matching
python src/training.py --cfg-path config.yml

Checkpoints are written to output/two_stage_encoding/:

• stage1_best.pt — best Stage 1 model by validation Pearson's r
• stage2_epoch_N.pt — Stage 2 snapshot every 5 epochs

Local debug (CPU, tiny model)

python src/training.py --cfg-path debug_config.yml

---

Evaluation

Pearson's r heatmaps

Loads all available Stage 1 and Stage 2 checkpoints, evaluates on the configured validation set, and saves per-subject per-voxel Pearson's r heatmaps to output/two_stage_encoding/heatmaps/.

python test/check_pearson.py

Output per checkpoint:

Stage 1 Overall Pearson's r: 0.1832
Stage 1 - Sub 1 Mean Pearson's r: 0.1754
Stage 2 Epoch 5 Overall Pearson's r: 0.2110
Stage 2 Epoch 5 - Sub 1 Mean Pearson's r: 0.2043
...

Tiny-batch overfit test

Confirms Stage 2 can memorize a single training batch. If loss does not approach 0 within 500 steps, the architecture cannot learn the task.

python test/overfit_test.py --cfg-path config.yml --subject-idx 0 --steps 500

---

Key Hyperparameters

---

Metric

Evaluation uses voxel-wise Pearson's r averaged across subjects:

$$r_v = \frac{\sum_t (y_v^t - \bar{y}_v)(\hat{y}_v^t - \bar{\hat{y}}_v)}{|\mathbf{y}_v| \cdot |\hat{\mathbf{y}}_v|}$$

The scalar reported is the mean over all V voxels and all S subjects.