The Free Transformer

TL;DR

The Free Transformer addresses limitations of purely autoregressive decoding by introducing latent conditioning through a conditional VAE for a decoder-only Transformer. It injects a latent Z into the model's middle layer and trains a lightweight encoder with a Binary Mapper to produce discrete latent codes, enabling unsupervised learning of latent structure while managing information flow with a free-bits KL objective. Empirically, Free Transformer improves performance on reasoning benchmarks for 1.5B and 8B models and shows further gains when trained on 1T tokens, all with a modest encoder overhead of roughly 3%. The approach enhances inductive bias in generation and opens avenues for more structured latent representations in large-scale autoregressive models.

Abstract

We propose an extension of the decoder Transformer that conditions its generative process on random latent variables which are learned without supervision thanks to a variational procedure. Experimental evaluations show that allowing such a conditioning translates into substantial improvements on downstream tasks.

Figures (6)

• Figure 1: A standard decoder Transformer (a) can be extended to utilize a random state $Z$ in inference (b, left), in which case it has to be trained as a conditional VAE with an encoder (b, right). The Free Transformer (c) reduces the overhead of the encoder by having the random state $Z$ injected in its middle layer (c, left), and using for encoder during training the combination of its first half combined with one non-causal layer specific to the encoder (c, right). See Figure \ref{['fig:free-transformer']} for a detailed depiction of that architecture.
• Figure 2: The Free Transformer. We omit the normalization layers and residual connections from the model and the batch size from the tensor shapes for clarity. The operators in orange are specific to the encoder and are evaluated only for training or KV cache pre-filling, those with a dashed contour have no trainable parameters. The Binary Mapper is described in § \ref{['sec:binary-mapper']}. During generation, the encoder is not evaluated and $Z$ is sampled uniformly among the one-hot vectors of dimension $2^H$.
• Figure 3: The synthetic sequences of § \ref{['sec:synthetic-dataset']} are of fixed length, with a "target" made of a random letter repeated 8 times at a random position, an i.i.d. noise of exclamation marks, and a prompt indicating the target's letter.
• Figure 5: Experiments with 1.5B models trained on 47B tokens. Comparison on standard benchmarks of the baseline and our models. The optimization hyperparameters were tuned for the baseline and kept unchanged, but the Free Transformers require $3.6\%$ more compute and parameters for the encoder.
• Figure 6: Experiments with 8B models trained on 200B tokens. Comparison on standard benchmarks of the baseline and our models. The training procedure was tuned for the baseline and kept unchanged, but the Free Transformers require $3.1\%$ more compute and parameters for the encoder.