cs.LGMay 31, 2026

ThinkSwitch: Context Distillation with LoRA and Weight Interpolation for Specific-Purpose Reasoning Tasks

Authors: Dhruv SainiRohan Pandey

Organizations: Bellevue High School · DigitalOcean

Abstract

Large language models often improve on difficult tasks by spending inference-time compute on a reasoning trace before producing the final answer. That extra computation can be useful, but it also raises latency, token cost, and deployment complexity. We introduce \textbf{ThinkSwitch}, a low-compute procedure for co-training paired instruct and thinking checkpoints. Starting from compatible Qwen3-4B instruct and thinking models, each iteration asks the thinking checkpoint to generate answers, removes the reasoning trace, distills the answer-only pairs into the instruct checkpoint with QLoRA, and reconstructs a thinking checkpoint with spherical weight interpolation. The only human-supplied inputs are task prompts; the labels are generated by the model itself. On a 30-question AIME 2026 evaluation, ThinkSwitch improves the instruct checkpoint from 10/30 to 20/30 and the thinking checkpoint from 14/30 to 22/30. On a 30-question PubMedQA subset, it improves the instruct checkpoint from 13/30 to 18/30 and the thinking checkpoint from 18/30 to 25/30. The complete experiment uses 15 training prompts per domain and costs $2.86 on a single cloud RTX 3070. The results are small-scale, but they indicate that targeted distillation loops can move part of the benefit of explicit reasoning into weights while preserving a separate thinking mode.

Explore similar work

Nov 7, 2025cs.AI

CoT-X: An Adaptive Framework for Cross-Model Chain-of-Thought Transfer and Optimization

Long Chain-of-Thought (CoT) traces can improve reasoning accuracy, but repeatedly generating them is costly for smaller or latency-constrained language models. This paper studies a practical alternative: produce a rich rationale once with a capable \emph{thinking} model, compress it, and reuse the compressed trace as context for a cheaper \emph{answering} model. We introduce CoT-X, an adaptive framework for cross-model CoT transfer. CoT-X segments reasoning traces into semantic units, scores their diagnostic and logical importance, selects budget-feasible evidence paths, and reconstructs a coherent compressed rationale for the answering model. On 7,5017,501 Japanese medical licensing questions spanning 1010 specialties, CoT-X improves accuracy over direct truncation by up to 40.5%40.5\% under the same token budget, with the largest gains at 6464--256256 tokens. Across 6464 thinking--answering pairs from eight DeepSeek-R1 and Qwen3 models (1.5B--32B parameters), reasoning transfer is most reliable within a model family, yet remains effective across families once compression normalizes the trace. A Gaussian Process Bayesian optimization layer finds near-optimal model--budget configurations with 1515 evaluations rather than an exhaustive search over all 6464 pairs, reducing evaluation cost by 84%84\%. These results show that reasoning quality, token budget, and model compatibility can be optimized jointly, making CoT-style reasoning more practical under realistic deployment constraints.
Ziqian Bi, Yinzhi Wang, Tianyang Wang +6
Jul 11, 2026cs.CL

Structured Thoughts For Improved Reasoning And Context Pruning

Large language models (LLMs) excel at generating long chains of thought, but long reasoning traces are often verbose and memory-inefficient. In this work, we introduce Structured Thoughts, a framework that organizes reasoning into alternating <try> and <outcome> blocks: <try> captures exploratory scratch work, while <outcome> contains the distilled conclusion of that step. We construct a dataset of structured thoughts by segmenting reasoning traces into <try> blocks and prompting an LLM to summarize each step into its corresponding <outcome>. Fine-tuning pretrained foundation models on this reformatted data produces models that adopt the structured reasoning style, leading to performance gains of up to 8.08% on reasoning benchmarks compared to standard SFT. The explicit structure also enables context pruning: after each <try>/<outcome> pair, the <try> can be pruned, allowing the model to retain conclusions without keeping the full scratch work in the context. A proof-of-concept pruning implementation achieves an average of 85% memory / context savings with an 8.67% performance drop across mathematical tasks.
Zain Sarwar, Supriyo Chakraborty, Berkcan Kapusuzoglu +5
Jun 3, 2026cs.CL

LoRi: Low-Rank Distillation for Implicit Reasoning

Implicit chain-of-thought (iCoT) methods aim to internalize reasoning in large language models, but often underperform explicit CoT prompting. We empirically find that hidden-state reasoning trajectories exhibit low-rank structure. Motivated by this observation, we propose a low-rank distillation framework that transfers reasoning by aligning teacher and student trajectories in a shared low-rank tensor subspace using first- and second-order statistics. The resulting formulation captures the global structure of reasoning while supporting a compact latent reasoning process. We evaluate the method across multiple model families, including LLaMA and Qwen, at different scales on mathematical reasoning benchmarks. Our approach consistently improves performance, especially on challenging multi-step tasks, approaching explicit CoT accuracy and outperforming prior iCoT distillation methods.
Ryan Solgi, Jiayi Tian, Zheng Zhang