LangVAE and LangSpace: Building and Probing for Language Model VAEs

TL;DR

This paper introduces LangVAE, a modular framework for constructing language-model VAEs (LM-VAEs) atop pre-trained LMs, paired with LangSpace for systematic latent-space probing and evaluation. A key contribution is the KV cache injection mechanism, which guides autoregressive decoders without altering their architecture, enabling substantial reductions in training and inference compute (average parameter reductions >95% for many decoder sizes) while maintaining compatibility with state-of-the-art LLMs. The authors demonstrate the framework with a case study across encoder/decoder combinations and SRL-annotated inputs, revealing complex interactions in generalisation and disentanglement across architectural families; LangSpace provides quantitative metrics and visualisations (e.g., disentanglement scores, traversal quality) to systematise such experiments. The work emphasizes accessibility and reproducibility through PyPI releases, HF Hub checkpoints, and a supporting video, highlighting the practical impact of modular, scalable LM-VAEs for controlled text generation and interpretability.

Abstract

We present LangVAE, a novel framework for modular construction of variational autoencoders (VAEs) on top of pre-trained large language models (LLMs). Such language model VAEs can encode the knowledge of their pre-trained components into more compact and semantically disentangled representations. The representations obtained in this way can be analysed with the LangVAE companion framework: LangSpace, which implements a collection of probing methods, such as vector traversal and interpolation, disentanglement measures, and cluster visualisations. LangVAE and LangSpace offer a flexible, efficient and scalable way of building and analysing textual representations, with simple integration for models available on the HuggingFace Hub. Additionally, we conducted a set of experiments with different encoder and decoder combinations, as well as annotated inputs, revealing a wide range of interactions across architectural families and sizes w.r.t. generalisation and disentanglement. Our findings demonstrate a promising framework for systematising the experimentation and understanding of textual representations.

Figures (5)

• Figure 1: LangVAE is a flexible framework designed to support arbitrary combinations of pretrained encoders and decoders for learning latent representations under either a categorical semantic prior or a Gaussian prior. LangSpace facilitates comprehensive analysis of the learned latent space through automated evaluation of key properties such as disentanglement and visualization (top) and enables controlled generation by leveraging these latent properties, such as latent traversal, interpolation, and arithmetic operations (bottom).
• Figure 2: Diagram of fundamental LLVAE architecture.
• Figure 3: Overview of the LangVAE framework.
• Figure 4: Illustration of the KV cache injection scheme. $W_m z$ projects hidden KV cache entries $X^{h}_{t}$ that are attended by the decoder when predicting the next token. The hidden cache entries are interleaved with the ones produced by the decoder.
• Figure 5: TSNE plots for the [bert-base-cased, Llama-3.2-3B] combination, without (a) and with (b) SRL annotated inputs. We can observe a better separation of the water and animal subjects on the annotated model.