Language Model Inversion through End-to-End Differentiation

Despite emerging research on Language Models (LM), few approaches analyse the invertibility of LMs. That is, given a LM and a desirable target output sequence of tokens, determining what input prompts would yield the target output remains an open problem. We formulate this problem as a classical gradient-based optimisation. First, we propose a simple algorithm to achieve end-to-end differentiability of a given (frozen) LM and then find optimised prompts via gradient descent. Our central insight is to view LMs as functions operating on sequences of distributions over tokens (rather than the traditional view as functions on sequences of tokens). Our experiments and ablations demonstrate that our DLM-powered inversion can reliably and efficiently optimise prompts of lengths $10$ and $80$ for targets of length $20$, for several white-box LMs (out-of-the-box).

💡 Research Summary

This paper addresses the Language Model Inversion (LMI) problem: given a frozen, pre-trained language model (LM) and a desired target output sequence of tokens, find an input prompt that would cause the LM to generate that target output. The authors frame this as a gradient-based optimization problem and propose a novel solution centered on creating an end-to-end Differentiable Language Model (DLM).

The core innovation is a shift in perspective: instead of viewing an LM as a function mapping sequences of tokens to other sequences of tokens, they reinterpret it as a function mapping sequences of distributions over tokens to other sequences of distributions over tokens. This abstract shift is operationalized by modifying two non-differentiable components in the standard LM autoregressive pipeline.

First, the non-differentiable hard embedding lookup is replaced with a soft embedding module. Here, the input is not a discrete token but a probability distribution over the entire vocabulary. The soft embedding computes a weighted sum of all embedding vectors, with weights given by this distribution, allowing gradients to flow back to the input distribution parameters.

Second, the non-differentiable next-token sampling (or argmax) step is replaced using the Gumbel-Softmax (GS) reparameterization trick. The GS provides a continuous, differentiable approximation to sampling from a categorical distribution. A critical enhancement is making the temperature parameter (τ) of the GS learnable per token position, via a parameter φ, following τ_eff = ε + τ_0 * (1 + tanh(φ)). This allows the optimization process to automatically navigate the bias-variance trade-off inherent in the GS estimator, improving sample efficiency.

The resulting DLMI algorithm initializes a matrix of logits (Z) representing the prompt and a set of temperature parameters (Φ). It then uses gradient descent (e.g., Adam) to optimize both Z and Φ jointly, minimizing a distance metric (e.g., cross-entropy) between the target output sequence and the sequence generated by the DLM when conditioned on the soft prompt derived from Z. Teacher forcing is used during the DLM’s forward pass for the target sequence.

The paper critically reviews prior LMI work, highlighting limitations in internal validity (assuming a single “ground-truth” prompt for a given output) and external validity (evaluating only on fluent text). Their evaluation protocol explicitly considers a spectrum of target output fluency, defined by the subject LM’s perplexity. Experiments demonstrate that their DLM-powered inversion can reliably and efficiently optimize prompts of lengths 10 and 80 for target sequences of length 20 across several white-box LMs, significantly outperforming a strong gradient-based baseline (GBDA) and working for both fluent and non-fluent targets.

In summary, the key contributions are: 1) a theoretically sound and practically efficient method to make frozen LMs end-to-end differentiable via the DLM framework, 2) a simple gradient-based algorithm that solves the LMI problem at reasonable scales, and 3) empirical demonstration of effective prompt inversion across a range of target fluencies, moving beyond the limitations of prior evaluation setups.