Source-linked AI summary

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

Nuwa Xi, Yuhan Chen, Sendong Zhao, Haochun Wang, Bing Qin, Ting Liu

arXiv:2403.01969v1cs.CL

TL;DR

Small models struggle to learn CoT efficiently because existing approaches rely on direct seq2seq training and additional LLM-generated data rather than fully exploiting existing CoT. The paper introduces AS-ES learning, which segments CoT into extractive and abstractive parts for iterative generation, and reports effectiveness across tasks and model sizes without data augmentation. It further argues that training-paradigm inefficiency, rather than inherent small-model capacity, explains the limitation, while noting that optimal settings vary by dataset.

  • Problem

    Existing CoT distillation methods often generate more LLM data and use direct seq2seq training, leaving efficient utilization of existing CoT and small-model capacity unresolved.

    Source: §1 · p. 1
  • Method

    AS-ES learning segments CoT into Extractive Segments and Abstractive Segments and uses them in an iterative generation process without additional data.

    Source: §1 · pp. 1–2, §6 · p. 8
  • Results

    AS-ES learning improves performance over direct seq2seq training on CoT-intensive MWP and PET summarization tasks and generalizes across model sizes and tasks.

    Source: §1 · p. 2, §6 · p. 8
  • Takeaways & Limitations

    The findings indicate that small-model CoT limitations stem from the training paradigm rather than inherent capacity, while one unified model can outperform two separate models in the studied setting.

    Source: §6 · p. 9, Limitations · p. 9

    The optimal AS-ES hyperparameter settings may vary across datasets, and the reported one-model advantage may not hold when DSM is trained simultaneously.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces AS-ES learning, which decomposes CoT into extractive and abstractive segments for iterative training, targeting efficient CoT learning in small models without additional data. Experiments on MWP and PET summarization examine whether data utilization, rather than intrinsic model capacity, limits small-model CoT learning.

1 Introduction

  • Existing approaches commonly generate more LLM-produced CoT data or use direct seq2seq training, overlooking both the cost of augmentation and small models’ difficulty learning complex reasoning.Some methods decompose reasoning or train separate models, but augmented CoT still adds data rather than fully exploiting existing information.Source: §1 · p. 1
  • AS-ES learning decomposes CoT into Extractive Segments that provide context and Abstractive Segments that infer unstated insights, then uses them for iterative generation without additional data.The paradigm is designed to maximize the latent potential of small models across model sizes and tasks.Source: §1 · pp. 1–2
  • The study tests whether one small model can perform both extraction and abstraction, and whether CoT-learning weaknesses reflect inefficient training paradigms rather than inherent capacity limits.It compares unified and two-model strategies for AS/ES generation.Source: §1 · p. 2
  • The evaluation covers PET impression generation and Math Word Problems, whose CoT structures differ between parallel processing across body regions and sequential step-by-step solutions.Figure 1 distinguishes report-organ sections and impression segments for PET, and question-answer segments for MWP.Source: §1 · p. 2, Figure 1 · p. 2
  • The paper claims AS-ES learning surpasses direct seq2seq training on CoT-intensive tasks and provides a theoretical explanation of its efficacy and CoT dynamics.The claimed benefits include data efficiency, applicability across model sizes and tasks, and insight into why small models struggle with CoT.Source: §1 · p. 2

2 Related Work

Prior work transfers CoT abilities to small models mainly by distilling LLM-generated rationales or solutions, often requiring additional teacher-generated data. The paper instead presents an automatic strategy intended to disentangle deduction efficiently without extra LLM-generated data.

Deduction in NLP Tasks

  • Deduction research has used end-to-end learning or separate validation modules, but prior alternatives may require human annotation or hand-engineered world models.The paper describes its own strategy as fully automatic.Source: §2 · p. 2

CoT in Small Models

  • CoT distillation methods for small models commonly use LLM-generated rationales or multi-step solutions as training data.Examples include rationale generation, consistency-based multiple rationales, and two-stage rationale-and-answer distillation.Source: §2 · pp. 2–3
  • These approaches require an LLM teacher and therefore add time and computational cost through additional generated data.The paper positions its method as requiring no additional LLM-generated data.Source: §2 · p. 3

3 Method

AS-ES learning segments chain-of-thought into extractive and abstractive parts, organizes them for iterative generation, and supports either separate or unified models. The method uses multiple segmentation criteria and adapts data construction to task-specific constraints.

3.1 AS-ES Segmentation

  • AS-ES segmentation splits each CoT into punctuation-delimited sub-sentences and scores them with metrics to classify extractive and abstractive segments.The method experiments with entropy-, location-, loss-, and similarity-oriented strategies; entropy-based segmentation treats lower model uncertainty as indicative of extractive content.Source: §3.1 · p. 3, Similarity-oriented Segmentation… · p. 3

Similarity-oriented Segmentation (bleu/rouge)

  • Similarity-oriented segmentation uses BLEU and ROUGE to measure how closely each CoT segment resembles the query.Extractive segments typically share more similarity with the original context, whereas abstractive segments introduce new reasoning.Source: Similarity-oriented Segmentation… · p. 3

3.2 AS-ES Dataset Construction

  • The AS-ES dataset merges adjacent segments of the same type and creates iterative inputs that generate either the next extractive or abstractive segment.The workflow labels raw data as ES or AS before constructing the corresponding datasets, with generation stopping at task-specific markers.Source: §3.2 · p. 4, Figure 2 · p. 4
  • PET data construction limits sequence length by splitting scans anatomically and selectively incorporating normal findings into the extractive dataset.Normal findings receive “No obvious anomaly” labels, and their inclusion is controlled by γ, the proportion of normal findings retained.Source: §3.2 · p. 4

3.3 AS-ES Learning

  • Dual-path learning trains separate extraction and reasoning models, whereas uni-path learning consolidates AS-ES data into one model to simplify training and reduce resources.The unified model performs the same iterative generation procedure while handling both extractive and abstractive responsibilities.Source: §3.3 · pp. 4–5

4 Experiment

Experiments on MWP and PET summarization evaluate AS-ES across datasets, model sizes, segmentation strategies, and unified or dual-model training. AS-ES improves task performance with strategy-dependent behavior, often yielding higher accuracy despite lower textual similarity and suggesting that training strategy contributes more than model capacity alone.

4.1 Dataset

  • The experiments use GSM8K-derived CoT data for MWP and the newly introduced 11.6k-pair cPET-11K Chinese PET/CT dataset for PET summarization.Models include Flan-T5 variants for MWP and mT5 for PET, with Flan-T5-base and mT5-base as defaults.Source: §4.1 · p. 5, §4.2 · p. 5

4.2 Implementation

  • Evaluation prioritizes accuracy for MWP and BLEU for PET, while checkpoint selection compares best training loss, validation loss, and validation BLEU.The authors note that standard validation metrics may not track iterative-generation quality or MWP accuracy directly.Source: §4.2 · p. 6

4.3 Overall Performance of AS-ES Learning

  • AS-ES learning improves performance on both MWP and PET when paired with appropriate segmentation and training strategies.Entropy-oriented segmentation generalizes across model sizes and tasks, while interleaving segmentation is more targeted.Source: §4.3 · p. 6

4.3.1 Lower BLEU with Higher Accuracy

  • AS-ES often produces lower BLEU but higher MWP accuracy than direct training, consistent with different correct solutions rather than memorized textual matches.The authors associate this pattern with finer-grained logical reasoning through a step-by-step process.Source: §4.3.1 · pp. 6–7

4.3.2 Different Model Size Reacts Differently to Segmentation Strategy

  • Entropy-oriented segmentation works best for Flan-T5-large, whereas interleaving segmentation works best for Flan-T5-small, making effectiveness dependent on model capacity and dataset quality.Smaller models may handle entropy-based segmentation less reliably, especially for highly sequential MWP reasoning.Source: §4.3.2 · p. 7

4.3.3 Segmentation Works Differently for Different Tasks

  • MWP generally favors interleaving segmentation, while PET summarization benefits from entropy segmentation because their CoT structures differ.MWP alternates contextual premises and reasoning steps, whereas PET combines observations with collective diagnostic insights without a consistent interleaving pattern.Source: §4.3.3 · p. 7

4.3.4 USM v.s. DSM: Is One Model Enough?

  • A unified model outperforms the two-model approach for both MWP and PET in AS-ES learning, despite both approaches experiencing error accumulation.The authors attribute this to integrated training of extraction and reasoning, while separate optimization can limit the tested checkpoint’s capacity.Source: §4.3.4 · p. 7

5 Discussion

The discussion shows that segmentation strategy and training dynamics are central to AS-ES performance. Hyperparameters matter within task-specific ranges, while lower loss boundaries provide an explanation for AS-ES gains over direct training.

5.1 Effect of Segmentation

  • Interleaving, loss-based, and fine-tuned-model entropy segmentation all outperform direct training, whereas BLEU/ROUGE segmentation does not provide the same improvement.Entropy from pretrained rather than fine-tuned models introduces no further improvement.Source: §5.1 · p. 7, Table 3 · p. 7
  • Figure 3 reports AS-ES learning results under different hyperparameter settings, with BLUE metric values on the right secondary Y-axis; γ=1.0 in Figure 3b and β=1.0 in Figure 3c.Source: Figure 3 · p. 8

5.2 Effect of Hyperparameters

  • A β threshold around 1 yields the best balance between extractive and abstractive segments for entropy segmentation across tasks and training strategies.The result is reported for both unified- and dual-model training.Source: §5.2 · p. 8
  • γ=1.0 reaches the lowest performance curve point for PET, while overall performance is less sensitive to γ than to segmentation and training strategy.The authors interpret γ=1.0 as enough normal-finding context without diluting anomaly-focused information.Source: §5.2 · p. 8, Figure 3 · p. 8

5.3 Why AS-ES Learning Works?

  • ESM reaches a generally lower best-training-loss boundary than ASM, supporting the view that extraction is easier for small models than logical reasoning.The direct approach has a substantially higher lowest loss boundary than AS-ES learning.Source: §5.3 · p. 8
  • Task-specific segmentation preferences follow the loss boundaries: interleaving is lower for MWP, whereas entropy segmentation is lower for PET.The authors connect these lower boundaries to the corresponding performance differences between tasks.Source: §5.3 · p. 8
  • Overall, AS-ES learning is explained by achieving a generally lower loss boundary than direct training.This mechanism also helps explain why segmentation strategy changes with the structure of the task.Source: §5.3 · p. 8

6 Conclusion

AS-ES learning improves data-efficient CoT distillation for small models through iterative generation without additional data, generalizing across model sizes and tasks. The paper attributes small models’ CoT limitations to training-paradigm inefficiency rather than inherent capacity.

  • AS-ES learning improves small-model CoT distillation through iterative generation without incorporating additional data.Source: §6 · p. 8
  • Its effectiveness generalizes across different model sizes and tasks.Source: §6 · p. 9
  • The authors conclude that small models’ CoT-learning limitation stems from the training paradigm rather than inherent capacity.Source: §6 · p. 9

Limitations

The paper’s experiments focus on AS-ES learning within the direct training approach, and optimal hyperparameter settings may vary across datasets.

  • The experiments mainly evaluate AS-ES learning based on the direct training approach.Source: Limitations · p. 9
  • Although a general hyperparameter range is identified, dataset-specific optimal AS-ES settings may require additional exploration.Source: Limitations · p. 9
  • The observed advantage of one unified model over two separate models may not hold when a solution for DSM is trained simultaneously.Source: Limitations · p. 9

Ethical Considerations

The study introduces a Chinese PET report-impression dataset under an approved ethics-review protocol, using deidentified datasets for research only.

A Details of Implementation

Implementation uses fixed training settings across experiments, standard dataset splits, repeated runs, and primarily A100 GPUs.

Loading 2403.01969v1…