CDW-CoT: Clustered Distance-Weighted Chain-of-Thoughts Reasoning

TL;DR

This work addresses the problem that uniform Chain-of-Thought prompts fail to capture the diversity within datasets. It introduces CDW-CoT, which clusters data into $K$ groups, learns an optimal prompt distribution per cluster, and computes a distance-weighted prompt distribution at test time to tailor reasoning for each instance. Empirical results across six datasets and two LLMs show substantial gains over Zero-Shot-CoT, Auto-CoT, and Manual-CoT, with improvements such as $25.34\%$ on LLaMA2 and $15.72\%$ on LLaMA3 on average. The approach reduces manual prompt engineering and demonstrates robust performance across commonsense, symbolic, and mathematical reasoning tasks, with potential for broader applicability and multimodal extensions.

Abstract

Large Language Models (LLMs) have recently achieved impressive results in complex reasoning tasks through Chain of Thought (CoT) prompting. However, most existing CoT methods rely on using the same prompts, whether manually designed or automatically generated, to handle the entire dataset. This one-size-fits-all approach may fail to meet the specific needs arising from the diversities within a single dataset. To solve this problem, we propose the Clustered Distance-Weighted Chain of Thought (CDW-CoT) method, which dynamically constructs prompts tailored to the characteristics of each data instance by integrating clustering and prompt optimization techniques. Our method employs clustering algorithms to categorize the dataset into distinct groups, from which a candidate pool of prompts is selected to reflect the inherent diversity within the dataset. For each cluster, CDW-CoT trains the optimal prompt probability distribution tailored to their specific characteristics. Finally, it dynamically constructs a unique prompt probability distribution for each test instance, based on its proximity to cluster centers, from which prompts are selected for reasoning. CDW-CoT consistently outperforms traditional CoT methods across six datasets, including commonsense, symbolic, and mathematical reasoning tasks. Specifically, when compared to manual CoT, CDW-CoT achieves an average accuracy improvement of 25.34% on LLaMA2 (13B) and 15.72% on LLaMA3 (8B).

Figures (4)

• Figure 1: Using the same prompts for all instances in the dataset resulted in significant performance variability across different clusters, highlighting the limitations of Auto-CoT in addressing diverse reasoning demands within different data categories. This underscores the need for tailored prompt strategies.
• Figure 2: Framework of the proposed CDW-CoT. (a) After clustering, prompt candidates are generated based on the cluster centers. $\text{CTD}^{(i)}$ and cluster center coordinates $(x_i, y_i)$ are also obtained. (b) For each cluster, $p^{(i)}$ is initially set to $p^{(\text{init})}$ and then optimized through Black-Box Prompt Learning (BBPL) to achieve the optimal distribution. (c) For test instance, a distance-weighted prompt probability distribution is constructed to select prompts and perform reasoning.
• Figure 3: Sensitivity analysis of temperature for CDW-CoT on StrategyQA and MultiArith datasets.
• Figure 4: Impact of pool size on CDW-CoT on CommonsenseQA and MultiArith datasets.