BrainVAE – Probabilistic fMRI Autoencoder

• BrainVAE is a probabilistic fMRI autoencoder that captures the distribution of variable yet functionally consistent neural responses to visual stimuli.
• It employs variational sampling with KL regularization and a strong CLIP-based semantic alignment to maintain visual semantics in the latent space.
• The architecture integrates convolutional encoders and attention modules to ensure stable high-dimensional fMRI reconstruction and effective visual-to-neural synthesis.

BrainVAE is the probabilistic fMRI autoencoding backbone inside SynBrain, introduced for visual-to-fMRI synthesis under the premise that a single visual stimulus does not evoke a single fixed cortical response but a distribution of plausible responses that vary across trials, contexts, and subjects while preserving shared functional and semantic structure. In that formulation, BrainVAE replaces deterministic visual-to-neural synthesis with a semantic-conditioned variational model whose latent uncertainty represents biological neural variability, and whose latent manifold is constrained to remain aligned with the paired image semantics through CLIP-based contrastive learning (Mai et al., 14 Aug 2025). In a broader and looser sense, several neuroimaging studies employ related VAE-style latent-variable models for 3D MRI generation, harmonization, regression, and multimodal brain-age estimation, although those works are introduced under distinct formal names rather than BrainVAE itself (Vogelsanger et al., 2021).

1. Definition and conceptual scope

In the strict sense documented by SynBrain, BrainVAE is not a generic label for any brain-imaging VAE; it is a specific probabilistic representation-learning module whose purpose is to capture the “variable but functionally consistent” nature of neural responses in visual cortex. The motivating observation is that repeated presentations of the same image can elicit noticeably different BOLD patterns because of trial noise, attentional fluctuations, and inter-individual variability, yet these responses are not arbitrary because they retain consistent functional and semantic structure (Mai et al., 14 Aug 2025).

This distinguishes BrainVAE from prior deterministic visual-to-fMRI synthesis methods. Deterministic models collapse the many possible responses associated with one stimulus into a single average pattern; BrainVAE instead models a stimulus as inducing a distribution over plausible fMRI patterns. Within SynBrain, that probabilistic latent structure is then used as the substrate for image-conditioned synthesis, few-shot adaptation to new subjects, and the generation of synthetic fMRI that can improve data-limited fMRI-to-image decoding (Mai et al., 14 Aug 2025).

The term also sits within a wider neuroimaging tradition of VAE-based latent representation learning. Related systems include volumetric MRI VAEs and Intro-VAEs for disease separability, hierarchical multi-modal VAEs for PET prediction, VQ-VAE-style brain-volume compression, supervised VAEs for age regression, InfoVAE variants for clinically informative MRI embeddings, and adversarial variational multimodal models for brain-age estimation (Vogelsanger et al., 2021, Lei et al., 2024, Tudosiu et al., 2020, Zhao et al., 2019, Huynh et al., 28 Sep 2025, Usman et al., 2024). This suggests that BrainVAE is best understood as a specific instantiation within a broader family of brain-imaging latent-variable models rather than as a universally standardized architecture.

2. Probabilistic formulation

BrainVAE is formulated as a variational autoencoder for fMRI signals. Given an input fMRI vector $y_{\mathrm{fMRI}}$ — denoted as $x$ in the appendix — the encoder produces a latent posterior distribution

$q(z \mid y_{\mathrm{fMRI}})$

parameterized by a mean $\mu$ and a log-variance $\log \sigma^2$. A latent sample is drawn as

$z \sim \mathcal{N}(\mu,\sigma^2),$

using the reparameterization trick, and the decoder reconstructs the signal as

$\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$

The prior is the standard normal

$p(z)=\mathcal{N}(0,I),$

and training regularizes the posterior toward that prior through a KL divergence term (Mai et al., 14 Aug 2025).

The training objective is a weighted sum of voxel reconstruction, KL regularization, and semantic alignment:

$$\mathcal{L}_{\mathrm{BrainVAE}} = \mathcal{L}_{\mathrm{MSE}} + \lambda_{\mathrm{KL}} \mathcal{L}_{\mathrm{KL}} + \lambda_{\mathrm{CLIP}} \mathcal{L}_{\mathrm{CLIP}}.$$

The reconstruction term is voxel-wise mean squared error,

$$\mathcal{L}_{\mathrm{MSE}} = \lVert \mathcal D(z) - y_{\mathrm{fMRI}} \rVert_2^2,$$

and the KL term is

$x$0

The semantic regularizer is a SoftCLIP contrastive loss,

$x$1

where $x$2 is the CLIP embedding of the paired image $x$3, extracted by a frozen CLIP visual encoder. The reported hyperparameters are $x$4 and $x$5, with semantic alignment heavily weighted and KL regularization kept soft to avoid harming reconstruction (Mai et al., 14 Aug 2025).

The functional interpretation of these terms is explicit. The KL term smooths the latent space and supports sampling from a well-behaved posterior family; the CLIP term grounds the latent neural representation in visual semantics, making the latent fMRI space not only reconstructive but also semantically aligned with the stimulus. In SynBrain’s framing, BrainVAE therefore learns a latent neural manifold that is simultaneously sampleable, semantically organized, and biologically variable (Mai et al., 14 Aug 2025).

3. Architecture and latent geometry

BrainVAE was designed to be more stable than the MLP-based VAE baseline, denoted MLP-VAE and borrowed from or inspired by MindSimulator. The reported failure modes of MLP-VAE are unstable training, collapsed reconstructions, and diverging MSE, attributed to the lack of spatial inductive bias and token-wise independence in MLPs. BrainVAE instead uses a convolutional encoder plus attention modules to better capture fMRI structure (Mai et al., 14 Aug 2025).

The encoder processes input fMRI $x$6 with an initial 1D convolution from $x$7 channel to $x$8 channels, kernel size $x$9, padding $q(z \mid y_{\mathrm{fMRI}})$0, followed by adaptive max pooling to length $q(z \mid y_{\mathrm{fMRI}})$1. It then applies a hierarchical ResNet-style backbone with three resolution stages, channel multipliers $q(z \mid y_{\mathrm{fMRI}})$2, two residual blocks per stage, one downsampling block, and a middle block with self-attention. A final LayerNorm, SiLU, and 1D convolution produce an encoder feature tensor

$q(z \mid y_{\mathrm{fMRI}})$3

Two parallel pre-projector MLPs produce $q(z \mid y_{\mathrm{fMRI}})$4 and $q(z \mid y_{\mathrm{fMRI}})$5, each using LayerNorm, GELU, and linear layers, and project to a latent size of $q(z \mid y_{\mathrm{fMRI}})$6. After sampling, a symmetric post-projector maps the latent back to $q(z \mid y_{\mathrm{fMRI}})$7, and the decoder reconstructs the subject-specific fMRI vector $q(z \mid y_{\mathrm{fMRI}})$8. The appendix states that the decoder mirrors the encoder with ResNet-style upsampling, self-attention in the middle block, and linear interpolation to match each subject’s voxel count (Mai et al., 14 Aug 2025).

The algorithmic summary is:

$q(z \mid y_{\mathrm{fMRI}})$9

Optimization uses

$\mu$0

This design couples high-dimensional fMRI reconstruction with a latent distribution whose uncertainty is intended to encode trial-level response variability (Mai et al., 14 Aug 2025).

A latent-dimensionality ablation varying the number of down blocks shows that one down block, corresponding to latent dimensionality $\mu$1, is preferred to deeper compression.

Down blocks	Latent dimensionality	Reported result
$\mu$2 block	$\mu$3	MSE $\mu$4, Pearson $\mu$5, Cosine $\mu$6, Inception $\mu$7, CLIP $\mu$8, Syn retrieval $\mu$9
$\log \sigma^2$0 blocks	$\log \sigma^2$1	slightly worse across all metrics
$\log \sigma^2$2 blocks	$\log \sigma^2$3	further degradation

The authors interpret this as evidence that preserving higher-dimensional latent structure is important for keeping fine-grained temporal or voxel information and semantic fidelity (Mai et al., 14 Aug 2025).

4. Role inside SynBrain

BrainVAE is the component that creates the semantic neural manifold on which the rest of SynBrain operates. The contrastive CLIP term aligns latent fMRI codes with CLIP image embeddings, so the latent space becomes semantically organized according to the visual stimulus space rather than being a purely reconstructive bottleneck. This is presented as the mechanism that makes the latent space interpretable and controllable, ensuring that generated responses preserve stimulus semantics (Mai et al., 14 Aug 2025).

Within the broader pipeline, the Semantic-to-Neural Mapper functions as a semantic transmission pathway that learns a point-to-distribution mapping from CLIP embeddings into the BrainVAE latent space. At inference time, an image’s CLIP representation is mapped directly into a plausible fMRI latent code and then decoded. In this division of labor, BrainVAE creates the latent neural response manifold, while the S2N Mapper learns how to reach that manifold from vision (Mai et al., 14 Aug 2025).

This latent design is also the formal mechanism by which SynBrain maintains functional consistency while modeling variability. The KL term prevents arbitrary latent fragmentation and supports smooth sampling; the CLIP term keeps samples semantically tied to the image. The qualitative and analysis sections argue that the resulting model can generate diverse but functionally equivalent responses: voxel-level patterns may vary, but category-selective regions remain consistently activated in ways that reflect the stimulus. In the “diverse synthesis” experiment, adding Gaussian noise around the mapped latent center produces multiple plausible fMRI samples that still decode to semantically consistent images, indicating that the learned latent space is structured rather than chaotic (Mai et al., 14 Aug 2025).

5. Empirical behavior and ablation evidence

The strongest experimental evidence for BrainVAE’s contribution comes from ablations against deterministic autoencoding and from comparisons with MindSimulator-style baselines. Removing variational sampling and turning SynBrain into a deterministic autoencoder causes a substantial drop in semantic metrics: Inception falls from $\log \sigma^2$4 to $\log \sigma^2$5, CLIP from $\log \sigma^2$6 to $\log \sigma^2$7, and retrieval from $\log \sigma^2$8 to $\log \sigma^2$9. This is the central empirical argument that probabilistic modeling itself matters, rather than only the downstream mapper (Mai et al., 14 Aug 2025).

Additional ablations isolate other components. Removing contrastive learning causes retrieval to collapse almost completely, indicating that BrainVAE alone can learn some structure but does not automatically align with semantic space. Removing the S2N Mapper also substantially hurts performance, showing that the BrainVAE latent space requires explicit semantic-to-neural projection to be fully effective in synthesis. BrainVAE-DiT, in which a diffusion transformer replaces the S2N Mapper while keeping the same VAE backbone, underperforms SynBrain, supporting the claim that the backbone works best when paired with a direct probabilistic semantic mapper rather than iterative denoising from noise (Mai et al., 14 Aug 2025).

The comparison with MindSimulator provides the clearest subject-specific synthesis numbers. SynBrain with BrainVAE reaches MSE $z \sim \mathcal{N}(\mu,\sigma^2),$0 and Pearson $z \sim \mathcal{N}(\mu,\sigma^2),$1, whereas MindSimulator yields MSE $z \sim \mathcal{N}(\mu,\sigma^2),$2 and Pearson $z \sim \mathcal{N}(\mu,\sigma^2),$3 at one trial, with corresponding gains in semantic-level metrics and retrieval. The paper also attributes the inferiority of the MLP-VAE baseline to unstable training, collapsed reconstructions, and diverging MSE, reinforcing the architectural claim that convolution and attention are important for stable variational training on high-dimensional fMRI (Mai et al., 14 Aug 2025).

Taken together, these results position BrainVAE not as a generic VAE backbone added for regularization, but as the specific mechanism enabling one-shot generation of diverse, semantically constrained, and biologically variable fMRI responses. A plausible implication is that SynBrain’s performance gain derives from the interaction of three elements rather than any single one in isolation: probabilistic latent sampling, strong visual-semantic alignment, and an explicit semantic-to-neural projection stage.

6. Relation to other brain-imaging latent-variable models

Outside SynBrain, closely related models show how VAE-family methods have been adapted to distinct neuroimaging objectives. A 3D VAE and Intro-VAE trained on whole FLAIR volumes from healthy controls, multiple sclerosis, and microvascular leukoencephalopathy use unsupervised latent embeddings followed by LDA to examine disease separability. In that study, the Intro-VAE produces sharper and more diverse samples than the standard VAE, the IVAE-256 latent space yields precision $z \sim \mathcal{N}(\mu,\sigma^2),$4 and recall $z \sim \mathcal{N}(\mu,\sigma^2),$5 for MS detection, and latent traversals reveal factors affecting ventricle shape and intensity; however, the authors explicitly note that acquisition-parameter differences may confound latent-space separation (Vogelsanger et al., 2021).

InVA extends the VAE framework in a different direction by introducing a two-level hierarchical architecture for harmonization of multi-modal neuroimaging and image-on-image prediction. For each modality, a modality-specific shallow encoder produces Gaussian latent features, these are passed through a shared deep encoder to support information borrowing across modalities, and a predictor head infers an output image such as amyloid-β PET SUVR from MRI cortical thickness and cortical volume. On ADNI data, the reported MSE is $z \sim \mathcal{N}(\mu,\sigma^2),$6 for InVA with both modalities, compared with $z \sim \mathcal{N}(\mu,\sigma^2),$7 for VAE with cortical thickness, $z \sim \mathcal{N}(\mu,\sigma^2),$8 for VAE with cortical volume, and $z \sim \mathcal{N}(\mu,\sigma^2),$9 for BART with both modalities (Lei et al., 2024).

A distinct line of work replaces continuous Gaussian latents with discrete codebooks. The volumetric VQ-VAE-style model for full-resolution 3D brain MRI uses 3D convolutions, three VQ bottlenecks, FixUp residual blocks, and VBM- and segmentation-based evaluation of morphology preservation. It reports compression to $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$0 of original bit size while maintaining fidelity, high Dice overlap for white matter, grey matter, and CSF, and statistically significant improvements over an $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$1-WGAN baseline with Wilcoxon signed-rank test $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$2 for Dice scores (Tudosiu et al., 2020).

Other variants couple latent-variable modeling to supervision. A regression-oriented VAE for brain aging replaces the unconditional latent prior with an age-conditioned prior $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$3, jointly regularizing reconstruction and age prediction. On 245 healthy subjects, it reports $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$4 and rMSE $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$5 for ROI measures, and $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$6 and rMSE $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$7 for ventricle-centered 3D images, with latent traversals showing ventricle enlargement with age (Zhao et al., 2019). InfoVAE-Med3D instead emphasizes mutual-information preservation in 3D MRI embeddings, using a 512-dimensional latent vector and downstream SVR; it reports, for example, MAE $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$8, $\hat{y}_{\mathrm{fMRI}} = \mathcal{D}(z).$9, and RMSE $p(z)=\mathcal{N}(0,I),$0 on the BrainAge dataset, as well as clear gender clustering and age gradients in latent-space projections (Huynh et al., 28 Sep 2025). A multimodal adversarial variational autoencoder further combines sMRI and fMRI, decomposes latent codes into shared generic and unique modality-specific parts, adds cross reconstruction and multitask age regression plus sex classification, and reports mean absolute error $p(z)=\mathcal{N}(0,I),$1 years on OpenBHB (Usman et al., 2024).

Relative to these systems, BrainVAE is unusual in that its primary object is not anatomical MRI, PET harmonization, or age regression, but a probabilistic and semantically constrained manifold for fMRI synthesis from vision. This suggests that its closest conceptual relatives are not simply generic neuroimaging VAEs, but latent-variable models in which uncertainty, semantic organization, and downstream controllability are all first-class design objectives.

7. Misconceptions, caveats, and open questions

One common misconception is that VAE-based brain models are merely compression devices. BrainVAE is explicitly not limited to compression: its latent space is shaped by voxel reconstruction, KL regularization, and SoftCLIP semantic alignment, and it is used as the basis for image-to-fMRI synthesis, few-shot subject adaptation, and synthetic-data augmentation for downstream decoding (Mai et al., 14 Aug 2025). Related neuroimaging VAEs likewise serve broader purposes, including disease separability, image-on-image prediction, supervised regression, and multimodal fusion (Lei et al., 2024, Zhao et al., 2019, Usman et al., 2024).

A second misconception is that probabilistic neural synthesis is equivalent to uncontrolled randomness. In BrainVAE, variability is bounded by two explicit constraints: KL regularization toward $p(z)=\mathcal{N}(0,I),$2, which smooths the latent space and makes it sampleable, and CLIP-based semantic grounding, which ties latent codes to the paired image semantics. The intended result is not arbitrary diversity but diverse yet functionally equivalent responses (Mai et al., 14 Aug 2025).

The broader literature also identifies important limitations. In whole-brain MRI latent-space analysis, disease separability may partly reflect acquisition bias rather than pathology, because metadata such as pixel bandwidth, repetition time, and echo time can correlate with class labels (Vogelsanger et al., 2021). VAE-family reconstructions may also remain blurry or artifact-prone even when improved variants are used: the standard VAE in volumetric FLAIR analysis is described as very blurry, the Intro-VAE as sharper but still artifact-laden, and InfoVAE-Med3D, despite better SSIM and PSNR than AE, VAE, and $p(z)=\mathcal{N}(0,I),$3-VAE baselines, still exhibits the usual VAE blurriness (Vogelsanger et al., 2021, Huynh et al., 28 Sep 2025).

These caveats matter for interpreting BrainVAE as well. SynBrain reports that synthesized signals capture interpretable patterns shaped by biological neural variability, but the wider literature suggests that interpretability claims in latent neuroimaging models must always be read alongside possible confounds from architecture, acquisition, and objective design. A plausible implication is that the most robust future BrainVAE-like systems will continue to combine explicit inductive bias, strong semantic or clinical supervision, and careful evaluation of whether latent structure reflects biology rather than nuisance variation.