Bayesian Image-on-Image Regression via Deep Kernel Learning based Gaussian Processes

1. Problem Statement

The paper addresses the challenge of Image-on-Image Regression (IIR) in neuroimaging, specifically the task of predicting task-evoked functional MRI (fMRI) activity using resting-state fMRI data.

• Core Difficulty: Integrating heterogeneous imaging predictors with varying resolutions and spatial domains (e.g., voxel-wise fractional amplitude of low-frequency fluctuation [fALFF] maps vs. region-wise connectivity matrices) to predict high-resolution outcome images (task contrast maps).
• Limitations of Existing Methods:
  • Deep Learning (CNNs, GANs, RBDN): Often lack interpretability, require massive datasets, struggle with different input/output resolutions, and lack rigorous statistical inference (uncertainty quantification).
  • Traditional Statistical Models (Linear Regression, SBLF): Often fail to capture complex non-linear associations, ignore spatial correlations, or suffer from overfitting and computational inefficiency when scaling to large neuroimaging datasets.
• Motivating Application: Predicting task-based fMRI contrast maps (Language: story-math; Social: random-baseline) from resting-state data (fALFF and connectivity) in the Human Connectome Project (HCP).

2. Methodology: BIRD-GP

The authors propose BIRD-GP, a Bayesian hierarchical framework that integrates Deep Neural Networks (DNNs) with Gaussian Processes (GPs) via deep kernel learning. The method operates in a two-stage analysis:

Stage 1: Image Projection via Basis Expansion

Instead of modeling raw high-dimensional voxels directly, the method projects images into a lower-dimensional Euclidean space using learned basis functions.

• Model: Both predictor ($X_i$) and outcome ($Y_i$) images are modeled as realizations of Gaussian Processes.
  • $X_i(v) = x_i^\top \psi_x(v) + \epsilon_x$
  • $Y_i(u) = y_i^\top \psi_y(u) + \epsilon_y$
• Deep Kernel Learning: Unlike traditional PCA or fixed kernels (e.g., Squared-Exponential), the orthonormal basis functions ($\psi_x, \psi_y$) are learned adaptively using a Deep Neural Network (DNN). The DNN learns the eigenfunctions of the covariance kernel, allowing the model to capture complex spatial structures and smoothness constraints inherent in brain images.
• Orthogonalization: The DNN output is orthogonalized (via Gram-Schmidt or SVD) to form the basis functions used for projection.

Stage 2: Nonlinear Regression via DNN

The projected coefficients ($x_i$ and $y_i$) are linked via a feed-forward DNN.

• Conditional Distribution: The projected outcome coefficients are modeled as:
  • $(y_i | x_i) \sim \mathcal{N}(\mathcal{N}(x_i; \theta_{xy}), \Lambda_y)$
  • Where $\mathcal{N}(\cdot)$ is a multi-layer DNN capturing the non-linear mapping between resting-state features and task-evoked activity.
• Bayesian Framework: Priors are placed on DNN weights, biases, and variance parameters. This allows for the generation of posterior predictive distributions, providing valid uncertainty measures (credible intervals) for predictions.

Posterior Computation

To handle the computational complexity of high-dimensional Bayesian inference, the authors develop a hybrid posterior computation algorithm:

1. Stage 1 (Gibbs Sampler): Efficiently samples the posterior of the projected coefficients and variance parameters. This step is parallelizable across subjects.
2. Stage 2 (Stein Variational Gradient Descent - SVGD): Uses SVGD to approximate the posterior of the DNN parameters ($\theta_{xy}$). SVGD is a gradient-based variational inference method that updates a set of particles to match the target posterior distribution, offering a scalable alternative to MCMC for deep learning models.

Feature Importance

The paper introduces a method to quantify the importance of specific dimensions in the projected predictor space. It defines an importance function based on the expected magnitude of the derivative of the log-likelihood with respect to the input, approximated via Monte Carlo sampling. This allows for the identification of which basis images (and underlying brain regions) drive the predictions.

3. Key Contributions

1. Novel Framework (BIRD-GP): A unified Bayesian framework combining DNNs and GPs for IIR, capable of handling heterogeneous input modalities (voxel-wise and region-wise) and different resolutions.
2. Deep Kernel Learning: A data-adaptive approach to learning covariance kernels (basis functions) via DNNs, overcoming the rigidity of fixed kernels and the noise sensitivity of PCA.
3. Scalable Inference: The development of a hybrid Gibbs-SVGD algorithm that makes Bayesian inference feasible for large-scale neuroimaging data (hundreds of subjects, thousands of voxels).
4. Interpretability: Unlike "black-box" deep learning, BIRD-GP provides interpretable basis images and quantifies prediction uncertainty, crucial for scientific discovery in neuroscience.
5. Feature Importance Metric: A novel metric to identify which projected dimensions (and thus which brain networks) are most critical for prediction.

4. Results

Synthetic Data Experiments

• Datasets: MNIST, Fashion MNIST, and HCP-based synthetic fMRI data.
• Performance: BIRD-GP outperformed state-of-the-art competitors (CNNs, DNNs, RBDN, Linear Regression, and SBLF) in terms of prediction accuracy (MSE and classification accuracy).
  • On MNIST, BIRD-GP achieved near-perfect classification accuracy (1.00) with significantly fewer parameters (126K) compared to RBDN (445K) and standard CNNs.
  • On HCP synthetic data, BIRD-GP demonstrated diagonal-dominant prediction correlation matrices, meaning it predicted individual subjects' outcomes more accurately than it predicted other subjects' outcomes, indicating strong subject-specific generalization.
• Robustness: BIRD-GP remained robust even when the data was generated from a mis-specified model (voxel-wise regression), performing comparably to the true generative model.

Real Data Analysis (HCP 1200)

• Task: Predicting Language (story-math) and Social Recognition (random-baseline) task contrast maps from resting-state fALFF and connectivity matrices.
• Predictor Comparison:
  • Connectivity vs. fALFF: Connectivity matrices alone provided superior predictive power compared to fALFF maps alone for both tasks.
  • Combination: Combining both modalities yielded the best performance, though connectivity was the dominant driver.
  • Task Difficulty: The Language task was more predictable than the Social Recognition task.
• Model Comparison: BIRD-GP significantly outperformed Linear Regression and Voxel-wise Regression (VR) in prediction accuracy and subject-specific correlation.
• Neurobiological Insights:
  • Visual Network: Highly active in predicting both tasks.
  • Dorsal Attention Network: Crucial for both language and social cognition.
  • Subcortical Network: Dominant in specific basis images, supporting language and social functions.
  • Default Mode Network: Specifically engaged in language tasks but deactivated in social recognition.

5. Significance

• Scientific Utility: BIRD-GP bridges the gap between the predictive power of deep learning and the interpretability required for scientific inference in neuroscience. It allows researchers to not only predict brain activity but also understand which resting-state features drive those predictions.
• Clinical & Research Impact: By successfully predicting task-evoked activity from resting-state data, the method offers a cost-effective alternative to expensive task-based fMRI studies, potentially enabling the use of resting-state scans as biomarkers for cognitive traits or clinical conditions.
• Methodological Advancement: The integration of SVGD with Bayesian hierarchical models for image regression sets a new standard for scalable, uncertainty-aware deep learning in high-dimensional spatial data.