# Enhancing Generalization in Chain of Thought Reasoning for Smaller Models

arXiv: https://arxiv.org/abs/2501.09804 · HTML: https://mlsanity.com/paper/2501.09804/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Prompting and chain-of-thought reasoning, one of 17 directions in LLM reasoning.
- Builds on 33 of its 35 references that are in the map.
- First posted on arXiv 16 Jan 2025.

## Placed in

- LLM reasoning / [Prompting and chain-of-thought reasoning](https://mlsanity.com/lr/prompting-and-chain-of-thought-reasoning/)

## Builds on (most cited first, 30 of 33)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Scaling Instruction-Finetuned Language Models](https://mlsanity.com/paper/2210.11416/) (2022)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [GPT Understands, Too](https://mlsanity.com/paper/2103.10385/) (2021)
- [P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks](https://mlsanity.com/paper/2110.07602/) (2021)
- [Large Language Models Can Self-Improve](https://mlsanity.com/paper/2210.11610/) (2022)
- [Large Language Models Are Reasoning Teachers](https://mlsanity.com/paper/2212.10071/) (2022)
- [Adversarial Training Methods for Semi-Supervised Text Classification](https://mlsanity.com/paper/1605.07725/) (2016)
- [FreeLB: Enhanced Adversarial Training for Natural Language Understanding](https://mlsanity.com/paper/1909.11764/) (2019)
- [SMART: Robust and Efficient Fine-Tuning for Pre-trained Natural Language Models through Principled Regularized Optimization](https://mlsanity.com/paper/1911.03437/) (2019)
- [Specializing Smaller Language Models towards Multi-Step Reasoning](https://mlsanity.com/paper/2301.12726/) (2023)
- [WARP: Word-level Adversarial ReProgramming](https://mlsanity.com/paper/2101.00121/) (2021)
- [Distilling Reasoning Capabilities into Smaller Language Models](https://mlsanity.com/paper/2212.00193/) (2022)
- [Domain Adaptation via Prompt Learning](https://mlsanity.com/paper/2202.06687/) (2022)
- [Explanations from Large Language Models Make Small Reasoners Better](https://mlsanity.com/paper/2210.06726/) (2022)
- [Adversarial Training for Large Neural Language Models](https://mlsanity.com/paper/2004.08994/) (2020)
- [InfoBERT: Improving Robustness of Language Models from An Information Theoretic Perspective](https://mlsanity.com/paper/2010.02329/) (2020)
- [Generalization v.s. Memorization: Tracing Language Models' Capabilities Back to Pretraining Data](https://mlsanity.com/paper/2407.14985/) (2024)
- [ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators](https://mlsanity.com/paper/2003.10555/) (2020)
- [PERT: Pre-training BERT with Permuted Language Model](https://mlsanity.com/paper/2203.06906/) (2022)

## Built on by (most cited first, 5 of 5)

- [Compositional Generalization from Learned Skills via CoT Training: A Theoretical and Structural Analysis for Reasoning](https://mlsanity.com/paper/2502.04667/) (2025)
- [Quantification of Large Language Model Distillation](https://mlsanity.com/paper/2501.12619/) (2025)
- [MENTOR: Reinforcement Learning via Flexible Teacher-Optimized Rewards for Tool-Use Distillation](https://mlsanity.com/paper/2510.18383/) (2025)
- [Does Reasoning Help LLM Agents Play Dungeons and Dragons? A Prompt Engineering Experiment](https://mlsanity.com/paper/2510.18112/) (2025)
- [Thinking Economically: A Hierarchical Framework for Adaptive-Complexity Reasoning in LLMs](https://mlsanity.com/paper/2606.01168/) (2026)

## Abstract (arXiv)

Chain-of-Thought (CoT) reasoning in smaller language models is a challenging natural language process problem yet highly desirable in many real-life applications. Existing CoT knowledge distillation methods often suffer from overly conservative memorization in smaller LLMs, leading to low generalization confidence. As fully preserving the CoT ability of teacher model is impossible, we hypothesize that adversarial CoT fine-tuning is crucial for developing smaller LLM with robust CoT generalization. To this end, we propose textitPRompt-Assisted Domain-Adversarial fine-tuning (PRADA), a principled fine-tuning framework that integrates diverse CoT domains. Specifically, PRADA pioneers two CoT improvements in smaller LLM: (1) Recovering the domain-invariant feature insight which typically lost during distillation with domain adversarial fine-tuning; (2) Enhancing the domain adaptability of CoT prompt engineering by employing domain-adversarial approaches. We theoretically demonstrate the effectiveness of our approach and empirically show that it significantly outperforms the state of the arts in a wide range of tasks. Moreover, our empirical findings reveal that the smaller LLM, when leveraging PRADA, aligns closely with domain knowledge, thereby improving the explainability of our approach.