AS-ES Learning: Towards Efficient CoT Learning in Small Models

TL;DR

This work tackles the challenge of teaching chain-of-thought reasoning to small models without relying on costly data augmentation from LLMs. It introduces AS-ES learning, which decomposes CoT into Extractive Segments and Abstractive Segments and supports iterative generation via either dual-path (DS M) or unified-path (USM) training. The approach demonstrates data-efficient improvements on Math Word Problems and PET report summarization, offering explanations for when segmentation strategies and model sizes help or hinder performance. The findings suggest that small models’ CoT limitations stem more from data utilization paradigms than intrinsic capacity, with AS-ES providing a principled mechanism to exploit CoT information more effectively. This work has practical impact for deploying reasoning-capable models in resource-constrained settings and offers a foundation for further theoretical and empirical exploration of CoT dynamics.

Abstract

Chain-of-Thought (CoT) serves as a critical emerging ability in LLMs, especially when it comes to logical reasoning. Attempts have been made to induce such ability in small models as well by distilling from the data with CoT generated by Large Language Models (LLMs). However, existing methods often simply generate and incorporate more data from LLMs and fail to note the importance of efficiently utilizing existing CoT data. We here propose a new training paradigm AS-ES (Abstractive Segments - Extractive Segments) learning, which exploits the inherent information in CoT for iterative generation. Experiments show that our methods surpass the direct seq2seq training on CoT-extensive tasks like MWP and PET summarization, without data augmentation or altering the model itself. Furthermore, we explore the reason behind the inefficiency of small models in learning CoT and provide an explanation of why AS-ES learning works, giving insights into the underlying mechanism of CoT.

Figures (4)

• Figure 1: A. Description of PET-scan summarization task: Each blue section denotes the part of the report involving a particualr organ, with the green representing the related impression. Organ names are bolded and italicized, while light and dark green distinguish between impression segments. B. Description of Math Word Problem task: The blue section highlights the question, while light and dark green sections denote distinct segments of the answer.
• Figure 2: The workflow for labeling raw data as either ES or AS parts, followed by constructing ES/AS datasets.
• Figure 3: results of different hyperparameter settings for AS-ES learning. The BLUE metric values (blue line) correspond to the right Y-axis (secondary axis). $\gamma$ is set to $1.0$ in Figure \ref{['fig:hyp2']}, while $\beta$ is set to $1.0$ in Figure \ref{['fig:hyp3']}.
• Figure 4: Smoothed training curve for AS-ES Learning. In Figure \ref{['fig:dis1']} and Figure \ref{['fig:dis2']}, Grey, pink and yellow denote direct approach, interleaving segmentation and entropy segmentation respectively. Figure \ref{['fig:dis3']} and Figure \ref{['fig:dis4']} show the curve of the best training loss among time. All curves for ASM is displayed in red with ESM in bleu in Figure \ref{['fig:dis3']}. In Figure \ref{['fig:dis4']}, curves for ASM, ESM, USM are displayed in green, blue and pink, respectively.