LoRi: Low-Rank Distillation for Implicit Reasoning

The gist: LoRi proposes a low-rank distillation framework that transfers reasoning from explicit chain-of-thought into compact implicit latent processes by aligning teacher and student trajectories in a shared low-rank subspace using first- and second-order statistics.

Background and Motivation

Implicit chain-of-thought (iCoT) methods aim to internalize reasoning in large language models, but often underperform explicit CoT prompting The paper empirically finds that hiddenstate reasoning trajectories exhibit low-rank structure Motivated by this observation, the authors propose a low-rank distillation framework that aligns the global geometry of the teacher’s reasoning trajectory through low-rank statistical representations By stacking hidden states across layers and CoT tokens, they found that the normalized cumulative singular values grow rapidly with rank, indicating that reasoning trajectories are well-approximated by low-dimensional subspaces

The LoRi Method

LoRi is a distillation framework that transfers reasoning from explicit chain-of-thought into a compact implicit latent process The method combines two complementary objectives: rationale-level alignment, which preserves the global geometry of the teacher’s trajectory, and anchor-level alignment, which regularizes the transition from latent reasoning to answer generation The composite objective trained is L = LLR + λLCE, where LCE supervises the student’s explicit output sequence and LLR aligns the student’s latent reasoning states with the teacher’s hidden representations in a low-rank subspace

Low-Rank Rationale Alignment (Lrationale)

To transfer reasoning capability, LoRi aligns the student’s hidden states with the teacher’s reasoning trajectory in a low-rank subspace This is achieved by projecting high-dimensional tensors into a shared low-rank subspace using Tucker decomposition to extract low-rank factor matrices via SVD The rationale-level loss, Lrationale, aligns student and teacher representations by matching the first- and second-order statistics of their projections in the shared low-rank space This objective transfers the global structure of the teacher’s reasoning process without tokenlevel alignment and remains invariant to reasoning length

Anchor-Level Alignment (Lanchor)

While Lrationale captures global reasoning structure, anchor-level alignment constrains the transition from latent reasoning to answer generation at a fixed position corresponding to an answer prediction token The anchor-level loss, Lanchor, is defined as 1/L sum of (U(a)⊤H hl − U(a)⊤H h′l)2 across layers l=1 to L This formulation leverages a shared low-rank subspace learned across training samples, ensuring the student’s representation at the answer prediction point aligns with the dominant structure of the teacher’s hidden states at that location

Implications and Results

The low-rank structure provides a geometric explanation for why long chain-of-thought reasoning can be compressed into a shorter latent trajectory The formulation is invariant to the length of the reasoning trajectory because alignment is performed through aggregated statistics rather than step-wise correspondence LoRi consistently outperforms prior iCoT methods, improving accuracy by up to ∼12% and achieving strong gains on GSM8K-Hard [up to ∼10%, as shown in Fig. 1 (b)] across model scales Furthermore, LoRi-E achieves performance comparable to SFT-CoT across all models and benchmarks, suggesting the proposed distillation method preserves explicit reasoning capability while enabling both implicit and explicit reasoning at inference time The method also significantly reduces training cost by using a two-stage procedure where teacher-derived quantities are precomputed once

Ablation Studies

Ablation on latent reasoning steps shows that performance consistently improves from K = 2 to K = 5, suggesting that a minimum number of latent reasoning iterations is needed to capture the underlying reasoning dynamics Ablation on training sample size indicates that accuracy improves rapidly in the lowdata regime and saturates after roughly 128 samples across all model scales The results support the view that distillation only needs to recover the dominant reasoning structure rather than the full trajectory The overall largest benefits are observed in smaller models, while the effects become more nuanced at larger scales

Conclusion

LoRi outperforms prior implicit reasoning methods across model families, scales, and benchmarks and reduces the gap of iCoT with explicit CoT while preserving the efficiency advantages of implicit reasoning The results also support a geometric perspective of reasoning: long chain-of-thought trajectories can be effectively compressed by preserving their low-rank structure This perspective provides a principled and efficient alternative to existing iCoT distillation approaches and suggests that the essential dynamics of reasoning are governed by a low-dimensional subspace that can be exploited for both learning and inference

Limitations

The proposed formulation is motivated by an empirical observation that hidden-state reasoning trajectories exhibit strong low-rank structures across models and benchmarks While LoRi improves reasoning accuracy and training efficiency, the theoretical relationship between low-rank representation geometry and reasoning capability is still not fully understood Future work could further investigate the geometry of reasoning trajectories and the role of low-dimensional structure in CoT reasoning from a theoretical perspective

References

Jinze Bai et al. 2023. Qwen technical report. arXiv:2309.16609 Yoshua Bengio, Aaron Courville, and Pascal Vincent. 2013. Representation learning: A review and new perspectives Zhuotong Chen, Qianxiao Li, and Zheng Zhang. 2022. Self-healing robust neural networks via closed-loop control Zhuotong Chen, Zihu Wang, Yifan Yang, Qianxiao Li, and Zheng Zhang. 2024. PID control-based selfhealing to improve the robustness of large language models Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems Lieven De Lathauwer, Bart De Moor, and Joos Vandewalle. 2000. A multilinear singular value decomposition Yuntian Deng, Yejin Choi, and Stuart Shieber. 2024a. From explicit cot to implicit cot: Learning to internalize cot step by step Yuntian Deng, Kiran Prsad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, and Stuart Shieber. 2024b. Implicit chain-of-thought reasoning via knowledge distillation Abhimanyu Dubey, Abhinav Jauhri, Ankur Pandey, et al. 2024. The llama 3 herd of models Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. 2016. Testing the manifold hypothesis Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. 2022. Pal: Program-aided language models Noah Golowich, Allen Liu, and Abhishek Shetty. 2025. Sequences of logits reveal the low rank structure of language models Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. 2024. Training large language models to reason in a continuous latent space Zhenghao He, Guangzhi Xiong, Bohan Liu, Sanchit Sinha, and Aidong Zhang. 2026. Reasoning beyond chain-of-thought: A latent computational mode in large language models Anna Kuzina, Maciej Pioro, Paul N. Whatmough, and Babak Ehteshami Bejnordi. 2025. Kava: Latent reasoning via compressed kv-cache distillation Bo Li, Guanzhi Deng, Ronghao Chen, Junrong Yue, Shuo Zhang, Qinghua Zhao, Linqi Song, and Lijie Wen. 2025a. Rema: A unified reasoning manifold framework for interpreting large language models Jindong Li, Yali Fu, Li Fan, Jiahong Liu, Yao Shu, Chengwei Qin, Menglin Yang, Irwin King, and Rex Ying. 2025b. Implicit reasoning in large language models: A comprehensive survey Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. 2024. Chain of thought empowers transformers to solve inherently serial problems Zicheng Lin, Tian Liang, Jiahao Xu, Qiuzhi Liu, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, and Zhaopeng Tu. 2025.

1. Bold header: Low-rank iCoT Distillation for Efficient Reasoning

The improved system can transfer complex reasoning from a computationally expensive explicit Chain-of-Thought (CoT) process into a compact implicit latent process, specifically by aligning teacher and student trajectories in a shared low-rank subspace using first- and second-order hidden-state statistics. This allows the student to achieve length invariant distillation, meaning it can capture the global structure of reasoning while supporting a compact latent reasoning process.

2. Bold header: Improved Reasoning Accuracy on Hard Tasks

The system exhibits enhanced performance on challenging problems by aligning with the dominant geometric structure of reasoning. The paper notes that LoRi consistently outperforms prior iCoT methods, improving accuracy by up to ∼12% and achieving strong gains on GSM8K-Hard, suggesting it is particularly effective where reasoning is most complex.

3. Bold header: Enhanced Inference Latency Efficiency

The system can maintain the efficiency benefits of implicit reasoning during inference while boosting performance. Specifically, iCoT consistently achieves lower latency across all model scales and LoRi maintains this benefit, with iCoT being about 6.9× faster for LLaMA-3.2-1B.

4. Bold header: Scalable and Cost-Effective Distillation Pipeline

The method improves training efficiency by using a two-stage procedure where teacher-derived low-rank factors are precomputed once, after which the student is fine-tuned independently on a subset of the training data, which substantially reduces training cost. This makes distillation a more scalable and accessible approach for implicit reasoning distillation.

5. Bold header: Robustness to Trajectory Length

The system's reliance on statistical matching rather than step-by-step correspondence ensures robustness against trajectory length. The formulation is explicitly described as being invariant to the length of the reasoning trajectory, enabling it to transfer long CoT reasoning into short latent trajectories.