PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Thought Models with Variational Bayes Inference-Time Computation

Deqian Kong, Minglu Zhao, Dehong Xu, Bo Pang, Shu Wang, Edouardo Honig, Zhangzhang Si, Chuan Li, Jianwen Xie, Sirui Xie, Ying Nian Wu

arXiv:2502.01567v2cs.CLcs.LGstat.ML

TL;DR

Data availability is a bottleneck to scaling language models. Latent Thought Models use explicit latent thought vectors and dual-rate training, and show stronger pretraining and zero-shot results alongside few-shot arithmetic in-context learning.

  • Problem

    As language models scale, data availability has emerged as a critical bottleneck to continued scaling.

    Source: §1 · p. 1

  • Method

    LTMs use explicit latent thought vectors with prior models in latent space, trained through a dual-rate optimization algorithm.

    Source: §6 · p. 9

  • Results

    LTMs show superior pretraining perplexity and zero-shot performance; at equivalent trFLOPs/tok, LTM-Medium achieves perplexity 10.95 versus GPT-2-Large at 11.5, using 6.7% of its parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Few-shot arithmetic in-context learning appears at scales significantly smaller than GPTs and scales further with model size.

    Source: §1 · p. 2

    LTMs currently lack a latent-space reward or verifier model to guide latent-vector optimization for reasoning.

    Source: §4 · p. 9

Abstract

from arXiv · show

We propose a novel class of language models, Latent Thought Models (LTMs), which incorporate explicit latent thought vectors that follow an explicit prior model in latent space. These latent thought vectors guide the autoregressive generation of ground tokens through a Transformer decoder. Training employs a dual-rate optimization process within the classical variational Bayes framework: fast learning of local variational parameters for the posterior distribution of latent vectors (inference-time computation), and slow learning of global decoder parameters. Empirical studies reveal that LTMs possess additional scaling dimensions beyond traditional Large Language Models (LLMs), such as the number of iterations in inference-time computation and number of latent thought vectors. Higher sample efficiency can be achieved by increasing training compute per token, with further gains possible by trading model size for more inference steps. Designed based on these scaling properties, LTMs demonstrate superior sample and parameter efficiency compared to autoregressive models and discrete diffusion models. They significantly outperform these counterparts in validation perplexity and zero-shot language modeling tasks. Additionally, LTMs exhibit emergent few-shot in-context reasoning capabilities that scale with model size, and achieve competitive performance in conditional and unconditional text generation.

1. Introduction

Latent Thought Models (LTMs) combine explicit latent thought vectors with a dual-rate variational training scheme. Their scaling dimensions and reported results suggest improved sample and parameter efficiency alongside gains across several language tasks.

1. Introduction

  • The approach targets a data-availability bottleneck in conventional scaling, where performance gains have relied on increasing model and data size.Source: §1 · p. 1
  • LTMs use explicit latent thought vectors with a prior to guide an autoregressive Transformer decoder, trained by fast posterior inference for local latent parameters and slower learning of global decoder parameters.Source: §1 · pp. 1–2
  • LTMs scale with data, model size, inference steps, and latent size; increasing inference steps improves sample and compute efficiency, while larger latent sizes offer additional headroom.At fixed training FLOPs per token, these dimensions can be optimized together.Source: §1 · p. 2
  • At equivalent trFLOPs/tok, LTM-Medium achieves perplexity comparable to GPT-2-Large (10.95 vs. 11.5) with only 6.7% of its parameters.The introduction also reports LTM-Small at 11.85 perplexity with 26% less trFLOPs/tok and 5.0% of GPT-2-Large’s parameters.Source: §1 · p. 2
  • LTM-Medium and LTM-Large achieve 52.2% and 91.7% reductions in perplexity, respectively, compared with state-of-the-art results at GPT-2 scale.Source: §1 · p. 2
  • LTMs show emergent few-shot in-context arithmetic learning at scales significantly smaller than GPTs, with the capability increasing with model size.The introduction also reports that increasing the number of latent thought vectors appears helpful.Source: §1 · p. 2
  • LTM-Large outperforms autoregressive and diffusion counterparts on conditional sentence completion by MAUVE, while unconditional-generation metrics are comparable to GPT-2-Large and generation is faster.Source: §1 · p. 2

2. Method

LTMs generate tokens with prior-governed latent thoughts and a Transformer decoder; training pairs sequence-specific variational inference with slower shared-decoder updates.

2.1. Latent Thought Models (LTMs)

  • LTMs generate tokens autoregressively conditioned on latent thoughts; layered vectors attend to corresponding decoder layers, and a 256-token context makes z carry information across segments.The model assumes an isotropic Gaussian prior over the latent thoughts.Source: §2.1 · p. 3

2.2. Learning and Posterior Inference

  • Classical variational Bayes optimizes sequence-specific q(z|x) parameters with fast gradient steps, alternating with slow shared-decoder updates; finite inference steps may leave inference unconverged.The variational posterior is Gaussian with a diagonal covariance matrix for computational efficiency.Source: §2.2 · pp. 3–4

2.3. Conditional and Unconditional Generation

  • Conditional generation infers z from prompt x before sampling a completion; unconditional generation samples prior latent thoughts, while updating z with each new token is left for future work.The alternative updating scheme is described as more computationally intensive.Source: §2.3 · pp. 4–5

2.4. Inference-Time Computation

  • Inference-time computation adds cost in training and testing; Tfast is the most influential trFLOPs/tok factor, while Nz has minimal impact and layer cost depends on attention relative to embeddings.For the adopted generation sampling scheme, the compute estimate is a worst-case estimate across tasks.Source: §2.4 · p. 5

3. Empirical Study

The empirical study examines LTMs on OpenWebText and language-modeling, reasoning, and generation tasks against autoregressive and diffusion baselines. It reports scaling benefits from inference steps and latent size, strong benchmark results, and layered semantic reconstruction in latent thought vectors.

3.1. Experimental Setup

  • Experiments use OpenWebText for pretraining, compare against autoregressive and discrete-diffusion baselines, and evaluate zero-shot perplexity across several dataset validation splits.OpenWebText contains around 8B tokens; the last 100K documents are reserved for validation.Source: §3.1 · p. 5

3.2. Scaling Behaviors

  • More inference steps improve sample efficiency and can become compute-efficient beyond training-compute thresholds; larger latent sizes further improve both efficiencies.Across model sizes and inference-step configurations, increasing latent vectors improves performance; inference-step gains appear from 16 to 128 steps.Source: §3.2 · p. 6, Figure 4 · p. 6

3.3. Comparison with Existing Language Models

  • 10.95 vs. 11.5 perplexity: LTM-Medium matches GPT-2-Large at equivalent training compute using 6.7% of its parameters; LTM-Large reaches 3.05 validation perplexity with 76M parameters.LTM-Large was trained on 3B tokens and selected for its tradeoff between inference speed and sample efficiency.Source: §3.3 · pp. 6–7
  • LTMs consistently outperform existing baselines across zero-shot language-modeling benchmarks at equivalent training-compute checkpoints.Table 1 reports zero-shot unconditional perplexity across datasets.Source: Table 1 · p. 7, §3.3 · p. 7
  • At 38M parameters, LTM-Small with five-shot demonstrations surpasses GPT-2-Medium finetuning and test-time-search baselines; larger models improve zero- and few-shot performance.Figure 5 compares arithmetic reasoning on GSM8K; the reported metric is pass@5.Source: Figure 5 · p. 7, §3.3 · p. 7
  • Table 2 reports better conditional completion than autoregressive and diffusion models; Table 3 reports comparable unconditional Gen PPL and entropy with 5× faster sampling.For unconditional results, the authors include only baselines with entropy above 5.6.Source: §3.3 · pp. 7–8, Table 2 · p. 8, Table 3 · p. 8

3.4. Ablation Studies

  • Variational Bayes learning mitigates posterior collapse, while the VAE baseline suffers collapse and Langevin sampling yields lower-quality generations than variational Bayes.The VAE baseline uses an identical decoder and a 12-layer encoder with full attention.Source: §3.4 · p. 8, Table 4 · p. 8

3.5. Probing Results on Latent Thought Vectors

  • Progressive reconstruction reveals distributed processing through layers 1–8, followed by accuracy above 95% at synthesis layers 9–10 in the 12-layer, 96-vector model.The case study describes a progression from disconnected terms to coherent phrases and then complete semantic integration.Source: §3.5 · p. 8

4. Limitations: Prior and Reward

The paper identifies limitations in its latent prior and in its support for inference-time reasoning. It points to more structured priors and latent-space reward or verifier models as future directions.

4. Limitations: Prior and Reward

  • The current model uses a simple Gaussian prior and separate thought-vector sets that cross-attend to different decoder layers.The authors identify more structured, learnable priors with more interpretable latents as a potential improvement.Source: §4 · p. 8
  • A latent reasoning model that generates a chain of thought vectors is offered as a possible way to model language of thought.This would make posterior inference a process of parsing, formalization, compression, and understanding.Source: §4 · p. 9
  • The model lacks a latent-space reward or verifier model to guide optimization of latent vectors during inference-time reasoning.The authors connect such models to prior applications in offline reinforcement learning and online molecule-design optimization.Source: §4 · p. 9

5. Related Work

Related work spans autoregressive and diffusion language modeling, variational latent-variable models, explicit latent-space methods, cognitive theories, and test-time computation. The paper situates LTMs among these lines of work.

Autoregressive and Diffusion Language Modeling

  • Autoregressive models such as GPT-3 and PaLM have broad language-task success, while discrete diffusion offers an alternative, including masked diffusion related to any-order autoregression.The cited masked-diffusion approach iteratively transitions tokens into a masked state in the forward process.Source: §5 · p. 9

Variational Bayes Language Modeling

Large Language Models with Explicit Latent Space

Declarative-Procedural Model in Cognitive Science

  • The declarative-procedural account distinguishes declarative memory for facts and associations from procedural memory for learned skills and language rules.In the paper's model, latent vectors parallel declarative or episodic memory, while the decoder generator corresponds to procedural memory.Source: §5 · p. 9

Language of Thought (LOT) Hypothesis

  • The language-of-thought hypothesis proposes a mental language with syntax and semantics underlying natural-language learning and use.The passage also points to computational work on LOT-like structures in cognitive modeling and program induction.Source: §5 · p. 9

Complementary Learning and Test-Time Computation

  • Complementary learning systems link rapid learning of specific experiences to the hippocampus and slower learning of general knowledge to the neocortex.Test-time computation work includes dynamic adjustment of recurrent-network computation, test-time parameter adaptation, and non-parametric context caches.Source: §5 · p. 9

6. Conclusion

The conclusion presents LTMs as language models with explicit latent thought vectors and dual-rate optimization, and highlights scaling behavior as a focus of evaluation. It positions structured priors and latent-space reward or verifier models as future directions for reasoning and planning.

6. Conclusion

  • LTMs combine explicit latent thought vectors with explicit prior models in latent space and a dual-rate optimization algorithm.The paper empirically investigates scaling with inference steps and latent dimensionality, drawing inspiration from cognitive-science theories.Source: §6 · p. 9
  • The paper frames its work as groundwork for more structured, interpretable priors and latent-space reward-verifier models for reasoning and planning.Its stated inspirations include declarative-procedural memory systems, the language-of-thought hypothesis, and complementary learning systems.Source: §6 · p. 9

Impact Statement

The impact statement describes LTMs as a model class for explicit latent thought vectors and inference-time computation, with potential for explicit internal representations and reasoning or planning. The appendix specifies implementation choices and the OpenWebText validation split.

Impact Statement

  • The model class has potential to learn more explicit internal representations and enable more explicit reasoning and planning based on them.The statement describes potential, not demonstrated outcomes.Source: Impact Statement · p. 10

A.1. Model Details

  • Training and posterior inference use FlashAttention and the Liger kernel; attention layers use RMS layer normalization and SwiGLU.The passage cites Dao et al. (2022), Hsu et al. (2024), and Zhang & Sennrich (2019).Source: §A.1 · p. 13
  • All LTMs use 512 hidden dimensions, 8 attention heads, a 1024-token maximum sequence length, and latent vectors matching hidden-vector dimensionality.The autoregressive generator uses a sliding window of 256, with rotary position embeddings for ground tokens and latent vectors in each layer.Source: §A.1 · p. 13
  • OpenWebText documents are concatenated and wrapped to length 1024 with EOS separators, and the last 100K documents form validation.The paper uses the GPT-2 tokenizer, adds one EOS token, and does not pad or truncate sequences.Source: §A.1 · p. 13

A.2. Training Details

Training uses separate slow and fast learning rates, with AdamW updating global parameters and Adam updating latent thought vectors.

A.2. Training Details

  • The slow learning rate starts at 4 × 10−4 and follows cosine decay to 4 × 10−5.The schedule applies to global-parameter training.Source: §A.2 · p. 13
  • Training warms up linearly for the first 1000 iterations and clips the gradient norm to 1.Source: §A.2 · p. 13
  • The fast learning rate increases linearly from 0.3 to 0.34.Source: §A.2 · p. 13
  • AdamW updates global parameters with β1 = 0.9 and β2 = 0.95, while Adam updates latent thought vectors.The authors state that Adam is used for latent thought vectors without introducing additional inductive bias in optimization.Source: §A.2 · p. 13

A.3. Experiment Details

The experiments evaluate zero-shot perplexity, arithmetic reasoning, and conditional and unconditional generation using distinct benchmark and decoding protocols.

A.3. Experiment Details

  • Zero-shot perplexity is measured on validation splits of PTB, Wikitext, LM1B, Lambada, AG News, PubMed, and Arxiv after training on OpenWebText.The evaluation uses the detokenizers from Sahoo et al. (2024) and inserts an [EOS] token between dataset sequences.Source: §A.3 · p. 13
  • GSM8K arithmetic reasoning is evaluated on its 1K test set using pass@5, with a problem counted solved if any of five generated candidates matches the final solution.Each candidate is limited to 50 new tokens.Source: §A.3 · p. 13
  • GPT-2 baselines use beam search with beam size 5, while LTMs infer z five times per prompt and sample multinomially after each inference.Few-shot prompts concatenate examples before response generation.Source: §A.3 · p. 13
  • Conditional generation is evaluated with MAUVE on five continuations per OpenWebText validation sample, conditioning on 50 tokens and generating 50 more.The evaluation uses 1,000 samples; baseline results in Table 2 come from Lou et al. (2024).Source: §A.3 · p. 13
  • Unconditional generation is evaluated by GPT-2-XL perplexity on 64 greedy-decoded sequences of 1024 tokens prompted with [BOS].The authors use GPT-2-XL rather than GPT-2-Large for evaluation, citing fair calculation of GPT-2-Large Gen PPL.Source: §A.3 · p. 13

A.4. Samples for Unconditional Generation

The section presents unconditional sample outputs for LTM models at three scales.

A.4. Samples for Unconditional Generation

A.5. Samples for Conditional Generation

The section includes a conditional-generation example and probes how latent thought vectors encode and reconstruct information across model layers.

A.5. Samples for Conditional Generation

  • The LTM-Large conditional-generation example marks generated tokens in blue.Source: Figure 9 · p. 17

B. Probing the Latent Thought Vectors

  • The probing study progressively includes latent thought-vector layers from bottom to top across 200 OpenWebText validation samples and also presents a layer-by-layer case study.Source: B. Probing the Latent Thought Vectors · p. 17

B.1. Progressive Layer Inclusion

  • ∼55% accuracy through layers 1–5 in 6-layer LTM-Medium rises to complete reconstruction at layer 6; LTM-Large reaches ∼65% through layers 1–8 and >95% after synthesis at layers 9–10.The figure compares models with 24 latent vectors (4 per layer) and 96 latent vectors (8 per layer), respectively, and describes deeper models as distributing information more gradually.Source: Figure 10 · p. 17

B.2. Case Study

As more layers are included, text reconstruction improves from scattered matches to complete accuracy, with semantic information distributed hierarchically across the latent space.

B.2. Case Study

  • Reconstruction accuracy rises from 22% at layers 0–3 to 100% with all 12 layers.At layers 0–3, only scattered words match the original.Source: Figure 11 · p. 18, §B.2 · p. 18
  • At layers 0–6, accuracy reaches 30%, and structural elements and phrases about the ocean and landscape emerge.Source: Figure 11 · p. 18
  • At layers 0–9, accuracy reaches 65%, with coherent phrases and key descriptive elements.The progression is described as showing semantic information distributed hierarchically across the model’s latent space.Source: Figure 11 · p. 18
Loading 2502.01567v2…