Modality-Aware and Anatomical Vector-Quantized Autoencoding for Multimodal Brain MRI

Mingjie Li; Edward Kim; Yue Zhao; Ehsan Adeli; Kilian M. Pohl

cs.CVApr 6, 2026

Modality-Aware and Anatomical Vector-Quantized Autoencoding for Multimodal Brain MRI

Authors: Mingjie Li, Edward Kim, Yue Zhao, Ehsan Adeli, Kilian M. Pohl

Organizations: Stanford University

Abstract

Learning a robust Variational Autoencoder (VAE) is a fundamental step for many deep learning applications in medical image analysis, such as MRI synthesizes. Existing brain VAEs predominantly focus on single-modality data (i.e., T1-weighted MRI), overlooking the complementary diagnostic value of other modalities like T2-weighted MRIs. Here, we propose a modality-aware and anatomically grounded 3D vector-quantized VAE (VQ-VAE) for reconstructing multi-modal brain MRIs. Called NeuroQuant, it first learns a shared latent representation across modalities using factorized multi-axis attention, which can capture relationships between distant brain regions. It then employs a dual-stream 3D encoder that explicitly separates the encoding of modality-invariant anatomical structures from modality-dependent appearance. Next, the anatomical encoding is discretized using a shared codebook and combined with modality-specific appearance features via Feature-wise Linear Modulation (FiLM) during the decoding phase. This entire approach is trained using a joint 2D/3D strategy in order to account for the slice-based acquisition of 3D MRI data. Extensive experiments on two multi-modal brain MRI datasets demonstrate that NeuroQuant achieves superior reconstruction fidelity compared to existing VAEs, enabling a scalable foundation for downstream generative modeling and cross-modal brain image analysis.

Figures & tables

Figure 1: Conventional approaches handle different MRI modalities by duplicating separate encoders. In contrast, we employ a unified 3D encoder that disentangles shared anatomical representations $z_{anat}$ from modality-specific features $z_{mod}$ within the same latent space.

Figure 2: Architecture of NeuroQuant: It learns to disentangle anatomical and modality-specific representations from multi-modal brain MRI. The encoder uses factorized multi-axis attention to extract shared anatomical features and modality-adaptive cues. Anatomical features are discretized via a shared vector-quantized codebook, while modality features generate FiLM parameters $(\gamma,\beta)$ through a global modality embedding $s_{m}$ and MLP to guide the decoder. The Modality-FiLM decoder reconstructs 3D MR volumes with modality-adaptive contrast and anatomical coherence. Cross-modal reconstruction and modality prediction ensure disentanglement and consistency across modalities. $\bigoplus$ refers to the concatenate operation.

			T1				T2
	Models	$f$	PSNR	SSIM	Accuracy	Dice	PSNR	SSIM	Accuracy	Dice
NCANDA	VQGAN	8x	26.09	90.77	76.23	90.08	25.57	88.28	66.25	89.80
	SD-VAE	8x	27.48	92.72	77.14	92.51	26.16	90.87	67.34	90.37
	MediTok	16x	28.03	94.11	77.78	93.88	27.34	93.08	67.55	92.53
	MedVAE	8x	28.26	94.38	78.12	94.83	27.59	92.51	67.87	92.31
	NeuroQuant (ours)	16x	28.89	95.27	79.09	95.25	28.12	93.57	68.36	93.45

Table 1: Quantitative comparison with state-of-the-art methods.

f

denotes the spatial compression ratio of each VAE. All the metrics are higher is better. For each dataset and metric, the best result is highlighted in bold , and the second-best result is underlined .

Figure 3: Qualitative comparison of T1- and T2-weighted MRI reconstructions across three views. Unlike existing models that exhibit blurring or structural artifacts (especially in non-axial planes), NeuroQuant achieves superior voxel-level fidelity and anatomical consistency. Notably, our joint 2D/3D training strategy (NeuroQuant) significantly enhances the recovery of fine-grained cortical details and sharp tissue boundaries compared to the 3D-only variant (w/o j.t.).

	T1			T2
Settings	PSNR	SSIM	Accuracy	PSNR	SSIM	Accuracy
Base model	27.56	92.18	77.38	26.43	91.24	67.51
+ FMA	28.17	94.37	78.22	27.39	92.08	67.85
+ M-FiLM	28.52	94.84	78.39	27.83	92.49	68.01
+ $\mathcal{L}_{\mathrm{cross}}$	28.63	95.01	78.45	27.94	92.76	68.06
+ $\mathcal{L}_{\mathrm{adv}}$ (full)	28.89	95.27	79.09	28.12	93.57	68.36

Table 2: Ablation study on the NCANDA dataset evaluating the contribution of key components in NeuroQuant. We incrementally add factorized multi-axis attention (FMA), modality-FiLM decoder (M-FiLM),

\mathcal{L}{\mathrm{cross}}

, and

\mathcal{L}{\mathrm{adv}}

Figure 4: SynthSeg-based cerebellum cortex and caudate volume comparison between GT, our full model, and ablated variants.

Vector-Quantized Variational Autoencoder

Vector-quantized variational autoencoders (VQ-VAEs) learn discrete codebook representations for compression, generation, and structured data modeling.

See all 34 papers

Latent graph encoding of multimodal neuroimaging features with generative AI architectures
Jul 8, 2026Ishaan Batta, Meenu Ajith, Vince CalhounVariational Autoencoders Multimodal Embedding
Center for Translational Research in Neuroimaging and Data Science (TReNDS): Georgia State University, Georgia Institute of Technology, and Emory University, Atlanta, USA
Modeling Local, Global, and Cross-Modal Context in Multimodal 3D MRI
Jun 25, 2026Minh Duc Do, Tillmann Rheude, Noel Kronenberg +23D ViTs Brain Age Prediction
1Berlin Institute of Health at Charité - Universitätsmedizin Berlin · 2Health Data Science Unit, Heidelberg University Hospital and BioQuant · 3Intelligent Medicine Institute, Fudan University +1
BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation
Jun 17, 2026Max Van Puyvelde, Ibrahim Gulluk, Wim Van Criekinge +1Masked Autoencoders Latent Diffusion Models
Department of Biomedical Data Science, Stanford University School of Medicine · Department of Mathematical Modelling, Statistics & Bioinformatics, Ghent University · Department of Electrical Engineering, Stanford University

Jul 8, 2026J/K move · Enter open · S save

Latent graph encoding of multimodal neuroimaging features with generative AI architectures

Ishaan Batta, Meenu Ajith, Vince Calhoun

Center for Translational Research in Neuroimaging and Data Science (TReNDS): Georgia State University, Georgia Institute of Technology, and Emory University, Atlanta, USA

Variational Autoencoders Multimodal Embedding Multimodal Generation Multimodal Graph Learning Neuroimaging

Open paper AskPDF arXiv

While generative models enable encoding of complex neuroimaging data for feature generation and reconstruction, developing optimal architectural frameworks with appropriate encoding and latent space processes is crucial for studying structural and functional properties of the brain. We design a multimodal generative framework for structural and functional magnetic resonance imaging (MRI) features through systematic evaluation of encoding strategies, latent multimodal fusion, and generative model selection. Using structural gray matter volume (GMV) and static functional network connectivity (sFNC) features from a large neuroimaging dataset, we analyze generative frameworks involving variational autoencoders (VAEs), transformers, generative adversarial networks (GANs), and diffusion models. Architectures that employ modality-aware graph encoding of functional connectivity into a lower-dimensional latent space outperform vectorized encoders or direct data space approaches. The proposed multimodal graph VAE (gMMVAE) surpasses alternative generative variants across multiple metrics for generation fidelity, reconstruction quality, efficiency, and latent space discriminability, highlighting its potential for robust multimodal neuroimaging analysis.

Abstract

Vector-Quantized Variational Autoencoder

Explore similar work

Latent graph encoding of multimodal neuroimaging features with generative AI architectures

Modeling Local, Global, and Cross-Modal Context in Multimodal 3D MRI

BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation