Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning | Summary

Contents

Summary
1 Introduction
2 Related Work
3 Methodology
- 3.1 Learning Latent Abstractions
- 3.2 Reasoning with Discrete Latent Tokens
4 Experiments
5 Conclusion
Impact Statement
Appendix
Brief Thoughts

This article explains the key points of Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning.

2025-02-05 (arXiv), 42nd International Conference on Machine Learning, PMLR 267, 2025
Su, DiJia, Zhu, Hanlin, Xu, Yingchen, Jiao, Jiantao, Tian, Yuandong, Zheng, Qinqing.
Meta AI, UC Berkeley, UCL
Paper

Read this article in Korean

Summary

Token Assorted addresses the computational cost of long chain-of-thought (CoT) traces, the explicit intermediate text steps used to solve a problem. Compared with complete textual CoT training on MetaMathQA, a dataset of generated mathematical solutions, its mixed latent-and-text training raises average accuracy across seven mathematical benchmarks from 14.1 to 17.8, 25.2 to 28.1, and 33.4 to 37.9 for the tested models with 1, 3, and 8 billion parameters, respectively, while the paper reports an average 17% reduction in generated response length.
The method learns discrete latent tokens, vocabulary entries representing compressed text chunks, using a vector-quantized variational autoencoder (VQ-VAE), an encoder–decoder that maps text through a finite set of learned vectors. It replaces a randomly sized prefix of each training trace with these tokens and preserves the remaining text; at inference, the language model generates both token types directly.
Experiments cover navigation planning, synthetic logical reasoning, and mathematical reasoning with models up to 8 billion parameters. Replacement ablations, controlled comparisons of replacement choices, support retaining text and varying the compressed prefix, but the efficiency evaluation measures response token counts rather than runtime. Analyses of attention, the model’s weights for combining token representations, and decoded latent examples offer limited evidence about how the representations work.

1 Introduction

The Introduction motivates compression by observing that textual reasoning contains words needed for linguistic coherence alongside task-relevant information. Its central adaptation problem is teaching a pretrained large language model (LLM), a language model trained on extensive text, to use previously unseen latent vocabulary entries without a multistage schedule for language-model training.

Related Work distinguishes explicit CoT prompting, training on reasoning traces, and reasoning in latent space. Alternatives include self-consistency, selecting an answer by voting over generated responses; continuous latent reasoning, feeding hidden representations back into the model; and pause tokens, learned tokens that provide computation before an answer.

The proposed method combines learned discrete compression with a remaining textual trace. The discussed continuous methods instead continue computation through hidden representations.
The experiments compare against iCoT, a curriculum that progressively removes textual reasoning, and Pause Token, alongside complete-CoT and answer-only training.

3 Methodology

The method separates each example into prompt P, reasoning trace C, and solution S, written X = P ⊕ C ⊕ S, where ⊕ denotes concatenation. After learning a compressed trace Z, it replaces the first m text tokens with m/r latent tokens, producing X̃ = P ⊕ [z₁, …, zₘ/ᵣ, cₘ₊₁, …, cₜ꜀] ⊕ S; r is the ratio of text-trace length to latent-trace length.

Prefix replacement with two discrete latent tokens at chunk size L = 16 and compression rate r = 16.

3.1 Learning Latent Abstractions

The VQ-VAE is trained on the complete example but subsequently compresses only the reasoning trace. Its encoder maps chunks of L text tokens to L/r vectors, quantization replaces each vector with its nearest vector in a codebook, the learned finite vector set, and the prompt-conditioned decoder reconstructs text. A latent token identifies the selected codebook entry.

The training objective combines reconstruction, vector-quantization, and commitment terms: these train text recovery, move codebook entries toward encoder outputs, and encourage encoder outputs to remain near their selected entries.
A straight-through gradient, passing the decoder gradient across the nondifferentiable quantization operation, trains the encoder. Both encoder and decoder use causal Transformers, sequence models whose attention is restricted to the current and preceding positions.

Text encoding, nearest-codebook quantization, and prompt-conditioned reconstruction in the VQ-VAE.

3.2 Reasoning with Discrete Latent Tokens

Reasoning with Discrete Latent Tokens trains next-token prediction on partially compressed examples. Three choices define the procedure: retain a textual suffix, replace a contiguous prefix from left to right, and randomize the prefix length separately for each example.

The tokens <boLatent> and <eoLatent> delimit the latent segment, and the language-model vocabulary includes the learned latent identifiers.
The paper describes a 50M-parameter VQ-VAE used during data preparation. Language-model training and inference do not run this auxiliary model; inference directly generates latent tokens followed by text.

4 Experiments

Experiments examine synthetic planning and logical reasoning, followed by fine-tuning Llama language models on mathematical reasoning data. Sol-Only uses questions and final answers, CoT uses complete reasoning traces, iCoT progressively removes those traces, and Pause Token adds learned computation tokens.

4.1 Benchmarks

4.1.1 SYNTHETIC BENCHMARKS includes Keys-Finding Maze, navigation through colored locked doors while carrying at most one key; ProntoQA, 9000 logical problems with exactly one correct proof path; and ProsQA, 17,886 logical problems generated from directed acyclic graphs, graphs with directed edges and no directed cycles, with distracting reasoning paths. Maze training contains 100k examples with intermediate search traces generated by nondeterministic A*, a heuristic search algorithm.

4.1.2 MATHEMATICAL REASONING uses MetaMathQA or Dart-MATH, another dataset of generated mathematical solutions, to train on questions derived from Math and GSM8K, mathematical problem-solving benchmarks. MetaMathQA contains 395k examples, comprising 155k derived from Math and 240k from GSM8K; the authors regenerate responses using Llama-3.1-405B-Inst.

Math and GSM8K are treated as in-domain evaluations, meaning their source datasets underpin the training data. Dart-MATH emphasizes difficult questions through difficulty-aware rejection tuning, a procedure for selecting generated training responses.
Out-of-domain evaluations, drawn from other datasets, are College-Math, 2818 problems from 9 textbooks; DeepMind-Math, 1000 school-level problems; OlympiaBench-Math, 675 English text-only competition problems; TheoremQA, 800 theorem-based problems; and Fresh-Gaokao-Math-2023, 30 examination questions.

4.2 Main Results

The default compressor is trained for 100k steps with Adam, the parameter-update optimizer, learning rate 10⁻⁵, and batch size 32, using L = 16 and r = 16 to encode 16 text tokens per latent token. For randomized replacement, the paper prints M = {0, 72, 128, 160, 192, 224, 256}, samples a maximum from this set, and then samples m uniformly from the stated sequence [0, 16, 32, …, m_max].

The printed value 72 conflicts with the accompanying statement that every member of M is a multiple of 16, leaving this sampling detail ambiguous.
Although the main text states a codebook size of 1024, Appendix A.1 specifies task-dependent sizes: 64 for logical reasoning, 512 for the maze, and 1024 for mathematics.

4.2.1 SYNTHETIC BENCHMARKS fine-tunes pretrained GPT-2, a language model, for 16k steps on ProntoQA and ProsQA, using learning rate 10⁻⁴, 100 warmup steps that gradually increase the learning rate, batch size 128, and greedy decoding, which selects the highest-probability next token. The maze uses a T5 encoder–decoder model trained from scratch for 100k steps at learning rate 7.5 × 10⁻⁴ and batch size 1024; its 1-Feasible-10 metric measures whether any of 10 responses sampled with top-k decoding, restricted to the 10 highest-probability tokens, reaches the goal feasibly.

Table 4.1 reports maze success of 62.8% versus 43% for CoT, ProntoQA accuracy of 100% versus 98.8%, and ProsQA accuracy of 96.2% versus 77.5%. These correspond to gains of 19.8, 1.2, and 18.7 percentage points.
Mean generated lengths fall from 1312.0 to 374.6 tokens for the maze, 92.5 to 7.7 for ProntoQA, and 49.4 to 10.9 for ProsQA.

Synthetic planning and logical reasoning performance with generated token counts.

4.2.2 MATHEMATICAL REASONING fine-tunes Llama-3.2-1B, Llama-3.2-3B, and Llama-3.1-8B on MetaMathQA for 1 epoch, one pass through the training data, with batch size 32 and sequence packing of 4096, combining examples into sequences of that length. Validation selects learning rate 10⁻⁵ for 8B and 2.5 × 10⁻⁵ for 1B and 3B; evaluation uses greedy decoding.

Table 4.2 gives seven-benchmark average accuracies of 17.8, 28.1, and 37.9, compared with CoT averages of 14.1, 25.2, and 33.4. The corresponding gains are 3.7, 2.9, and 4.5 percentage points.
Examples include Math improving from 10.5 to 14.7 for 1B, GSM8K from 69.7 to 73.8 for 3B, and Fresh-Gaokao-Math-2023 from 16.7 to 30.0 for 8B. The method does not lead every comparison: the 8B Pause Token model reaches 39.6 on Math versus 37.2 for the latent model.
Table 4.3 reports mean response lengths of 518, 513.6, and 513.7 tokens versus CoT lengths of 655.2, 642.9, and 578.5. The paper summarizes these as an average 17% reduction; iCoT produces shorter responses but substantially lower average accuracy.

Mathematical reasoning accuracy after MetaMathQA fine-tuning at three model sizes.

Generated response lengths across mathematical benchmarks and model sizes.

4.3 Ablation & Understanding Studies

Replacement ablations compare AR-Replace, the proposed randomized prefix replacement, against All-Replace, replacing the entire trace; Curriculum-Replace, progressively replacing it over 8 epochs; and Poisson-Replace, independently replacing each 16-token chunk with probability 0.5. Table 4.4 reports the highest average accuracy for AR-Replace at every tested size: 17.8, 28.1, and 37.9, versus 15.1, 20.5, and 36.3 for Poisson-Replace.

All-Replace and Curriculum-Replace obtain substantially lower averages, supporting mixed traces under these conditions. The explanation that remaining text helps align latent tokens with pretrained linguistic representations is the authors’ interpretation.
Attention analysis examines two questions using the final attention layer, averaging across heads, parallel attention computations, and normalizing column-summed attention weights. Figure 4.1 shows greater attention to selected numbers and operation-related words, but these two examples do not establish a causal mechanism.

Accuracy comparison of complete, curriculum, random-chunk, and randomized-prefix replacement.

Normalized prompt-token attention for two mathematical questions under latent and CoT models.

4.4 Ablations on the Latent r parameters

The compression ablation tests r = 2, 16, and 32 on Llama-3.2-3B. Table 4.5 reports average accuracies of 27.08, 28.1, and 27.9 with mean response lengths of 595.6, 513.6, and 481.4 tokens, compared with CoT accuracy of 25.2 and length of 642.9.

Increasing compression reduces average response length, while r = 16 gives the highest average accuracy among these three settings.
The measured accuracy–length trade-off applies to this model and training setup. Individual datasets do not uniformly improve: DeepMind-Mathematics accuracy remains below the CoT baseline at every tested compression rate.

Accuracy and generated length for compression rates 2, 16, and 32 on Llama-3.2-3B.

Average mathematical accuracy and token count as functions of latent compression rate.

4.5 Additional Examples and Interpretability Result

Additional Examples and Interpretability Result points to four generated responses in Appendix D, decoded latent examples in Appendix E, and Dart-MATH experiments in Appendix F. The supplementary training-data comparison also produces shorter responses, although its average accuracy improvement over CoT is much smaller than in the MetaMathQA experiment.

5 Conclusion

The Conclusion attributes the observed accuracy and token-count improvements to combining compressed initial reasoning with detailed textual continuation and randomizing compression during training. The experiments support this combination across the tested planning, logic, and mathematics tasks, with exceptions on individual mathematical benchmarks.

Impact Statement

The Impact Statement identifies the opacity of latent reasoning as a potential safety concern and proposes decoding latent tokens into readable text as a mitigation. The paper supplies qualitative decoding examples but does not evaluate safety outcomes or whether decoded text faithfully explains the model’s computation.

Appendix

A Experiment Details, including A.1 VQ-VAE Model Details, specifies 2-layer encoder and decoder Transformers with 4 attention heads, embedding size 512, and block size 512, the configured token-block length. Codebook sizes are 64 for ProntoQA and ProsQA, 512 for Keys-Finding Maze, and 1024 for mathematics; maximum sequence lengths are 2048 for synthetic experiments and 256 for mathematics. A.2 Keys-Finding Maze covers A.2.1 ENVIRONMENT DETAILS, A.2.2 DATASET DETAILS, and A.2.3 MODEL DETAILS. Its illustrated environment has 3 × 3 rooms of 5 × 5 cells, and the dataset contains 100k training, 500 validation, and 300 test examples; structured prompts encode walls, agent, objective, keys, and doors, while responses contain search-node records and trajectories. The model uses 6 layers, 3 attention heads, embedding size 64, rotary position embeddings that encode token positions, and no dropout, the training procedure that randomly suppresses activations. A.3 ProntoQA and ProsQA identifies pretrained GPT-2 with 12 layers, hidden size 768, 12 attention heads, vocabulary size 50,257, and 117 million parameters. A.4 LLM experiments names Llama Cookbook, https://github.com/meta-llama/llama-cookbook, and repeats the fine-tuning settings; its reference to Llama-3.2-8B differs from the Llama-3.1-8B identifier used in the main experiments.
B Notations collects the prompt, trace, solution, compression, codebook, encoder, quantizer, and decoder symbols.
C Details of Attention Weights Analysis and C.1 Generated Responses provide the two attention-analysis examples: both approaches answer the percentage-difference problem correctly, while the latent approach answers the savings problem with 25 days and CoT gives 35.
D Other Text Generation Examples shows responses on exponent multiplication, a binomial coefficient, vertical asymptotes, and directed cycles; these are illustrative cases rather than aggregate evaluations.
E Interpreability Examples decodes latent prefixes for proportional conversion and greatest-common-divisor problems. The decoded fragments contain grammatical errors and malformed text, including “We find the greatest of CD of::63 common find,” while the subsequent response correctly computes the greatest common divisor of 315 and 108 as 9. These examples show inspectable fragments without establishing reliable semantic reconstruction.
F Additional Experiments fine-tunes Llama-3.1-8B on DART-MATH. Table F.1 reports average accuracy of 40.8 for the latent method versus 40.4 for CoT, with lower latent accuracy on GSM8K, DeepMind-Mathematics, and TheoremQA. Table F.2 reports mean response lengths of 452.7 versus 529.1 tokens and labels the reduction as 16%, although these displayed averages imply approximately 14.4%.

Brief Thoughts

The strongest evidence is the joint improvement in benchmark accuracy and generated length, reinforced by replacement ablations in which complete compression performs poorly. The DART-MATH result, only 0.4 percentage points above CoT on average, shows that the accuracy benefit varies with the training-data condition.

Response length does not directly establish latency or total training cost, including compressor training and data preparation. The results also lack reported uncertainty estimates, and the 30-question Fresh-Gaokao-Math-2023 benchmark makes its prominent gains sensitive to a small number of answers.

Reproduction requires resolving source inconsistencies: the replacement set includes 72 despite the stated 16-token alignment, and some prose and gain annotations disagree with table entries. For example, Table 4.2 gives an 8B average gain over CoT of 4.5 points, whereas the surrounding prose says 5.3; the numerical comparisons above use the table entries.