sabertoaster/flow-matching

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:

Stage	Model	Goal
1	MultiSubjectConvLinearEncoder	Predict a Mean Anchor — a deterministic per-voxel fMRI estimate shared across subjects
2	CFM (Conditional Flow Matching)	Learn a per-subject neural vector field that refines the Mean Anchor into a sharper, stochastic fMRI prediction

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

Source	Content	Usage
Friends (seasons 1–7)	fMRI BOLD + multimodal features	Train (S1), Val (S6, S7)
Movie10	fMRI BOLD + multimodal features	Supplementary val (Figures, Life, Bourne, Wolf)

Subjects used: 1, 2, 3, 5.

Feature Models

The encoder can ingest intermediate activations from any combination of:

Key	Model
internvl3_8b, internvl3_14b	InternVL3 vision-language model
qwen-2-5-omni-3b, qwen-2-5-omni-7b	Qwen2.5-Omni audio-video model
whisper	OpenAI Whisper (audio)
llama_3.2_1b, llama_3.2_3b	LLaMA 3.2 (text)
vjepa2	V-JEPA 2 (video)

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

Parameter	Value	Location
Stage 1 embed dim	192	config.yml / stage1.model.embed_dim
Stage 1 encoder kernel	45	config.yml / stage1.model.encoder_kernel_size
Stage 1 LR	3e-4	config.yml / stage1.lr
Stage 2 latent dim	128	config.yml / stage2.latent_dim
Stage 2 U-Net channels	[256, 256]	config.yml / stage2.decoder.channels
Stage 2 block type	Conformer	config.yml / stage2.decoder.*_block_type
Stage 2 LR	3e-4	config.yml / stage2.lr
Euler integration steps	25	config.yml / stage2.n_timesteps
CFM σ_min	1e-4	config.yml / stage2.cfm.sigma_min

---

Metric

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

$$r_v=\frac{\sum _t(y_v^t-{\bar{y}}_v)({\widehat{y}}_v^t-{\bar{\widehat{y}}}_v)}{\mathrm{\parallel }{\mathbf{y}}_v\mathrm{\parallel }\cdot \mathrm{\parallel }{\widehat{\mathbf{y}}}_v\mathrm{\parallel }}r\_v\; =\; \backslash frac\{\backslash sum\_t\; (y\_v^t\; -\; \backslash bar\{y\}\_v)(\backslash hat\{y\}\_v^t\; -\; \backslash bar\{\backslash hat\{y\}\}\_v)\}\{\backslash |\backslash mathbf\{y\}\_v\backslash |\; \backslash cdot\; \backslash |\backslash hat\{\backslash mathbf\{y\}\}\_v\backslash |\}$$

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