Sub-Task Decomposition Enables Learning in Sequence to Sequence Tasks

TL;DR

Problem: end-to-end sequence-to-sequence models struggle on multi-hop tasks. Approach: formalize intermediate sub-task supervision by appending sub-task labels to inputs and provide positive theoretical guarantees of polynomial-time learnability under decomposition. Key contributions: (i) a general learnability result for decomposition-based seq-to-seq learning over a broad class of P-time functions, (ii) a concrete bit-subset parity example showing an exponential gap between end-to-end and decomposition-based learning, and (iii) a universality claim with a circuit-based construction, plus Transformer-level empirical validation. Significance: supplies theoretical justification for chain-of-thought and decomposition strategies in NLP, linking practical intermediate supervision techniques to provable learning advantages.

Abstract

The field of Natural Language Processing has experienced a dramatic leap in capabilities with the recent introduction of huge Language Models. Despite this success, natural language problems that involve several compounded steps are still practically unlearnable, even by the largest LMs. This complies with experimental failures for end-to-end learning of composite problems that were demonstrated in a variety of domains. An effective mitigation is to introduce intermediate supervision for solving sub-tasks of the compounded problem. Recently, several works have demonstrated high gains by taking a straightforward approach for incorporating intermediate supervision in compounded natural language problems: the sequence-to-sequence LM is fed with an augmented input, in which the decomposed tasks' labels are simply concatenated to the original input. In this paper, we prove a positive learning result that motivates these recent efforts. We show that when concatenating intermediate supervision to the input and training a sequence-to-sequence model on this modified input, unlearnable composite problems can become learnable. We show that this is true for any family of tasks which on the one hand, are unlearnable, and on the other hand, can be decomposed into a polynomial number of simple sub-tasks, each of which depends only on O(1) previous sub-task results. Beyond motivating contemporary empirical efforts for incorporating intermediate supervision in sequence-to-sequence language models, our positive theoretical result is the first of its kind in the landscape of results on the benefits of intermediate supervision for neural-network learning: Until now, all theoretical results on the subject are negative, i.e., show cases where learning is impossible without intermediate supervision, while our result is positive, showing that learning is facilitated in the presence of intermediate supervision.

Figures (4)

• Figure 1: An illustrative example of the prominent method for introducing sub-task decomposition and intermediate supervision for math word problems ling-etal-2017-programcobbe2021gsm8k. Intermediate sub-tasks and their labels are concatenated to the original task’s input to form a new input sequence. At training time, the likelihood of the entire sequence following the original input is maximized conditioned on the input, and at test time only the original input is given to the model.
• Figure 2: Illustration of the proposed input and output for learning the $d=8$ bit-subset parity problem with sequence-to-sequence models.
• Figure 3: The number of steps until a BERT-base sized Transformer learns bit-subset parities with and without intermediate supervision. By learning we mean validation accuracy higher than $60\%$. While this definition is somehow arbitrary, in practice we observed a grokking phenomenon power2021grokking where very soon after the accuracy became higher than random level it also became almost perfect (accuracy $>95\%$).
• Figure 4: Illustration of a typical bit-subset parity learning trajectory. In these tasks we observed a grokking phenomenon power2021grokking where very soon after the validation accuracy became higher than random level it also became almost flawless (accuracy $>95\%$). Therefore, successful learning is not sensitive to the exact cutoff.

Theorems & Definitions (38)

• Theorem 1
• Lemma 1
• proof
• Theorem 2
• Lemma 2
• proof
• Corollary 1
• proof
• Lemma 3
• proof