Bayesian Distance-to-Set Models: from Latent Variable to Latent Projection

1. Problem Statement

Standard Bayesian latent variable models (LVMs) assume data are generated near a structured, low-dimensional set $Z_\theta$. In canonical LVMs, each observation $y_i$ is associated with a latent coordinate $z_i \in Z_\theta$, and the data is modeled as a noisy deviation from this coordinate (e.g., $y_i = z_i + \epsilon_i$).

Key Challenges:

• Computational Inefficiency: Inference requires sampling from the joint posterior $\Pi(\theta, z | y)$. As the sample size $n$ grows, the dimension of the latent space ($\Theta \times \mathbb{R}^{nd}$) increases, leading to slow mixing in Markov Chain Monte Carlo (MCMC) algorithms.
• Integration Difficulty: The marginal likelihood $L(y|\theta) = \int L(y|\theta, z)\pi_L(z|\theta)dz$ often lacks a closed form, necessitating data augmentation which exacerbates the dimensionality issue.
• Approximation Trade-offs: Existing remedies (custom MCMC, Laplace approximations, variational inference) are often model-specific, rely on potentially spurious approximations, or fail to preserve the geometric structure of the latent space.

The authors propose a paradigm shift: instead of modeling the deviation from a specific latent coordinate, model the distance from the data point to the set itself.

2. Methodology: Distance-to-Set (DTS) Models

Core Concept

The DTS framework replaces the latent coordinate $z_i$ with a latent projection $\hat{z}_i$, defined as the point in the set $Z_\theta$ closest to the observation $y_i$.

• Likelihood Definition: The likelihood is defined directly on the distance between the observation and the set:
  $$L(y_i; \theta, \sigma) = \phi(\text{dist}(y_i, Z_\theta); \sigma)$$
  where $\text{dist}(y_i, Z_\theta) = \min_{z \in Z_\theta} \|y_i - z\|$ and $\phi$ is a distance kernel (e.g., Gaussian).
• Latent Projection: The projection $\hat{z}_i = \text{argmin}_{z \in Z_\theta} \|y_i - z\|$ is computed via optimization (e.g., soft-thresholding for $\ell_1$-balls) rather than sampled. It serves as an auxiliary variable for likelihood evaluation but is not a parameter to be sampled in the MCMC chain.

Generative View & Geometric Properties

The authors provide a generative interpretation where:

1. A projection $\hat{z}_i$ is drawn from the image of the observation space under projection.
2. A residual $\epsilon_i$ is drawn from a constrained space (the normal cone of $Z_\theta$ at $\hat{z}_i$) such that $\hat{z}_i + \epsilon_i$ projects back to $\hat{z}_i$.
3. Independence: For convex sets, the deviation (residual) is independent of the set-shift parameters (e.g., fixed effects $\mu$), a property that holds due to the translation equivariance of the projection operator.

Regularization via Normalizing Constant

A critical theoretical contribution is the analysis of the normalizing constant $m_{\theta, \sigma}$.

• Occam's Razor Effect: As the set $Z_\theta$ expands, the distance to observations decreases, which would naturally favor larger sets. However, the normalizing constant $m_{\theta, \sigma}$ (which depends on the volume/surface area of the set's neighborhood) increases with set size.
• Penalty: The term $m_{\theta, \sigma}^{-1}$ in the posterior acts as an implicit regularizer, penalizing overly large sets and preventing overfitting without requiring explicit priors on set size.

3. Key Contributions

1. Dimensionality Reduction: By integrating out the latent coordinates analytically (via projection), the DTS model reduces the inference space from $(\theta, z)$ to just $(\theta)$. This eliminates the high-dimensional augmented space, enabling efficient MCMC sampling even for large $n$.
2. Tractable Marginal Likelihood: The approach yields a marginal likelihood that is computationally tractable (or close-to-tractable) because the projection $\hat{z}_i$ can often be computed in closed form (e.g., via proximal operators) or via fast optimization.
3. Theoretical Guarantees:
   • Posterior Consistency: The authors establish sufficient conditions (A1–A6) for the posterior of the set parameter $\theta$ to concentrate around the true parameter $\theta_0$, even when the likelihood is not factorizable over data indices.
   • Independence: Proved that for convex sets, the residual distribution is independent of the location parameters.
4. Flexible Framework: The method decouples the geometry of the set (defining sparsity, structure, etc.) from the noise distribution, allowing for flexible modeling (e.g., inducing sparsity via $\ell_1$-balls while using Gaussian kernels).

4. Empirical Results

The authors validated the DTS model through simulations and two real-world applications:

A. Multi-Environment Study (STAR Dataset)

• Task: Assess the generalizability of a "small-class" effect across 79 schools.
• Comparison: DTS vs. Pooled, Standard Hierarchical, Horseshoe, and Spike-and-Slab models.
• Results:
  • Interpretability: The radius parameter $r$ in DTS directly quantifies the degree of heterogeneity (generalizability).
  • Sparsity: DTS successfully identified sparse school-specific deviations, performing comparably to Horseshoe and Spike-and-Slab in estimating the global effect $\beta$.
  • Efficiency: DTS was significantly faster (e.g., ~107s vs. ~1045s for Spike-and-Slab on real data) because it avoids sampling high-dimensional latent variables.

B. Bayesian Transfer Learning

• Task: Transfer knowledge from a large source dataset ($S$) to a small target dataset ($T$) while preventing "reverse transfer" (where noisy target data distorts source estimates).
• Mechanism: Modeled the target parameters as lying within a distance-$r$ ellipsoid of the source parameters.
• Results:
  • Robustness: Unlike naive joint hierarchical models, DTS prevented the noisy target data from contaminating the source parameter estimates ($\beta_S$), even under large deviations.
  • Adaptivity: The model automatically tuned the transfer strength based on the distance between source and target.
  • Application (CTR Prediction): In a click-through rate prediction task, DTS achieved the lowest Root Mean Squared Error (RMSE) and lowest variance compared to OLS, Transfer-LASSO, and naive full-transfer methods.

5. Significance and Conclusion

The paper introduces a fundamental shift in Bayesian modeling for structured data. By replacing latent coordinates with latent projections, the authors solve the "curse of dimensionality" inherent in data-augmented MCMC for latent variable models.

• Computational Impact: The method enables scalable Bayesian inference for complex structured sets (e.g., sparse vectors, manifolds) that were previously computationally prohibitive.
• Statistical Insight: It provides a rigorous geometric interpretation of residuals (constrained to normal cones) and demonstrates how the geometry of the set itself induces regularization (Occam's razor).
• Practical Utility: The framework is highly applicable to transfer learning, multi-environment studies, and any scenario requiring structured, low-dimensional representations of high-dimensional data, offering a superior balance of interpretability, predictive accuracy, and computational speed.

The authors provide open-source code (Python/NumPyro) to facilitate adoption, and the framework opens new avenues for research in optimization-based likelihoods and geometric Bayesian inference.