Latent Diffusion Energy-Based Model for Interpretable Text Modeling

TL;DR

This work addresses interpretability in text generation by combining a symbol–vector energy-based prior with diffusion-based latent-space recovery, forming a Latent Diffusion Energy-Based Model (LDEBM). It presents a variational framework that integrates a diffusion process in the latent space, a symbol-conditional prior, and a geometric clustering regularization with information bottleneck to produce well-structured, interpretable latent representations. Across synthetic and real data tasks, LDEBM demonstrates superior generation quality, robust sampling, and enhanced controllable generation and attribute discovery, including semi-supervised classification capabilities. The approach is train-from-scratch and applicable to text with or without labels, offering a principled path toward interpretable, controllable, and scalable text modeling.

Abstract

Latent space Energy-Based Models (EBMs), also known as energy-based priors, have drawn growing interests in generative modeling. Fueled by its flexibility in the formulation and strong modeling power of the latent space, recent works built upon it have made interesting attempts aiming at the interpretability of text modeling. However, latent space EBMs also inherit some flaws from EBMs in data space; the degenerate MCMC sampling quality in practice can lead to poor generation quality and instability in training, especially on data with complex latent structures. Inspired by the recent efforts that leverage diffusion recovery likelihood learning as a cure for the sampling issue, we introduce a novel symbiosis between the diffusion models and latent space EBMs in a variational learning framework, coined as the latent diffusion energy-based model. We develop a geometric clustering-based regularization jointly with the information bottleneck to further improve the quality of the learned latent space. Experiments on several challenging tasks demonstrate the superior performance of our model on interpretable text modeling over strong counterparts.

Figures (5)

• Figure 1: Graphical illustration of the latent diffusion process. We construct the forward and reverse diffusion processes in the latent space. The symbolic one-hot vector is coupled with the initial latent vector $\mathbf{z}_0$. The latent and diffused latent variables are highlighted by the red and blue plates, respectively. The cyan arrows indicate that $\mathbf{z}_0$ is connected with only $\mathbf{z}_1$. We learn a sequence of ebm to model the reverse diffusion process $p_\alpha(\mathbf{z}_t|\mathbf{z}_{t+1})$.
• Figure 2: Evaluation on 2D synthetic data: a mixture of 16 Gaussians (upper panel) and a 10-arm pinwheel-shaped distribution (lower panel). In each panel, the top, middle, and bottom row display densities learned by svebm-ib, our model w/o geometric clustering, and our full model, respectively. In each row, from left to right, it displays the data distribution and the kde of: $\mathbf{x}$ generated by amortized posterior $\mathbf{z}$ samples, $\mathbf{x}$ by mcmc sampled prior $\mathbf{z}$ samples, posterior $\mathbf{z}$ samples, and prior $\mathbf{z}$ samples.
• Figure 3: Visualization of color-coded data points. We visualize data points and the corresponding inferred latent variables of two 2D synthetic datasets (gaussian and pinwheel). Data points with different labels are assigned with different colors.
• Figure A1: Visualization of $p_\alpha(\mathbf{y}|\mathbf{z}_t)$ over $t$. $p_\alpha(\mathbf{y}|\mathbf{z}_t)$ is constantly around the probability of $0.5$ over $t$.
• Figure A2: Full evolution of svebm-IB and our models. In each sub-figure, we provide the typical states of the model trained on the corresponding dataset, sequentially from the top row to the bottom row.