Prediction under Latent Subgroup Shifts with High-Dimensional Observations

TL;DR

This paper presents a scalable method using recognition-parametrised models to identify latent structures in high-dimensional image data, enabling accurate prediction under distribution shifts caused by unobserved confounders.

Contribution

It introduces a novel application of RPM to recover discrete latent variables from complex images, facilitating prediction adaptation in high-dimensional settings with latent shifts.

Findings

Successfully identifies causal latent structures in image data.

Adapts predictions accurately under latent distribution shifts.

Scales to high-dimensional observations where previous methods fail.

Abstract

We introduce a new approach to prediction in graphical models with latent-shift adaptation, i.e., where source and target environments differ in the distribution of an unobserved confounding latent variable. Previous work has shown that as long as "concept" and "proxy" variables with appropriate dependence are observed in the source environment, the latent-associated distributional changes can be identified, and target predictions adapted accurately. However, practical estimation methods do not scale well when the observations are complex and high-dimensional, even if the confounding latent is categorical. Here we build upon a recently proposed probabilistic unsupervised learning framework, the recognition-parametrised model (RPM), to recover low-dimensional, discrete latents from image observations. Applied to the problem of latent shifts, our novel form of RPM identifies causal latent structure in the source environment, and adapts properly to predict in the target. We demonstrate results in settings where predictor and proxy are high-dimensional images, a context to which previous methods fail to scale.