Enhancing Latent Computation in Transformers with Latent Tokens

TL;DR

The paper introduces latent tokens, lightweight dummy tokens inserted into decoder-only Transformers to provide latent computation and steer autoregressive generation while freezing the backbone. It defines a principled inference and training setup, including a specialized positional encoding that preserves existing token geometry and a loss that targets verbal tokens, enabling parameter-efficient fine-tuning of the latent vectors. Through synthetic tasks and benchmark evaluations, latent tokens demonstrate improved generalization, particularly in out-of-distribution scenarios, and three proposed mechanisms—self-prompting, information retrieval, and instruction adherence—are empirically explored. The findings suggest latent tokens offer a flexible, scalable augmentation that enhances intermediate computation with minimal infrastructure overhead, though further work is needed to map their internal mechanisms and optimize adaptive insertion strategies.

Abstract

Augmenting large language models (LLMs) with auxiliary tokens has emerged as a promising strategy for enhancing model performance. In this work, we introduce a lightweight method termed latent tokens; these are dummy tokens that may be non-interpretable in natural language but steer the autoregressive decoding process of a Transformer-based LLM via the attention mechanism. The proposed latent tokens can be seamlessly integrated with a pre-trained Transformer, trained in a parameter-efficient manner, and applied flexibly at inference time, while adding minimal complexity overhead to the existing infrastructure of standard Transformers. We propose several hypotheses about the underlying mechanisms of latent tokens and design synthetic tasks accordingly to verify them. Numerical results confirm that the proposed method noticeably outperforms the baselines, particularly in the out-of-distribution generalization scenarios, highlighting its potential in improving the adaptability of LLMs.

Figures (11)

• Figure 1: A decoder-only Transformer model with latent tokens. Left: The inference process. Right: The training process. In this visualization, we assume that one latent token is inserted periodically, once every two verbal tokens, into the original sequence, though it can be much more general, as discussed in Section \ref{['sec:design_choices']}.
• Figure 2: Comparisons of position IDs (a) with the standard Transformer, (b) with a naive way of inserting latent tokens, and (c) with the proposed method.
• Figure 3: Several ways of inserting latent tokens into the sequence. The white squares denote verbal tokens, while those with different colors refer to different latent tokens. In particular, Start_1 and End_1 add one latent token at the start or end of the query, respectively; 2_1 inserts one latent token every two verbal tokens; FS means function specialization. In 2_1 (w/ FS), we show four groups of latent tokens with different functions for completeness; however, it is possible to select several groups to accommodate the need of the task.
• Figure 4: Numerical results for the Generation task, averaged over three random runs. The number $23\%$ represents a relative improvement over the best baseline.
• Figure 5: Attention maps for the proposed Comma_2 method. Two latent tokens are denoted by < LATENT1> and < LATENT2>, respectively. For comparisons, the attention maps of other methods can be found in Appendix \ref{['sec:attention_maps']}.

Theorems & Definitions (1)

• Remark 1