Training Continuous Chain of Thought Models: A Tale of Two Regimes
Artificial IntelligenceComputation and Language
Summary
The authors introduce C-MTP, a new method to make computer reasoning more efficient by compressing detailed thought steps into shorter, dense summaries. Unlike earlier methods that need slow, step-by-step training, C-MTP directly models these summaries as averages of detailed steps, speeding up learning. Their method performs better than some previous direct methods and nearly as well as slower indirect ones on short reasoning tasks. However, both direct and indirect methods struggle with longer, more complex reasoning tasks, showing current techniques have limits.
Authors
Varun Yerram, He He, Eunsol Choi
Abstract
Continuous Chain-of-Thought methods replace verbose reasoning traces with a short sequence of dense latent representations. Earlier continuous CoT methods indirectly supervise the latent representations such that its final state match that of verbose reasoning traces, requiring autoregressive, slow generation during training. We introduce C-MTP, a simpler, faster direct supervision approach that models each latent as an average of the embeddings in the CoT traces to be compressed. Our approach outperforms a prior direct supervision method that approximates the distribution of compressed tokens, and performs competitively to slower indirect supervision approaches in existing evaluation setup with simplified CoT traces (less than 100 tokens). Lastly, we extend the evaluation of Continuous CoT methods to complex tasks with longer reasoning traces ($\ge$ few hundreds reasoning tokens). We find both direct and indirect supervision training methods perform poorly (roughly 65\% performance drop) in this setting, revealing the limitations of current continuous CoT methods. The code and checkpoints are released at https://github.com/Varun221/cmtp_research