Bayesian Active Learning for Discrete Latent Variable Models

TL;DR

The paper addresses the data-efficiency challenge of fitting discrete latent variable models by formulating Bayesian active learning around mutual-information utility $I(y; \theta\mid \mathbf{x}, \mathcal{D})$ and developing two practical inference-based pipelines (Gibbs-sampling and variational inference) to estimate this utility. It applies these infomax strategies to two model classes—mixtures of linear regressions (MLR) and input-output GLM-HMMs—showing substantial reductions in required data and faster parameter recovery compared with random or amortized approaches. Across extensive simulations and a real-world CA housing dataset, the Gibbs-based infomax method consistently achieves faster decreases in posterior entropy and lower RMSE, while the variational variant offers speedups at some cost to accuracy. The framework enables efficient adaptive experiments in neuroscience and other domains, and highlights the importance of tailoring active learning to latent-state structure, with potential extensions to higher-dimensional outputs and continuous inputs.

Abstract

Active learning seeks to reduce the amount of data required to fit the parameters of a model, thus forming an important class of techniques in modern machine learning. However, past work on active learning has largely overlooked latent variable models, which play a vital role in neuroscience, psychology, and a variety of other engineering and scientific disciplines. Here we address this gap by proposing a novel framework for maximum-mutual-information input selection for discrete latent variable regression models. We first apply our method to a class of models known as "mixtures of linear regressions" (MLR). While it is well known that active learning confers no advantage for linear-Gaussian regression models, we use Fisher information to show analytically that active learning can nevertheless achieve large gains for mixtures of such models, and we validate this improvement using both simulations and real-world data. We then consider a powerful class of temporally structured latent variable models given by a Hidden Markov Model (HMM) with generalized linear model (GLM) observations, which has recently been used to identify discrete states from animal decision-making data. We show that our method substantially reduces the amount of data needed to fit GLM-HMM, and outperforms a variety of approximate methods based on variational and amortized inference. Infomax learning for latent variable models thus offers a powerful for characterizing temporally structured latent states, with a wide variety of applications in neuroscience and beyond.

Figures (7)

• Figure 1: Discrete latent variable regression models and infomax learning. (A) Schematic of a discrete latent variable model for regression settings. The response $y$ of the model given a stimulus $x$ and a latent $z$ is produced by generalized linear models. Here the discrete latent variable $z$ determines which of the three generalized linear models at the bottom determines the input-output mapping on any trial. (B) Infomax learning for discrete latent variable models. On trial $t$, present an input $\mathbf{x}_{t}$ to the system of interest (e.g., a mouse performing a decision-making task) and record its response $y_{t}$. We assume this response depends on the stimulus (input) as well as an internal or latent state $z_{t}$, as specified by the model $P(y_t \mid \mathbf{x}_t, z_t, \theta)$. Second, update the posterior distribution over model parameters $\theta$ given the data collected so far in the experiment, $\mathcal{D}_{t} = \{\mathbf{x}_{1:t}$, $y_{1:t}\}$ using either MCMC sampling or variational inference. Third, select the input for the next trial that maximizes information gain, or the mutual information between the next response $y_{t+1}$ and the model parameters $\theta$.
• Figure 2: Infomax learning for mixture of linear regressions (MLR) models. (A) Model schematic. At time step $t$, the system is in state $z_{t} = k$ with probability $\pi_{k}$. The system generates output $y_{t}$ using state-dependent weights $\mathbf{w}_{k}$ and independent additive Gaussian noise (Eq. \ref{['eq:MLR']}). (B) Example 2-state model with two-dimensional weights $\mathbf{w}_{1} = (1,\,0)$ and $\mathbf{w}_{2} = (-1,\, 0)$. We consider possible inputs on the unit circle, which are the information-maximizing inputs for linear Gaussian models under an $L_2$ norm constraint. (C) Fisher information as a function of the angle between $\mathbf{w}_{1}$ and the input presented to the system, for different noise variances $\sigma^{2}$. (D) Comparison between infomax active learning (using Gibbs sampling and VI), DAD and random sampling for the 2D MLR model shown above with mixing probabilities $\pi = [0.6, 0.4]$ and noise variance $\sigma^2=0.1$. Error bars reflect 95% confidence interval (standard error) of the mean across 20 experiments. (E) Performance comparison for the same 2-state model but with 10-dimensional weight vectors and inputs. The possible inputs to the system were uniform samples from the 10-D unit hyper-sphere.
• Figure 3: Histogram showing inputs selected by our active learning method (over the course of 200 trials) on mixture of linear regressions (MLRs), when inputs lie on a 2D circle (see Fig. \ref{['fig:results_mlr']}). We find a drop in probability at 90$^\circ$, this is also predicted by the Fisher Information analysis discussed in text for infomax (using Gibbs sampling). However, we do not see such a trend while using DAD. Inputs selected by DAD were distributed over the unit circle with modes at multiple of 30$^\circ$. (DAD requires a continuous range of inputs, hence it select inputs from all over the unit circle as opposed to a discrete list.)
• Figure 4: Application of infomax learning to CA Housing Dataset kelley_pace_sparse_1997. (A) Best fitting mixing weights for 3 state MLR to 5000 samples of the dataset. (B) Best fitting state weights for 3 state MLR to 5000 samples of the CA housing dataset. Orange, green and blue represent states 1, 2 and 3 respectively. Black represents the linear regression fit. (C) BIC as number of MLR states is varied from 1 (standard linear regression) to 5. We select the 3 state model as BIC begins to level off beyond 3 states. (D) Posterior entropy between the 3 state MLR parameters obtained using 5000 samples (parameters shown in (A) and (B)) and recovered parameters as a function of the number of samples for random sampling (blue) and infomax with gibbs sampling (red). Error bars reflect 95% confidence interval of the mean across 10 experiments. (E) The same as in (D) but for the RMSE (root mean squared error). (F) Visualization of standard deviation of 500 inputs selected by both infomax (red) and random sampling (blue). Each dot corresponds to a different experiment. Examining (B), it is clear that the 3 states differ most according to the weights placed on the 'AveOccup', 'Latitude' and 'Longitude' covariates. All 10 infomax experiments select inputs with greater variance for the latitude and longitude covariates than are selected by the random sampling experiments.
• Figure 5: Infomax for GLM-HMMs. (A) Data generation process for the GLM-HMM. At time step $t$, a system generates output $y_{t}$ based on its input, $\mathbf{x}_{t}$, as well as its latent state at that time step, $z_{t}$. The system then either remains in the same state, or transitions into a new state at trial $t+1$, with the transition probabilities given by matrix $A$. (B) Example settings for the transition matrix and state GLMs for a 3 state GLM-HMM. These are the settings we use to generate output data for the analyses shown in panels C and D. (C) Left: posterior entropy over the course of 1000 trials for random sampling (blue), infomax with a single GLM (grey), infomax for the full GLM-HMM using variational inference (VI) and Gibbs sampling (magenta and red respectively). Middle: root mean squared error for the recovered transition matrix for each of the three input-selection schemes (random/infomax with GLM/infomax with GLM-HMM (Gibbs)/infomax with GLM-HMM (VI)). Right: root mean squared error for the weight vectors of the GLM-HMM for each of the input-selection schemes. (D) Selected inputs for random sampling (blue), active learning when there is model mismatch and the model used for infomax is a single GLM (gray), active learning with infomax (using Gibbs sampling) and the full GLM-HMM (red). Selected inputs over the course of 1000 trials are plotted, and are shown on top of the generative GLM curves.