Source-linked AI summary
Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning
DiJia Su, Hanlin Zhu, Yingchen Xu, Jiantao Jiao, Yuandong Tian, Qinqing Zheng
TL;DR
Textual reasoning traces improve LLM reasoning but are long and computationally costly. This paper replaces initial trace segments with VQ-VAE-derived discrete latent tokens while retaining later text and randomly mixing both forms during training. Across planning, logical, and mathematical benchmarks, the approach consistently outperforms text-only reasoning-trace baselines.
Problem
Textual chain-of-thought traces are computationally costly because lengthy sequences contain tokens devoted mainly to linguistic coherence rather than core reasoning information.
Method
The method uses VQ-VAE-derived discrete latent tokens to replace initial reasoning steps and randomly varies replacement lengths during single-stage LLM training.
Results
Across planning, logical, and mathematical benchmarks, the models consistently outperform baselines trained with text-only reasoning traces.
Takeaways & Limitations
Mixing latent abstractions with textual details improves token efficiency while preserving a representation of essential reasoning information.
Takeaways & Limitations
Reasoning with opaque latent tokens may raise safety concerns, although the VQ-VAE decoder can convert them into human-readable text.
Abstract
from arXiv · showhide
Large Language Models (LLMs) excel at reasoning and planning when trained on chainof-thought (CoT) data, where the step-by-step thought process is explicitly outlined by text tokens. However, this results in lengthy inputs where many words support textual coherence rather than core reasoning information, and processing these inputs consumes substantial computation resources. In this work, we propose a hybrid representation of the reasoning process, where we partially abstract away the initial reasoning steps using latent discrete tokens generated by VQ-VAE, significantly reducing the length of reasoning traces. We explore the use of latent trace abstractions in two scenarios: 1) training the model from scratch for the Keys-Finding Maze problem, 2) fine-tuning LLMs on this hybrid data with an extended vocabulary including unseen latent tokens, for both logical and mathematical reasoning problems. To facilitate effective learning, we introduce a simple training procedure that randomly mixes latent and text tokens, which enables fast adaptation to new latent tokens. Our approach consistently outperforms the baselines methods in various benchmarks.
1 Introduction
Explicit chain-of-thought improves LLM reasoning but creates long, costly sequences dominated partly by linguistic coherence. The paper proposes replacing initial reasoning steps with discrete latent abstractions while retaining later text and training on randomized mixtures.
- Explicit chain-of-thought training improves reasoning and planning by making intermediate steps available as text.
- Lengthy textual reasoning traces impose computational costs because many tokens mainly maintain linguistic coherence rather than encode core reasoning information.
- The method replaces initial chain-of-thought tokens with VQ-VAE-derived discrete latent abstractions and leaves later tokens unchanged.
- Randomly varying the number of replaced tokens exposes models to latent-text mixtures and supports adaptation to unseen latent vocabulary items.
- Across planning, logical, and mathematical benchmarks, the approach consistently outperforms text-only reasoning-trace baselines across model architectures.
2 Related Work
Prior work studies explicit textual reasoning traces and latent-space reasoning, including continuous representations and pause tokens. These approaches motivate compressing reasoning while preserving useful intermediate computation.
- Explicit chain-of-thought prompting uses textual reasoning examples or generated alternatives to improve question answering.
- Research also trains language models on high-quality reasoning traces and synthesizes datasets with diverse intermediate steps across domains.
- Continuous latent-space methods reuse hidden states as subsequent inputs, allowing models to continue reasoning without representing every step as text.
3 Methodology
The methodology represents each sample as prompt, chain-of-thought, and solution tokens, then uses a two-stage procedure to obtain latent abstractions and train an LLM with them.
- Each input is decomposed into prompt tokens P, reasoning tokens C, and solution tokens S joined by concatenation.
- Training proceeds in two stages: learning a latent representation of reasoning steps, then training the LLM with that representation.
1. Learning latent discrete tokens to abstract the rea-
The method learns discrete latent reasoning codes with a VQ-VAE, then partially replaces the earliest chain-of-thought tokens while preserving later text. Randomized replacement trains the LLM on varied abstraction levels.
- Learning latent discrete tokens: With L = 16 and r = 16, 32 textual chain-of-thought tokens become 2 discrete latent tokens, while remaining tokens stay textual.
- Learning latent discrete tokens: The compression rate r = tc/tz controls how many latent tokens represent the reasoning sequence.
- Training with partial latent traces: For each example, randomized m values replace prefixes of the reasoning trace, exposing the model to varied mixtures of latent and textual tokens.
- Learning latent discrete tokens: A VQ-VAE encodes input chunks into latent embeddings, quantizes them using a codebook, and decodes the quantized embeddings back to text.
- Learning latent discrete tokens: The architecture includes a codebook, encoder, nearest-neighbor quantizer, prompt embedder, and decoder producing text from latent embeddings.
- Learning latent discrete tokens: Its objective combines reconstruction, vector-quantization, and commitment losses, keeping encoder outputs near codebook vectors while supporting decoding.
- Training with partial latent traces: The second stage trains next-token prediction on modified samples that replace only initial reasoning tokens, retaining later text and marking latent spans.
- Training with partial latent traces: The 50M-parameter VQ-VAE is used during data preparation, not LLM training or inference, where the LLM directly generates latent tokens.
4 Experiments
The experiments evaluate latent-token reasoning across synthetic planning and logical benchmarks, plus mathematical reasoning with fine-tuned Llama models. Across these settings, partial latent replacement improves accuracy while reducing generated reasoning length, with randomized replacement and a fixed left-to-right pattern performing particularly well.
- Evaluation Setup: The evaluation spans synthetic planning and logical benchmarks, as well as in-domain and out-of-domain mathematical reasoning tasks.T5 or GPT-2 models are trained from scratch for synthetic tasks, while Llama models are fine-tuned for mathematical reasoning.
- Evaluation Setup: The randomized procedure replaces an initial number of CoT tokens with latent codes while retaining later tokens as text.For each example, the procedure samples a replacement limit and replaces the first m tokens, where m is selected in increments of 16.
- Synthetic and Mathematical Results: The latent approach improves synthetic-task accuracy, including a 15% absolute gain on Keys-Finding Maze, 100% accuracy on ProntoQA, and a 17.5% gain over CoT on ProsQA.The CoT baseline reaches 77.5% on ProsQA.
- Synthetic and Mathematical Results: Across mathematical benchmarks and model sizes, the latent approach generally outperforms baselines, with average gains of +5.3 points for 8B, +2.9 points for 3B, and +3.7 points for another reported size.On tasks without improvement, performance remains comparable to the best baseline.
- Efficiency: The latent approach reduces generated response length by 17% on average relative to CoT while surpassing it in prediction accuracy.iCoT produces shorter responses but significantly worse performance than CoT, whereas the latent approach combines shorter traces with higher accuracy.
- Ablations and Analysis: AR-Replace outperforms alternative replacement strategies, while partial replacement provides text anchors that help connect latent and linguistic representations.The alternatives include All-Replace, Curriculum-Replace, and Poisson-Replace.
- Ablations and Analysis: The latent model attends more strongly than CoT to numbers and words representing mathematical operations, suggesting focus on tokens supporting the internal logical flow.Reported examples include numerical tokens and operation-related words such as “earns” and “cost.”
- Ablations and Analysis: On Llama-3.2-3B, accuracy rises from 25.2 at the baseline setting to 28.1 at r = 16 before falling to 27.9 at r = 32.Increasing r reduces token counts, while r = 16 is identified as the balance between compression and preservation of task-critical information.
5 Conclusion
The paper combines discrete latent and text tokens to compress initial reasoning steps while retaining textual details, improving reasoning token efficiency across domains.
- The approach compresses initial reasoning steps using discrete latent tokens obtained from VQ-VAE.
- Training integrates abstract representations with textual reasoning details to capture essential reasoning information more efficiently.
- Randomizing how many text tokens are compressed supports fast adaptation to unseen latent tokens.
- The evaluation spans multiple domains and reports performance above standard methods using complete textual reasoning traces.
Impact Statement
The paper identifies safety concerns from opaque latent reasoning tokens and describes a VQ-VAE decoder as a way to recover human-readable representations.
- Opaque latent tokens may raise safety concerns because their reasoning content is not directly human-readable.
- A VQ-VAE decoder can decode latent tokens into human-readable format.
- The keys-finding maze includes colored keys and doors, requiring an agent to collect matching keys before reaching the goal.
- The maze dataset contains 100k training examples, 500 validation examples, and 300 test examples.
- Maze responses represent paths with create-node entries and agent positions over successive time steps.
A.4 LLM experiments
The Llama experiments use a codebase for fine-tuning and select learning rates by validation error, while the appendix records mathematical notation used throughout the paper.
- The experiments fine-tune Llama models using the Llama Cookbook codebase.
- Learning rates of 10^-5, 2.5 × 10^-5, 5 × 10^-5, and 10^-4 are compared, with the lowest-validation-error choice retained.
- The final learning rate is 10^-5 for Llama-3.2-8B and 2.5 × 10^-5 for Llama-3.2-1B and Llama-3.2-3B.
- Appendix B summarizes the mathematical notations used throughout the paper.
C Details of Attention Weights Analysis
The examples compare latent-token responses with textual chain-of-thought responses on mathematical word problems, showing that both can produce the same answers with different reasoning representations.
- The latent response replaces initial reasoning text with latent codes while retaining later textual calculations.
- For the percentage problem, both the latent approach and the CoT baseline calculate a positive difference of 10.
- The latent approach solves the savings problem by computing the $250 shortfall and dividing by $10 daily earnings to obtain 25 days.
- The CoT baseline uses a fully textual explanation for the same savings problem.
D Other Text Generation Examples
The examples show latent tokens replacing initial reasoning while later text completes logical and mathematical solutions. Across examples, the generated continuations derive correct answers or conclusions from compact latent prefixes.
- Mathematical reasoning: A latent prefix initiates exponent-rule reasoning that evaluates 5^3 · 5^2 to 3125.The continuation applies a^m·a^n = a^(m+n), substitutes a = 5, and computes the result.
- Mathematical reasoning: A latent prefix initiates simplification of 15!/[2!(15−2)!], whose continuation concludes with 105.The shown continuation cancels common factors before reaching the final value.
- Mathematical reasoning: A latent prefix begins the denominator analysis for a rational graph whose vertical asymptotes are x = −3 and x = 2.The continuation factors x^2 + x − 6 as (x + 3)(x − 2) and identifies the two zeros.
- Logical reasoning: A latent prefix is followed by a textual proof that the graph condition implies no directed cycles.The continuation argues that a directed cycle would conflict with the absence of links to lower-numbered nodes.
E Interpreability Examples
Decoded latent tokens are generally interpretable and capture the initial conceptual step of reasoning, although their wording can contain grammatical errors. The examples show abstractions for proportional conversion and Euclidean-algorithm initiation.
- Interpretability: Decoded latent tokens generally capture long or verbose initial reasoning sentences that may trigger useful follow-up thinking.The paper reports that decoded latent content is generally interpretable.
- Mathematical examples: In the bottle example, the decoded prefix captures proportional conversion from an 8-ounce bottle to a 12-ounce bottle despite a grammar mistake.The subsequent text computes 187.5 calories from the proportion.
- Mathematical examples: In the GCD example, the decoded prefix abstracts initiation of the Euclidean algorithm and leads into the process determining the answer 9.The continuation performs successive divisions until the remainder is zero.
F Additional Experiments
Additional DART-MATH experiments report that the approach outperforms iCoT and Sol-Only while marginally outperforming CoT, and generates shorter reasoning traces than CoT.
- Accuracy: On DART-MATH, the approach surpasses the iCoT and Sol-Only baselines while marginally outperforming CoT.This result is summarized in Table F.1 for fine-tuning a Llama-3.1-8B model.
- Efficiency: The approach generates shorter reasoning traces than the CoT baseline.Table F.2 reports average generated-response token counts and trace-length reduction relative to CoT.