From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency

TL;DR

This paper demonstrates that chain-of-thought enhances transformer sample efficiency by inducing sparse, interpretable attention patterns, which simplifies learning and improves reasoning performance beyond mere expressiveness gains.

Contribution

It reveals that CoT's benefits stem from inducing sparsity in attention, not just increased expressiveness, supported by theoretical analysis and experiments.

Findings

CoT improves sample efficiency in parity learning tasks.

CoT induces sparse, interpretable attention patterns.

Sparsity in attention layers is key to CoT's effectiveness.

Abstract

Chain-of-thought (CoT) significantly enhances the reasoning performance of large language models (LLM). While current theoretical studies often attribute this improvement to increased expressiveness and computational capacity, we argue that expressiveness is not the primary limitation in the LLM regime, as current large models will fail on simple tasks. Using a parity-learning setup, we demonstrate that CoT can substantially improve sample efficiency even when the representation power is sufficient. Specifically, with CoT, a transformer can learn the function within polynomial samples, whereas without CoT, the required sample size is exponential. Additionally, we show that CoT simplifies the learning process by introducing sparse sequential dependencies among input tokens, and leads to a sparse and interpretable attention. We validate our theoretical analysis with both synthetic and real-world experiments, confirming that sparsity in attention layers is a key factor of the improvement induced by CoT.