Bayesian Kernel Machine Regression via Random Fourier Features for Estimating Joint Health Effects of Multiple Exposures

Technical Summary: Bayesian Kernel Machine Regression via Random Fourier Features for Estimating Joint Health Effects of Multiple Exposures

Problem Statement
Environmental epidemiology has traditionally focused on single-exposure analyses, yet humans are simultaneously exposed to multiple correlated risk factors. Bayesian Kernel Machine Regression (BKMR) is a widely used method for estimating the joint, nonlinear effects of multiple exposures by modeling the exposure-response surface as a realization of a Gaussian process. However, BKMR faces significant computational bottlenecks when applied to large datasets. The method requires repeated inversion of an $n \times n$ kernel matrix during Markov Chain Monte Carlo (MCMC) iterations, where $n$ is the sample size. This makes BKMR computationally infeasible for large administrative health databases (e.g., $n > 5,000$). While existing approximations like the predictive process (using knots) exist, they suffer from challenges such as oversmoothing and a lack of clear guidance for knot selection in multivariate settings.

Methodology: Fast BKMR
The authors propose "Fast BKMR," a method that replaces the Gaussian process random effects in BKMR with supervised Random Fourier Features (RFF). This approach leverages Bochner's theorem, which states that any positive definite, stationary kernel can be expressed as the Fourier transform of its spectral density.

Key methodological components include:

• Basis Expansion: The random effect function $h(x_i)$ is approximated by a finite set of basis functions:
  $$h(x_i) \approx \sum_{j=1}^{J} a_j \cos(\omega_j^T x_i) + b_j \sin(\omega_j^T x_i)$$
  where $\omega_j$ are frequencies sampled from a spectral density, and $a_j, b_j$ are random-effect coefficients.
• Supervised Learning: Unlike standard data-independent RFF, this method treats the frequencies $\omega_j$ as random variables to be learned from the data. This allows the basis functions to be adaptive to the specific exposure-response surface.
• Model Reframing: By using RFF, the kernel machine regression is reframed as a linear mixed-effect model. This eliminates the need to invert the $n \times n$ kernel matrix, significantly reducing computational complexity.
• Inference: Bayesian inference is conducted using MCMC with Hamiltonian Monte Carlo (HMC) algorithms. The frequencies $\omega_j$ and the coefficients $(\gamma, a_j, b_j)$ are updated in blocks. The number of basis functions ($J$) is selected using the Watanabe-Akaike Information Criterion (WAIC).

Key Contributions and Results
The paper evaluates Fast BKMR through simulation studies and a real-world application, comparing it against standard BKMR and BKMR with a predictive process (PP).

1. Computational Efficiency: Fast BKMR demonstrates substantial reductions in computation time. In simulations with $N=10,000$ and 10 exposures, standard BKMR with PP took approximately 6.6 days to fit a model, whereas Fast BKMR required only about 1 day (with $J=1,000$). For smaller basis sets (e.g., $J=5$), the reduction in time was even more dramatic (up to 99% reduction). The computation time for Fast BKMR scales almost linearly with sample size and the number of basis functions, whereas BKMR with PP scales poorly with both sample size and the number of exposures.
2. Estimation Accuracy:
   • Strong Correlation: When the exposure-response surface exhibits strong correlations (common in air pollution data), Fast BKMR outperforms BKMR with PP, even with a small number of basis functions.
   • Kernel Misspecification: In scenarios where the true kernel differs from the assumed Gaussian kernel, Fast BKMR maintains robust performance, often yielding lower Root Mean Squared Error (RMSE) than BKMR with PP, particularly with higher numbers of exposures.
   • Parametric Scenarios: When the true exposure-response function is parametric (non-Gaussian), Fast BKMR achieves performance comparable to standard BKMR as the number of basis functions increases, while maintaining superior speed.
3. Application to Birthweight Data: The method was applied to a dataset of 273,711 birth records from Georgia to analyze the joint effects of CO, NO$_2$, and PM$_{2.5}$ on birthweight. Fast BKMR successfully captured complex, nonlinear joint effects and provided narrower posterior intervals compared to multi-pollutant linear models with interaction terms. The analysis confirmed that increasing levels of NO$_2$ and PM$_{2.5}$ are associated with reduced birthweight, while the association for CO was positive but not statistically significant.

Significance and Claims
The authors claim that Fast BKMR addresses the critical scalability issue of BKMR, enabling the analysis of large administrative health databases that were previously intractable. The method offers a computationally efficient alternative that does not sacrifice the flexibility of modeling nonlinear joint effects and interactions.

The paper highlights several specific advantages:

• Scalability: It facilitates the analysis of datasets with tens of thousands of observations and multiple exposures.
• Robustness: It performs well even when the exposure-response surface has strong dependencies or when the kernel is misspecified.
• Numerical Stability: By avoiding the inversion of the kernel matrix, the method circumvents numerical singularity issues that arise in Gaussian process regression when exposures are highly correlated (e.g., individuals in the same census tract sharing identical exposure data).
• Flexibility: The supervised nature of the RFF allows the basis functions to adapt to the data, offering an advantage over deterministic expansions like the Karhunen–Loève expansion.

The authors note limitations, including that the method is still more time-consuming than non-Bayesian alternatives (like WQS or q-gcomp) and currently applies only to continuous outcomes. They suggest future work could explore variable selection via spike-and-slab priors and extensions to generalized linear models for binary or count outcomes.