A Hybrid NUTS-Gibbs Sampler with State Space Marginalization for Estimation of Dynamic Structural Equation Models with Binomial Outcomes

1. Problem Statement

Dynamic Structural Equation Modeling (DSEM) is a standard framework for analyzing intensive longitudinal data (ILD), where participants are measured frequently over time. However, estimating DSEMs with discrete outcomes (specifically binomial or Bernoulli data) presents significant computational and statistical challenges:

• Scalability Issues: Traditional Metropolis-within-Gibbs samplers (e.g., as implemented in Mplus) require sampling $O(N \cdot T)$ parameters (where $N$ is the number of participants and $T$ is the number of timepoints) in every iteration. This scales poorly for large datasets (e.g., $N=269, T=104$), making estimation computationally prohibitive.
• Link Function Limitations: Existing Gibbs samplers for DSEMs are largely restricted to the probit link and the Bernoulli distribution. They struggle with the logit link and binomial data (counts of successes out of $n$ trials), which are common in fields like educational technology and cognitive testing.
• Sampling Inefficiency: Gibbs sampling suffers from high autocorrelation in high-dimensional, correlated parameter spaces. It also requires conditionally conjugate priors, limiting the flexibility of Bayesian workflows (e.g., using realistic non-conjugate priors or complex model extensions).
• Latent State Burden: Standard approaches sample latent states at every timepoint for every participant. While Kalman filters can analytically marginalize these states for Gaussian data, this exact marginalization is not directly applicable to non-Gaussian (discrete) responses without data augmentation.

2. Methodology

The authors propose a Hybrid NUTS-Gibbs Sampler that combines the strengths of Hamiltonian Monte Carlo (HMC) and Gibbs sampling, specifically tailored for DSEMs with discrete responses.

A. Model Formulation

The model extends DSEM to handle discrete responses ($y_{it}$) via a linear predictor ($y^*_{it}$) connected through a link function ($H$):

• Probit Link: Uses a latent response formulation with truncated normal distributions (Albert & Chib, 1993).
• Logit Link: Uses Pólya-Gamma (PG) data augmentation (Polson et al., 2013). This transforms the binomial likelihood into a conditionally Gaussian form by introducing auxiliary variables $\omega_{it} \sim PG(b, c)$. This allows the use of a logit link while maintaining a linear Gaussian state-space structure for the latent variables.

B. State Space Marginalization

The core innovation is the reformulation of the within-level model as a Linear Gaussian State Space Model (LG-SSM) using data augmentation.

• Marginalization: Instead of sampling the $O(N \cdot T)$ within-level latent states ($\eta_{1,it}$), the algorithm uses a Kalman Filter to analytically integrate them out.
• Result: The NUTS step targets a marginal posterior distribution over only $O(N + T)$ parameters (population, between-participant, and timepoint-specific parameters), rather than the full joint posterior including latent states.

C. The Hybrid Algorithm

The sampler alternates between two steps:

1. Gibbs Step (Latent Responses):
   • Samples the augmented latent responses ($\tilde{y}_{it}$) and auxiliary PG variables ($\omega_{it}$) conditional on current parameter values.
   • Uses Forward Filtering Backward Sampling (FFBS) to sample the latent state trajectories.
   • Parallelization: This step is highly parallelizable across participants ($N$) and timepoints ($T$).
2. NUTS Step (Model Parameters):
   • Executes one step of the No-U-Turn Sampler (NUTS) to update population-level ($\Theta$), between-participant ($\eta_2$), and timepoint-specific ($\eta_3$) parameters.
   • The target density is the marginal posterior computed exactly via the Kalman filter (integrating out $\eta_1$).
   • Efficiency: NUTS only simulates Hamiltonian dynamics for $O(N + T)$ parameters, avoiding the "curse of dimensionality" associated with sampling latent states.

3. Key Contributions

1. First Hybrid Sampler for Binomial DSEMs: Introduces a method that supports both probit and logit links for binomial data, utilizing Pólya-Gamma augmentation for the latter.
2. Computational Scalability: Reduces the dimensionality of the NUTS step from $O(N \cdot T)$ to $O(N + T)$. This makes estimation feasible for datasets with large $N$ and $T$ that were previously intractable.
3. Flexibility: By using NUTS for the structural parameters, the method supports non-conjugate priors and complex model extensions (e.g., asymmetric thresholds, non-linear effects) that are difficult to implement in pure Gibbs samplers.
4. Implementation: The algorithm is implemented in NumPyro (based on JAX), leveraging Just-In-Time (JIT) compilation and GPU acceleration for the Kalman filter operations.

4. Results

The authors conducted extensive simulation experiments and a real-world application:

• Simulation 1 (Participant-Invariant Dynamics):
  • The hybrid sampler significantly outperformed pure NUTS and pure Gibbs samplers.
  • Efficiency Gains: For $N=20, T=500$, the hybrid sampler was 2.4x more efficient than NUTS (probit) and 2.6x more efficient (logit).
  • Scalability: It successfully handled $T=2500$ timepoints, a scenario where pure NUTS and Gibbs samplers failed due to memory/time constraints.
• Simulation 2 (Participant-Varying Dynamics):
  • When dynamics vary by participant, pure Gibbs is impossible (lack of conjugacy).
  • The hybrid sampler outperformed pure NUTS when $N > T$ (efficiency gain up to 16.6x for $N=200$). Pure NUTS was slightly better when $N < T$.
• Simulation 3 (VAR(1) Model):
  • In a high-dimensional 9-indicator VAR(1) model, the hybrid sampler was 5.1x more efficient in bulk-ESS and 4.7x in tail-ESS compared to pure NUTS.
  • Time to reach a Bulk-ESS of 1,000 was 33 minutes (Hybrid) vs. 2.8 hours (NUTS).
• Application (Panic Attacks):
  • Applied to an Ecological Momentary Assessment (EMA) dataset ($N=43, T \approx 92$) predicting panic attacks from physiological data.
  • Successfully handled missing data via the Kalman filter prediction step.
  • Found no significant cross-lagged effects of heart rate/step count on panic attacks, demonstrating the model's practical utility.

5. Limitations and Future Directions

• Ordinal Data: The hybrid sampler is inefficient for ordinal responses. The constraint on threshold parameters creates a narrow acceptance region for NUTS, forcing tiny step sizes. Pure NUTS (without latent response formulation) performs better here.
• Negative Binomial Data: While theoretically possible via PG augmentation, the dispersion parameter ($r$) must be an integer for the standard PG sampling algorithm. Since NUTS cannot handle discrete parameters, estimating the dispersion parameter efficiently remains a challenge.
• GPU Acceleration: While the Kalman filter is GPU-friendly, the Pólya-Gamma rejection sampler may cause branch divergence on GPUs due to variable rejection rates across parallel threads.

6. Significance

This paper represents a major advancement in the computational feasibility of Dynamic Structural Equation Modeling for real-world, large-scale longitudinal data. By combining the exact marginalization capabilities of the Kalman filter with the efficient exploration of NUTS, the authors have solved the scalability bottleneck that previously forced researchers to use sub-optimal models (e.g., sum scores or treating discrete data as continuous). The method enables rigorous Bayesian inference for complex, high-dimensional DSEMs with binary and binomial outcomes, opening new avenues for research in psychology, education, and health sciences.