UniCoTT: A Unified Framework for Structural Chain-of-Thought Distillation

Authors: Xianwei Zhuang, Zhihong Zhu, Zhichang Wang, Xuxin Cheng, Yuexian Zou
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Chain-of-Thought; Structural Thought; Distillation; Unified Framework

Abstract:
Chains of thought (CoTs) have achieved success in enhancing the reasoning capabilities of large language models (LLMs), while their effectiveness is predominantly observed in LLMs.  Existing solutions methods adopt distillation to inject chain-of-thought capabilities into small models (SLMs). However, they: (1) can not guarantee the rationality of the generated explanation due to hallucinations;  (2) ignore diverse structures of CoT during knowledge transfer. In this paper, we propose a unified CoT distillation framework termed UniCoTT for considering diverse structural CoTs (\emph{i.e.}, chain, tree, and graph). UniCoTT contains two core strategies: iterative construction for structured CoTs and the structural constraint strategy. Specifically, UniCoTT prompts LLMs to iteratively produce accurate explanations with answers and unifies structured explanations as UniCoT which is seen as a bridge for knowledge transfer. Furthermore, UniCoTT utilizes the proposed unified supervised learning and structural consistency learning strategies to transfer knowledge of structured CoT to SLMs. Experimental results show that UniCoTT can significantly improve the performance of SLMs on multiple datasets across different NLP tasks. Our code is available at https://github.com/mengchuang123/UniCoTT.

Decision: accept (poster)
Reviewer scores: 8, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730094676885

thread_type:
official_review

<thread_object>
creator: reviewer_546m

<content>
summary:
This paper  introduces UniCoTT, a teacher-student framework aimed at transferring complex reasoning abilities from large language models (LLMs) to smaller language models (SLMs). UniCoTT extends traditional chain-of-thought (CoT) reasoning by leveraging diverse structured reasoning paths, such as chains, trees, and graphs, within a unified distillation process. This approach involves iterative CoT construction, node-level supervised contrastive learning, and structural consistency learning to reinforce reasoning capabilities in SLMs. Experimental results on factual reasoning, multiple-choice QA, and natural language understanding tasks demonstrate that UniCoTT outperforms existing methods, enhancing SLM performance across several benchmarks.

strengths:
- Extends CoT reasoning with diverse structures, which broadens the reasoning capabilities of SLMs.
- Implements structural consistency and contrastive learning, effectively aligning SLMs with complex CoT reasoning paths.
- Demonstrates superior performance on multiple tasks, showing effectiveness and generality in knowledge transfer.

weaknesses:
UniCoTT’s increased complexity and computational requirements could make real-world deployment challenging. To be fair, as distillation strategy proposed in this paper uses three types of reasoning and more compute to create dense supervision. The baselines like CoT may also uses more compute like more chains in self-consistency to increase the quality of distillation data.

questions:
No
</content>
</thread_object>

</thread>
<thread>
creation date:
1730490130699

thread_type:
official_review

<thread_object>
creator: reviewer_llay

<content>
summary:
This paper proposes UniCoTT, a unified distillation framework to transfer the diverse reasoning structures with CoT to smaller language models such as BERT and RoBERTa. Firstly, UniCoT is proposed as a unified bridge of various CoT structures, which is constructed by iteratively prompting LLMs to produce explanations with correct answers. After that, a node-level supervised contrastive loss and a structural consistency loss are designed as part of the training objective. Experiments on multiple reasoning datasets verified the effectiveness of UniCoTT by surpassing the baseline methods by a large margin.

strengths:
1. The proposed method is technically sound and intuitively makes sense. It is very interesting to transfer the knowledge from structured CoT texts into smaller models that can leverage rationale knowledge in a unified manner.
2. Experimental results on several benchmarks show some improvement upon baselines and the authors conduct extensive ablation studies and analyses on various design choices.
3. The paper itself is generally well-written.

weaknesses:
1. The generalizability of KNIFE is yet to be known. The proposed framework is only verified in multiple-choice datasets. Whether it could be extended to other task settings like text generation remains a concern.
2. The process of iteratively constructing UniCoT is hard to understand from the main body of the current version. I would suggest the authors move some content from the appendix to the main body. Meanwhile, it would be helpful if the authors could provide some overall statistics on the constructed UniCoT. For example, the averaged nodes and edges of the structure.

questions:
1. The authors list "hallucinations" as one of the major drawbacks of previous works, and motivate the design of UniCoTT in ``introduction'' section. I am wondering how the designed UniCoTT framework helps to alleviate this issue.
2. In lines 385-386, why $\alpha$ and $\beta$ is set to 0.5 and 0.2 respectively? Is it an intuitive trial or a result of a grid search?
3. It would be interesting to test the annotation efficiency of CoT with the teacher model. An empirical conclusion of how many annotations are enough for great distillation performance would be insightful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730600693325

thread_type:
official_review

<thread_object>
creator: reviewer_wq9a

<content>
summary:
The paper presents a novel framework for transferring the reasoning capabilities of large language models (LLMs) to small language models (SLMs) through a structured chain-of-thought (CoT) distillation approach. The authors propose UniCoTT, which considers diverse structural CoTs (chain, tree, and graph) and employs two core strategies: iterative construction for structured CoTs and a structural constraint strategy. The framework aims to address the challenges of ensuring the rationality of generated explanations and ignoring diverse structures of CoT during knowledge transfer. The experimental results demonstrate significant performance improvements of SLMs on multiple NLP tasks across various datasets.

strengths:
1. The paper introduces a unified framework that handles diverse structural CoTs, which is a significant advancement over existing methods that focus solely on chain structures.
2. The authors provide extensive experimental evidence to support the effectiveness of UniCoTT, showing improvements across different NLP tasks and datasets.
3. The consideration of structured reasoning pathways (tree and graph) in addition to chains is a strength, as it better captures the complexity of human reasoning processes.

weaknesses:
1. The paper could benefit from a discussion on the computational complexity of UniCoTT and its scalability, especially when dealing with very large datasets or more complex reasoning tasks.
2. The construction of UniCoT relies on APIs of LLMs, which may not be accessible or feasible in all situations. The paper could address potential alternatives or mitigation strategies. Besides, SLMs usually refer to small language models, e.g., 2B and 3B. The authors mainly conducted experiments on BERT and RoBERTa, which were not convincing enough.
3. While the results are promising, the paper primarily focuses on question-answering and NLU tasks. It would be beneficial to see how UniCoTT generalizes to other types of tasks.

questions:
1. How does the performance of UniCoTT scale with the size and complexity of the knowledge to be transferred? Are there diminishing returns as the complexity increases?
2. What are the limitations of the current implementation of UniCoTT, and how might these be addressed in future work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705339066

thread_type:
official_review

<thread_object>
creator: reviewer_adyg

<content>
summary:
This paper focuses on distilling the reasoning capability, specifically chain-of-thought reasoning, from large language models into smaller models. Specifically, the paper uses prompts to guide a larger teacher model to generate multiple explanations, or "thoughts," for given questions and answers. These explanations are represented in a graph structure. Then, the small student model is trained using traditional cross-entropy loss along with a novel structural consistency loss and supervised contrastive loss proposed by the authors.

strengths:
- The writing is clear and easy to follow, with a well-defined motivation for the research.
- The distillation framework proposed is innovative, especially in using a graph structure to represent different chains of thought and introducing corresponding training methods.
- The approach is extensively tested on multiple benchmark datasets, demonstrating strong empirical performance.

weaknesses:
- The framework mainly focuses on distilling explanation and reasoning abilities into base models like BERT. A concern is the limited application scope of such encoder-based models. To further validate the effectiveness of the proposed distillation framework for reasoning abilities, it would be interesting to distill the chain-of-thought reasoning from larger models into smaller decoder-based models and test them on complex reasoning tasks.

questions:
- Why focus on using an encoder as the student model?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes UniCoTT, a unified distillation framework to transfer the diverse reasoning structures with CoT to smaller language models such as BERT and RoBERTa. Firstly, UniCoT is proposed as a unified bridge of various CoT structures, which is constructed by iteratively prompting LLMs to produce explanations with correct answers. After that, a node-level supervised contrastive loss and a structural consistency loss are designed as part of the training objective. Experiments on multiple reasoning datasets verified the effectiveness of UniCoTT by surpassing the baseline methods by a large margin.

The proposed method is technically sound and intuitively makes sense. It is very interesting to transfer the knowledge from structured CoT texts into smaller models that can leverage rationale knowledge in a unified manner. Experimental results on several benchmarks show some improvement upon baselines and the authors conduct extensive ablation studies and analyses on various design choices.

On the other hand, there has been concern on the computational complexity of UniCoTT and its scalability, as well as its generalization on non-discriminative tasks. The latter is particularly important to incorporate given the topic of the work. Through the rebuttal phase, some of the other issues were addressed, while these aforementioned issues still remain.