On Learning Prediction-Focused Mixtures

TL;DR

This paper introduces a prediction-focused approach for mixture models that automatically selects relevant features to improve prediction performance in low-capacity settings, enhancing interpretability and optimization.

Contribution

It proposes a novel prediction-focused modeling method for mixtures that identifies relevant input signals and outperforms traditional models in low-capacity scenarios.

Findings

Outperforms non-prediction-focused models in relevant signal detection

Eases optimization of mixture models with prediction focus

Provides theoretical characterization of when prediction-focused modeling is effective

Abstract

Probabilistic models help us encode latent structures that both model the data and are ideally also useful for specific downstream tasks. Among these, mixture models and their time-series counterparts, hidden Markov models, identify discrete components in the data. In this work, we focus on a constrained capacity setting, where we want to learn a model with relatively few components (e.g. for interpretability purposes). To maintain prediction performance, we introduce prediction-focused modeling for mixtures, which automatically selects the dimensions relevant to the prediction task. Our approach identifies relevant signal from the input, outperforms models that are not prediction-focused, and is easy to optimize; we also characterize when prediction-focused modeling can be expected to work.