Distributional Latent Variable Models with an Application in Active Cognitive Testing

TL;DR

This work tackles the inefficiency and rigidity of traditional cognitive testing by introducing a nonlinear Distributional Latent Variable Model (DLVM) that embeds heterogeneous test observations into a shared latent space and maps them to distributional test parameters. The method combines a neural decoder with variational inference to capture cross-test and cross-participant correlations at the item level, and incorporates an active-learning framework based on mutual information to select the most informative tests for a new participant. Empirical results across multiple cohorts show DLVM achieves accuracy comparable to conventional approaches while requiring significantly fewer test items, improves robustness to outliers, and provides quantified uncertainty via distributional predictions. The approach has practical implications for rapid, uncertainty-aware cognitive assessments and scalable deployment in educational and clinical settings.

Abstract

Cognitive modeling commonly relies on asking participants to complete a battery of varied tests in order to estimate attention, working memory, and other latent variables. In many cases, these tests result in highly variable observation models. A near-ubiquitous approach is to repeat many observations for each test independently, resulting in a distribution over the outcomes from each test given to each subject. Latent variable models (LVMs), if employed, are only added after data collection. In this paper, we explore the usage of LVMs to enable learning across many correlated variables simultaneously. We extend LVMs to the setting where observed data for each subject are a series of observations from many different distributions, rather than simple vectors to be reconstructed. By embedding test battery results for individuals in a latent space that is trained jointly across a population, we can leverage correlations both between disparate test data for a single participant and between multiple participants. We then propose an active learning framework that leverages this model to conduct more efficient cognitive test batteries. We validate our approach by demonstrating with real-time data acquisition that it performs comparably to conventional methods in making item-level predictions with fewer test items.

Figures (7)

• Figure 1: Two example tests in the battery. Left: Span---Recall the order aliens appeared (binary accuracy). Right: Numerical Stroop---Select the box with more animals (reaction time).
• Figure 2: DLVM model architecture. The model only utilizes the decoder component of the traditional autoencoder framework. Given a latent vector $x$, it outputs the distributional parameters $\theta$. The model is trained to maximize the likelihood of the observed item-level data $Y_i$ under the output distributional parameters $\theta_i$.
• Figure 3: Median Fit of the trained DLVM model to the conventionally acquired data (colored) compared with the fit to data (gray). The median fit was identified based on the log probability of the data given the predicted parameters by the DLVM model. Timing outputs in milliseconds (left column) were fit with log-normal distributions. Item response times are indicated by short vertical lines, means by diamonds, and standard deviations by horizontal lines. Psychometric outputs (middle column) were fit by logistic sigmoids as a function of items to be remembered. Proportion correct for test times is indicated by circles, thresholds by diamonds and spreads by horizontal lines. Accurate completion of test items (right column) is fit by binomial distributions reflecting a hypothetical test containing 40 items. The proportion of test items successfully completed is indicated by vertical dashed lines.
• Figure 4: Distributions fitting test battery outputs for conventional test battery data (, gray) and for machine learning test battery data (, colors) for a median-fit participant (ID: MATML106). The median is based on the root-mean-square error (RMSE) between the and DLVM parameters. Plotting conventions as in \ref{['fig:in-sample_fits']}.
• Figure 5: Similarity of primary test outputs (i.e., means and thresholds) between the conventional test battery + independent maximum likelihood estimation () and machine learning battery + distributional latent variable modeling (). Timing task mean responses (left column) are in units of log milliseconds, psychometric task thresholds (center column) are in units of items recalled, and accuracy tasks (right column) are in units of probability correct. Numerical values are intra-class correlation coefficients.