Chain of Thought in Order: Discovering Learning-Friendly Orders for Arithmetic

Yuta Sato, Kazuhiko Kawamoto, Hiroshi Kera

Keywords: chain of thought, learning-friendly order, transformer

Abstract:

The chain of thought, i.e., step-by-step reasoning, is one of the fundamental mechanisms of Transformers. While the design of intermediate reasoning steps has been extensively studied and shown to critically influence performance, the ordering of these steps has received little attention, despite its significant effect on the difficulty of reasoning.
This study addresses a novel task of unraveling the chain of thought—reordering decoder input tokens into a learning-friendly sequence for Transformers, for learning arithmetic tasks. The proposed pipeline first trains a Transformer on a mixture of target sequences arranged in different orders and then identifies benign orders as those with fast loss drops in the early stage.
As the search space grows factorially in sequence length, we propose a two-stage hierarchical approach for inter- and intra-block reordering.
Experiments on four order-sensitive arithmetic tasks show that our method identifies a learning-friendly order out of a few billion candidates. Notably, on the multiplication task, it recovered the reverse-digit order reported in prior studies.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 3731