Identifiable Autoregressive VAE

• The paper introduces IAVAE, a deep generative model that ensures identifiability by modeling latent AR processes under nonlinear mixing conditions.
• It leverages auxiliary spatio-temporal variables to parameterize latent exponential family priors, achieving interpretable blind source separation and robust forecasting.
• Empirical evaluations show that IAVAE outperforms linear methods like FastICA in recovering latent sources and predicting high-dimensional data trends.

An identifiable autoregressive variational autoencoder (IAVAE) is a deep generative modeling architecture designed to recover independent latent components comprising nonstationary autoregressive processes from observed multivariate spatio-temporal data, even under nonlinear and nonstationary mixing conditions (Sipilä et al., 15 Sep 2025). The principal innovation of IAVAE lies in ensuring identifiability of the latent representations, enabling dimension reduction, interpretable decomposition (blind source separation), and improved forecasting and interpolation. This approach combines advances in identifiable variational autoencoders (iVAE) with autoregressive (AR) latent dynamics, leveraging auxiliary information to fulfill provable identifiability criteria when modeling complex dependencies in high-dimensional data.

1. Architectural Foundations

The IAVAE extends the iVAE framework by explicitly modeling the latent sources as nonstationary AR processes and allowing for nonlinear mixing into the observed space. The generative model consists of:

• Nonlinear Mixing Function: Observed data $\mathbf{x}$ is produced via a smooth, injective function $f$ acting on independent latent variables $\mathbf{z}$, that is, $\mathbf{x} = f(\mathbf{z}) + \epsilon$, where $\epsilon$ is noise. In blind source separation, $\epsilon$ is often omitted.
• Autoregressive Latent Processes: Each latent dimension $z_i$ follows a nonstationary AR(R) process:

$$z_i(s, t) = \mu_i(s, t) + \sum_{r=1}^R \gamma_{i,r}(s,t)\,[z_i(s, t - r) - \mu_i(s, t - r)] + \omega_i(s, t),$$

where $s$ indexes spatial location, $t$ time, $\mu_i(s, t)$ is a spatio-temporally varying trend, $\gamma_{i,r}(s, t)$ AR coefficients, and $\omega_i(s, t)$ innovations.

• Conditional Exponential Family Priors: The conditional distribution of latents is parameterized as:

$$p(\mathbf{z} | \mathbf{z}^-, \mathbf{u}) = \prod_{i=1}^P \frac{Q_i(z_i, z_i^-)}{Z_i(\mathbf{u})} \exp\left\{ \sum_{j=1}^k T_{i,j}(z_i, z_i^-)\,\lambda_{i,j}(\mathbf{u}) \right\},$$

where $\mathbf{z}^-$ denotes previous AR states, $\mathbf{u}$ is an auxiliary variable encoding spatio-temporal location and past information, $T_{i,j}$ are sufficient statistics, $\lambda_{i,j}(\mathbf{u})$ the natural parameters, and $Q_i(\cdot)$ the base measure.

• Neural Submodules: The model comprises
  • Encoder $q_{\theta_g}(\mathbf{z} | \mathbf{x}, \mathbf{u})$
  • Decoder (mixing network) $f$
  • Auxiliary network $w$ estimating conditional AR prior parameters from $\mathbf{u}$

Optimization proceeds by maximizing the evidence lower bound (ELBO):

$$\mathbb{E}_{q_{\theta_g}(\mathbf{z} | \mathbf{x}, \mathbf{u})}\left[ \log p_{\theta_h}(\mathbf{x} | \mathbf{z}) + \log p_{\theta_w}(\mathbf{z} | \mathbf{z}^-, \mathbf{u}) - \log q_{\theta_g}(\mathbf{z} | \mathbf{x}, \mathbf{u}) \right].$$

2. Theoretical Identifiability

Identifiability refers to the ability to recover the true latent components, up to allowable indeterminacies (e.g., permutation, scaling, and shift for Gaussian AR sources), from the observed data. The IAVAE achieves identifiable decomposition under the following conditions:

• The latent prior (exponential family) parameters $\lambda(\mathbf{u})$ vary sufficiently with the auxiliary variable $\mathbf{u}$. Specifically, for $P$ latent dimensions and $k$ sufficient statistics per dimension, the matrix $$L = [\lambda(\mathbf{u}_1) - \lambda(\mathbf{u}_0), \dots, \lambda(\mathbf{u}_{Pk}) - \lambda(\mathbf{u}_0)]$$ must be invertible—ensuring enough nonstationarity.
• The mixing function $f$ must be injective and smooth.
• The sufficient statistics $T_{i,j}$ must be differentiable and linearly independent.

Under these conditions, theorems in (Sipilä et al., 15 Sep 2025) establish that the learned parameters $(f, T, \lambda)$ are identifiable up to an affine transformation. For Gaussian AR latent processes, identifiability specializes to permutation, scaling, and offset. This ensures that the latent sources correspond to well-defined physical or statistical processes, crucial for scientific interpretability and reliable temporal prediction.

3. Spatio-Temporal Auxiliary Variables and Nonstationarity

Auxiliary variables $\mathbf{u}$ are engineered to encode both spatio-temporal coordinates and historical observations to induce nonstationarity in the latent prior parameters. Two schemes are demonstrated:

• Radial Basis Function Approach: The auxiliary variable is constructed via RBF kernels spanning spatial and temporal locations at multiple resolutions, which models smooth variations in AR coefficients and variance.
• Segmentation Approach: Spatio-temporal locations are partitioned into segments, each represented by an indicator vector, capturing abrupt or blockwise changes in dynamics.

This parameterization allows the model to learn nonstationary behavior in latent trends and AR coefficients, which is necessary for identifiability in complex real-world data exhibiting local variance or regime transitions.

4. Blind Source Separation (BSS)

Blind source separation refers to recovering the independent latent “sources” from observed mixtures. The IAVAE enables nonlinear and nonstationary BSS by:

• Modeling the mixing transformation $f$ as a generic MLP, dispensing with restrictive linearity.
• Enforcing independence (apart from time dependence) in the latent components via the AR prior and proper auxiliary conditioning.
• Solving for the unmixing via the encoder $q(\mathbf{z}|\mathbf{x}, \mathbf{u})$ subject to the identifiability conditions.

Empirical comparisons demonstrate that IAVAE-based BSS (especially iVAEar_r and iVAEar_s variants) outperforms linear methods such as symmetric FastICA and STBSS in recovering true sources, especially when either AR coefficients or latent variances are highly nonstationary in space or time.

5. Multivariate Spatio-Temporal Prediction

IAVAE supports forecasting and interpolation in high-dimensional spatio-temporal settings by:

• Projecting observations $\mathbf{x}$ into latent AR processes.
• Predicting future latent states with the learned AR dynamics.
• Mapping latent forecasts through the learned nonlinear decoder $f$ to obtain predictions of the original variables.

This approach enables flexible prediction of air pollution concentrations and meteorological variables across spatial locations and time horizons. Weighted mean squared error (wMSE) metrics demonstrate superiority over baseline models such as ARIMA, VARIMA, and spatio-temporal kriging, especially when both latent variance and AR coefficients vary nonstationarily.

6. Empirical Evaluation and Practical Implications

Simulation studies and real-world applications validate the architecture:

Method	Latent Recovery (MCC)	Prediction Performance (wMSE)
iVAEar_r	Highest	Best
iVAEar_s	High	Consistently better than linear BSS
FICA, STBSS	Lower	Inferior when nonstationarity is present

• IAVAE variants yield the highest mean correlation coefficients between recovered and true latent sources across a range of scenarios.
• Forecasting accuracy benefits from explicit modeling of nonstationary latent dynamics.
• Accurate estimation of the latent dimension and flexible auxiliary design are crucial for optimal operation.

A plausible implication is that introducing autoregressive structure in the latent space, conditioned on sufficiently variable auxiliary data, is broadly applicable for interpretable dimension reduction, forecasting, and source separation in spatio-temporal data domains.

7. Extensions and Generalization

The theoretical and methodological advances in IAVAE provide a template for generative sequence modeling under identifiability guarantees. The integration of AR priors with nonlinear mixing and nonstationary auxiliaries can be adapted to other autoregressive architectures, including language, video, longitudinal biomedical signals, and complex dynamical systems. These insights are transferable to architectures seeking interpretable latent dynamics and robust forecasting capabilities in nonstationary environments.