Source-linked AI summary

CoT-Valve: Length-Compressible Chain-of-Thought Tuning

Xinyin Ma, Guangnian Wan, Runpeng Yu, Gongfan Fang, Xinchao Wang

arXiv:2502.09601v1cs.AIcs.CL

TL;DR

Long chain-of-thought reasoning improves model reasoning but raises inference costs and may allocate tokens inefficiently across task difficulties. CoT-Valve introduces parameter-space control and MixChain-based tuning to generate reasoning chains of varying lengths and compress them. Experiments report controllable, compressible reasoning with better performance than prompt-based control, including substantial token reductions on GSM8K and AIME.

  • Problem

    Long reasoning chains increase inference costs, while reasoning models may need long paths for hard tasks and shorter paths for easy tasks.

  • Method

    CoT-Valve manipulates a parameter-space update direction through LoRA and uses long-short chain data with precise tuning and progressive compression.

  • Results

    CoT-Valve enables controllable and compressible reasoning and achieves better performance than prompt-based control and previous chain-compression baselines.

  • Takeaways & Limitations

    A single model can dynamically generate reasoning chains of varying lengths, supporting inference-efficiency improvements within the reported evaluation scope.

  • Takeaways & Limitations

    The method does not yet design how to select segments for localized compression, leaving that segment-selection problem for future work.

Abstract

from arXiv · show

Chain-of-Thought significantly enhances a model's reasoning capability, but it also comes with a considerable increase in inference costs due to long chains. With the observation that the reasoning path can be easily compressed under easy tasks but struggle on hard tasks, we explore the feasibility of elastically controlling the length of reasoning paths with only one model, thereby reducing the inference overhead of reasoning models dynamically based on task difficulty. We introduce a new tuning and inference strategy named CoT-Valve, designed to allow models to generate reasoning chains of varying lengths. To achieve this, we propose to identify a direction in the parameter space that, when manipulated, can effectively control the length of generated CoT. Moreover, we show that this property is valuable for compressing the reasoning chain. We construct datasets with chains from long to short for the same questions and explore two enhanced strategies for CoT-Valve: (1) a precise length-compressible CoT tuning method, and (2) a progressive chain length compression approach. Our experiments show that CoT-Valve successfully enables controllability and compressibility of the chain and shows better performance than the prompt-based control. We applied this method to QwQ-32B-Preview, reducing reasoning chains on GSM8K from 741 to 225 tokens with a minor performance drop (95.07% to 94.92%) and on AIME from 6827 to 4629 tokens, with only one additional incorrect answer.

1 Introduction

CoT-Valve addresses the cost and task-difficulty mismatch of long reasoning chains by enabling one model to generate controllably short or long paths. It combines parameter-space length control with MixChain-based tuning and achieves controllable, compressed reasoning across models.

  • Motivation: Long reasoning chains improve complex-task reasoning but can contain redundant intermediate steps and increase inference costs.The paper motivates compression for simple tasks while retaining long reasoning for difficult ones.
  • Motivation: Current reasoning models may allocate excessive tokens to simple tasks while providing insufficient tokens for complex tasks.The proposed goal is therefore a model that can generate both long and short reasoning paths.
  • Method: CoT-Valve controls reasoning length by manipulating an update direction in parameter space through a LoRA branch.Large steps produce short sequences, whereas small steps preserve longer and more complex trajectories.
  • Method: MixChain pairs long and short reasoning paths for each question to refine tuning and support progressive chain compression.The dataset is constructed from reasoning chains of varying lengths.
  • Evaluation: One training run enables models ranging from LLaMA variants to QwQ-32B-Preview and DeepSeek-R1 to generate reasoning paths of varying lengths.The method also reports better results than previous chain-compression baselines.
  • Findings: CoT-Valve achieves controllable reasoning generation and state-of-the-art results for compressed CoT.The authors also report that short reasoning paths can sometimes outperform longer ones, while excessively long or short chains complicate CoT distillation.

2 Related Work

Related work covers methods that improve reasoning through explicit search, scoring, and chain-of-thought generation, alongside approaches aimed at reducing the cost of long reasoning chains. These efforts include distillation, hidden-state internalization, and implicit token-wise compression.

  • Chain-of-Thought: Chain-of-thought reasoning and test-time scaling have motivated self-reasoning, Best-of-N, beam search, and Monte Carlo Tree Search methods.Related approaches also evaluate complete solutions and reasoning-path quality.
  • Chain Compression: Chain compression targets the high inference cost of long-chain reasoning in reasoning models.Prior work explores distilling chains into System 1, internalizing reasoning steps in hidden states, and compressing token-wise generation implicitly.

3 Method

CoT-Valve controls reasoning-chain length by manipulating a parameter-space direction, then uses length-varied chains to improve controllability and compression. Its MixChain dataset supports precise length control and progressive shortening while retaining correct answers.

  • Length-Compressible CoT Tuning: CoT-Valve targets the limited control of prompt-based methods, which may generate over 350 tokens after requesting fewer than 20.The method is designed to generate both long and short reasoning paths with one tuned model.
  • Length-Compressible CoT Tuning: A parameter-space update ∆θ acts as an external, parameter-efficient branch that steers the model toward more concise reasoning chains while preserving the final answer.The update is interpreted as a task vector because models with and without ∆θ output the same final answer.
  • Length-Compressible CoT Tuning: Adjusting the update magnitude α interpolates between longer and shorter reasoning paths, while α > 1 extrapolates to shorter paths beyond the training range.Each α value corresponds to a different CoT length at inference.
  • Length-Compressible CoT Tuning: The method leaves segment selection for localized compression as future work, despite supporting finer-grained control than prompt-based approaches.This design would retain long-chain reasoning in selected thoughts while compressing simpler segments.
  • Construct the MixChain Dataset: MixChain associates each question with progressively shorter reasoning paths, avoiding repeated sampling while providing a spectrum of chain lengths.It can be constructed from human-labeled solutions or by deriving ∆θ between a base LLM and a corresponding reasoning model when explicit solutions are unavailable.
  • Improved Tuning for CoT-Valve: CoT-Valve++ adds length-specific training constraints across the update direction, using normalized β and question-specific mmin and mmax to improve compressibility and controllability.The enhanced objective is trained on MixChain and addresses inconsistency between single-position training and varying-length inference.
  • Improved Tuning for CoT-Valve: Progressive compression trains successively shorter reasoning paths at each iteration instead of directly training on the shortest chain.The strategy gradually reduces reasoning length and is compared conceptually with iterative pruning in model compression.

4 Experiments

Experiments evaluate CoT-Valve across model types, datasets, and length-transition settings, measuring accuracy, token count, and efficiency. Results show controllable generation, effective compression, and advantages over baselines in several settings.

  • Experimental settings: The evaluation covers long-to-short, short-to-long, and short-long-short CoT settings across QwQ, DeepSeek-R1, LLaMA, and Qwen models.Experiments use GSM8K and AIME24, with accuracy, answer-token count, and ACU reported.
  • Controllability: CoT-Valve matches prompt-based control while generating shorter chains, and extrapolated parameter updates produce still shorter reasoning paths.CoT-Valve++ further surpasses the baseline and generalizes better during extrapolation.
  • Compression results: 5155 to 4630 tokens on AIME preserves the baseline accuracy, despite training on an easier dataset.Greedy decoding was used to reduce randomness in the AIME experiments.
  • Observations: 55.50 accuracy with 267.0 tokens versus 52.69 with 759.3 tokens illustrates that shorter chains can outperform longer ones on simple datasets.This pattern is reported for LLaMA-3.2-1B and is not observed on more complex datasets.
  • Compression results: 92.19 to 94.92 accuracy improves short-CoT performance under progressive compression while token counts decrease across turns.The comparison uses five total epochs, with the final epoch using ground-truth data.

5 Conclusion

The paper proposes CoT-Valve to generate reasoning chains of varying lengths and constructs long- and short-chain data to improve controllability and compression efficiency.

  • CoT-Valve enables models to generate reasoning chains of varying lengths instead of relying on prompt control.
  • The authors construct a dataset containing both long and short reasoning chains to enhance controllability and compression efficiency.
  • Experimental results demonstrate the method's effectiveness for dynamic reasoning-chain control and CoT compression.
  • Future work could explore finer-grained control strategies to improve reasoning efficiency and model controllability.

A.1 Evaluation Metric.

The evaluation uses model-specific answer-extraction and matching procedures, including strict matching for LLaMA-3.1-8B and flexible matching for LLaMA-3.2-1B-Instruct.

  • LLaMA-3.1-8B is evaluated with strict matching because response repetition can cause flexible matching to extract incorrect numerical values.
  • LLaMA-3.2-1B-Instruct is evaluated using the flexible match metric.
  • For QwQ-32B-Preview, DeepSeek-R1-Distill-Llama-8B, and Qwen-2.5B-LIMO, answers enclosed within \boxed{} are extracted first.

A.2 Training Setting.

Training settings vary across models, datasets, and experiments, with LoRA used in several setups and full fine-tuning used for the Qwen2.5-32B-LIMO refinement.

  • LLaMA-3.1-8B uses eight A5000 24GB GPUs, batch size 64, peak learning rate 4e-5, cosine decay, and LoRA rank 32.
  • LLaMA-3.2-1B-Instruct uses eight A5000 24GB GPUs, with batch size 8 for CoT-Valve and 64 for other experiments.
  • QwQ-32B-Preview inference limits are 4192 tokens for GSM8K and 8192 tokens for AIME.
  • DeepSeek-R1-Distill-Llama-8B uses the MixChain-zero-shot-GSM8K dataset, batch size 128, at most five epochs, and a 30K-token inference limit.

A.3 Dataset Explanation

The paper constructs MixChain-C and MixChain-Z datasets using interpolated model directions and filters out incorrect responses; dataset quality depends strongly on the base model.

  • MixChain-C and MixChain-Z are constructed with selected α values, excluding all incorrect responses.
  • For LIMO, α ranges from 0.6 to 0.8, while other datasets use α values from 0.2, 0.4, 0.6, and 0.8.
  • MixChain-Z includes model pairs such as Qwen-32B-Instruct → QwQ-32B-Preview and LLaMA-3.1-8B → R1-Distill-Llama-8B.
  • The performance of the base model significantly influences MixChain-Z dataset quality.
  • Table 9 reports dataset statistics using the QwQ-32B-Preview tokenizer to count tokens.
  • Table 10 presents results for LLaMA-3.2-1B-Instruct trained with DoRA under different α interpolation values.

B More Analysis

The analysis finds that CoT-Valve controls reasoning length through parameter interpolation and that MLP-related components influence chain shortening more than attention projections. Prompt-based constraints remain imprecise, while CoT-Valve achieves higher accuracy with substantially shorter chains in the DoRA experiment.

  • DoRA Experiments: 55.72 accuracy with 257.7 tokens surpasses 54.5 accuracy with 772.7 tokens for DoRA.The CoT-Valve model uses α=0.5, while direct training on QwQ synthesized data uses α=1.0.
  • Module Analysis: Attention projections affect chain shortening less than MLP layers and the final attention projection.Query, key, and value projections had less effect than other linear projections in the LoRA module comparison.

C Examples

Examples illustrate CoT-Valve's ability to generate reasoning chains across lengths, including extrapolated short outputs. The appendix also documents discrepancies between prompt-specified and actual chain lengths on GSM8K.

  • Generated Examples: 103 tokens is achieved through extrapolation after shortening multiple reasoning rounds and streamlining the language.The original model's longer reasoning includes an extra reflection phase that is removed during compression.
  • Generated Examples: CoT-Valve produces generated CoT from short to long and includes extrapolation cases demonstrating generalization ability.The example compares generated solutions for the same question across different chain lengths.
  • Prompt-Control Examples: GSM8K prompt conditions show significant discrepancies from the number of generated tokens.Table 12 is specifically described as documenting these mismatches.
Loading 2502.09601v1…