Knowledge Distillation and Speculative Decoding within the Latent Reasoning Space
- Labardini, Luca
- Advisor(s): Shang, Jingbo
Abstract
Large Language Models (LLMs) increasingly rely on complex reasoning chains to solve search-intensive tasks. Recent advancements, such as the Chain of Continuous Thought (Coconut), have absorbed the reasoning process from natural language tokens into the continuous latent space, allowing for the dense exploration of multiple reasoning paths. However, inference costs remain high. While speculative decoding reduces latency for standard LLMs, it relies on discrete token distributions, making it incompatible with continuous latent reasoning. This work proposes Speculative Reasoning, a framework for adapting speculative decoding to the continuous latent space. We introduce a knowledge distillation training recipe that enables a smaller draft model to generate latent thought vectors aligned with a larger target model. Additionally, we adopt a hybrid acceptance criterion utilizing cosine similarity to verify latent vectors and rejection sampling for natural language tokens. This ensures the integrity of the reasoning path while providing speedup for both latent and language token generation. We evaluate the proposed method using a GPT-2 Medium teacher and GPT-2 Mini student on a generated ProntoQA dataset. Experiments demonstrate a perfect acceptance rate on latent thoughts and a 75% acceptance rate on natural language tokens. This results in a 2.15x speedup on latent reasoning and a 1.49x overall speedup. These findings confirm that speculative decoding can be effectively extended to the continuous reasoning paradigm, offering a path toward efficient inference for next-generation reasoning models.