Source-linked AI summary
Sub-Task Decomposition Enables Learning in Sequence to Sequence Tasks
Noam Wies, Yoav Levine, Amnon Shashua
TL;DR
Multi-step natural-language problems remain difficult for large language models, and empirical work has used concatenated sub-task labels as intermediate supervision. This paper analyzes that sequence-to-sequence approach and proves that sufficient supervision makes some otherwise unlearnable tasks efficiently learnable, with guarantees for efficiently decomposable tasks.
Problem
Multi-hop problems requiring compounded operations remain difficult for language models, while empirical intermediate-supervision approaches lack theoretical guarantees motivating their benefits.
Source: §1 · pp. 1–2Method
Results
With sufficient intermediate supervision, sequence-to-sequence models can efficiently learn tasks that are unlearnable without it, including efficiently decomposable tasks with polynomially many sub-tasks.
Source: §1 · pp. 2–3Takeaways & Limitations
The result provides a positive theoretical guarantee for the intermediate-supervision approach used in sequence-to-sequence language-model work.
Source: §1 · p. 2, §7 · p. 9For bit-subset parity, the demonstrated exponential advantage uses O(d) supervised sub-tasks; whether O(1) sub-tasks suffice remains open.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Multi-hop tasks remain difficult for large language models, while concatenating sub-task labels to the input has shown empirical gains. This paper gives positive theoretical guarantees that sequence-to-sequence models can learn efficiently decomposable tasks with intermediate supervision, including tasks that are unlearnable without it.
1 INTRODUCTION
- 61% accuracy on StrategyQA is reported for Gopher, compared with around 87% human-level performance on this multi-hop benchmark.The example illustrates the difficulty of tasks requiring implicit reasoning steps.Source: §1 · p. 1
- Concatenating intermediate task labels to the input has been used in natural-language tasks, with reported gains over training without those labels.For 8-digit addition, Nye et al. (2022) report > 99% accuracy with intermediate calculations versus around ∼35% without them.Source: §1 · p. 2
- With intermediate supervision, a sequence-to-sequence model can learn efficiently decomposable tasks with polynomially many sub-tasks, each depending on only a small fraction of the input.The paper contrasts this with the existence of efficiently decomposable tasks that no polynomial-time algorithm can learn without intermediate supervision.Source: §1 · p. 2
- For bit-subset parity, the paper proves efficient learning with intermediate supervision and an exponential gap from end-to-end neural-network learnability.The target is the parity of an unknown subset of d/2 input bits.Source: §1 · pp. 2–3
- The paper generalizes its positive result to polynomial-time functions and identifies tasks that are unlearnable in polynomial time without supervision but efficiently learnable with it.The stated hardness follows under standard cryptographic or computational-complexity assumptions.Source: §1 · p. 3
- The paper’s positive guarantee is the first theoretical result showing neural-network learnability facilitated by intermediate supervision for a problem unlearnable without it.The authors also state that the sequence-to-sequence setting can learn any function in P with sufficient intermediate supervision.Source: §1 · p. 3
2 RELATED WORK
Prior work has explored intermediate supervision in reinforcement learning and demonstrated theoretical limits of end-to-end learning. The paper distinguishes its contribution by proving that supervision can make learning feasible, not only that learning without decomposition is difficult.
2 RELATED WORK
- Earlier theoretical work found end-to-end multi-hop learning inefficient or unfeasible, but did not establish polynomial-time learning made possible by intermediate supervision.Shalev-Shwartz et al. (2017) found higher gradient signal-to-noise with supervision, without a polynomial-update guarantee; Shalev-Shwartz & Shashua (2016) studied verification sample complexity.Source: §2 · pp. 3–4
- The paper presents its result as the first theoretical demonstration that neural networks benefit from sub-task decomposition in a supervised-learning setting.Prior reinforcement-learning analyses include a linear reduction in complexity from sub-goal rewards, unlike the paper’s exponential gap.Source: §2 · pp. 3–4
3 THE ANALYZED SEQUENCE-TO-SEQUENCE LEARNING ALGORITHM
The analyzed sequence-to-sequence procedure concatenates intermediate labels after the original input and predicts the final answer autoregressively. Training uses ground-truth prior sub-task results as inputs, whereas testing feeds the model’s own predictions.
3 THE ANALYZED SEQUENCE-TO-SEQUENCE LEARNING ALGORITHM
- The combined sequence contains the original input followed by intermediate labels, and the model predicts the sequence autoregressively.During training, previous sub-task results are ground truth; at test time, they are the model’s earlier predictions.Source: §3 · p. 4
- The analyzed sequence-to-sequence model is a classical Elman recurrent neural network with ReLU activations.Its input sequence has length T = d + len(S).Source: §3 · p. 4
- The binary cross-entropy loss excludes predictions of the original input and scores intermediate labels and the final outcome.The analyzed SGD variant trains only hidden weights W while keeping A, B, and M0 frozen at initialization.Source: §3 · p. 4
4 COMPOUNDED SEQUENCE TO SEQUENCE LEARNABILITY
The section establishes that teacher-forced sequence-to-sequence learning can generalize with polynomial sample complexity and gradient updates for a suitable hypothesis class. A lemma connects low error on supervised sub-tasks to test-time accuracy despite predicted intermediate results.
4 COMPOUNDED SEQUENCE TO SEQUENCE LEARNABILITY
- Teacher forcing yields generalization with polynomial sample complexity and gradient updates when the task’s sub-tasks belong to the analyzed hypothesis class.The authors use this guarantee to establish their positive results for later task families.Source: §4 · pp. 4, 6
- Low zero-one loss on every supervised sub-task bounds the chance of error in any sub-task, linking training with ground-truth prior results to test-time predictions.The proof uses a union bound over output locations.Source: §4 · p. 5
- The test-time error increase from not having ground-truth previous sub-task results is at most a factor of T − d, which is modest when T is polynomial in d.This qualifies the transfer from teacher-forced training to iterative prediction.Source: §4 · p. 5
- The analyzed hypothesis class represents each output as a sign of composed linear functions over at most N prior input locations, with polynomial activations.Its complexity is polynomial in T and the largest coefficient magnitude, but exponential in activation degree and N.Source: §4 · p. 5
- The convergence and generalization argument combines the teacher-forcing lemma with Wang et al. (2021)’s RNN learning result for low-complexity final answers.The paper treats intermediate steps as supervision, unlike that earlier work’s sequence-output setup.Source: §4 · p. 5
5 LEARNING BIT-SUBSET PARITY WITH SEQUENCE TO SEQUENCE MODELS
Bit-subset parity, difficult to learn end to end, becomes efficiently learnable when its sub-task labels are supplied. A BERT-base-sized Transformer experiment also shows a large gap between supervised and unsupervised settings.
5 LEARNING BIT-SUBSET PARITY WITH SEQUENCE TO SEQUENCE MODELS
- The decomposition recursively computes parities of adjacent groups of input bits, forming a binary tree whose nodes represent parity sub-tasks.The tree’s intermediate labels are converted into a sequence and concatenated after the original input.Source: §5 · p. 6
5.1 LEARNABILITY OF BIT-SUBSET PARITY WITH AND WITHOUT INTERMEDIATE SUPERVISION
- With sub-task supervision, a sequence-to-sequence model learns bit-subset parity to arbitrarily low zero-one loss in polynomially many updates; without it, learning requires exponentially many updates.The negative result is for full gradient descent; the cited work notes it does not hold for stochastic gradient descent with complex non-random initialization.Source: §5.1 · p. 7, §5.2 · p. 7
5.2 BIT-SUBSET PARITY EXPERIMENTS
- For 32-bit subset parity, a BERT-base-sized Transformer learned with intermediate supervision but did not learn without it after more than 2M steps.The experiment varied input sizes from 8 to 256 bits and trained models for 100k iterations, with exceptions for some unsupervised settings.Source: §5.2 · pp. 7–8
6 UNIVERSALITY OF DECOMPOSITION BASED SEQUENCE-TO-SEQUENCE LEARNING
The paper extends its positive learning result to functions in P by supplying intermediate labels for circuit sub-tasks. Under standard cryptographic or computational-hardness assumptions, some such functions cannot be learned in polynomial time without that supervision.
6 UNIVERSALITY OF DECOMPOSITION BASED SEQUENCE-TO-SEQUENCE LEARNING
- The construction represents a polynomial-size circuit as a directed acyclic graph with O(1) input-degree gates, then assigns intermediate tasks to non-leaf vertices in topological order.The resulting intermediate-label sequence has polynomial length in d.Source: §6 · p. 8
- Each intermediate target can be represented by a low-degree polynomial of O(1) inputs, with bounded coefficients that do not grow with d.This lets the supervised targets fall within the function class used to establish the polynomial learning guarantee.Source: §6 · pp. 8–9
- With sub-task decomposition supervision, a sequence-to-sequence model can learn any function in P with polynomial network size, sample complexity, and gradient updates.Under standard cryptographic or computational-complexity hardness assumptions, some binary classification problems cannot be learned in polynomial time without sub-task supervision.Source: §6 · pp. 8–9
7 DISCUSSION
The discussion presents an exponential learning gap for a broad family of functions and connects the theory to sequence-to-sequence methods that concatenate intermediate supervision to the input. It also identifies sufficient supervision and practical-architecture understanding as open boundaries.
7 DISCUSSION
- The paper reports an exponential learning gap between algorithms with and without intermediate supervision for a broad family of functions.It frames the analysis around concatenating intermediate labels to the input and jointly training a sequence-to-sequence model on those labels and the overall output.Source: §7 · p. 9
- For bit-subset parity, the demonstrated exponential advantage uses O(d) supervised sub-tasks; whether a similar advantage holds with only O(1) remains open.The discussion also leaves open which sub-tasks practical architectures can solve end-to-end and which require decomposition.Source: §7 · p. 9
A COMPOUNDED SEQUENCE TO SEQUENCE LEARNABILITY DETAILS
- The framework’s test-time error remains low when each single-hop sub-task has low zero-one loss, by a union-bound argument that accounts for accumulated errors.At test time, the model uses its previous predictions rather than ground-truth sub-task results.Source: A COMPOUNDED SEQUENCE TO SEQUENCE… · p. 14
B SUB TASKS LEARNABILITY PROOFS
The proofs establish that the parity and polynomial-time task decompositions fit the paper’s sequence-to-sequence hypothesis class. For parity, the decomposition is encoded as a sequence of intermediate nodes, with predicted sub-task results used at test time.
B SUB TASKS LEARNABILITY PROOFS
- The parity decomposition is converted to a sequence by inverse breadth-first traversal, with the input bits and decomposition nodes (except the root) forming the input sequence.The root represents the final answer.Source: B SUB TASKS LEARNABILITY PROOFS · p. 14
- At training time, each sub-task receives the preceding ground-truth result, while at test time it receives the model’s preceding prediction.The model is causal, so a prediction at one position does not depend on later sub-task results.Source: B SUB TASKS LEARNABILITY PROOFS · p. 14
- The proof shows that length-2 parity intermediate steps and intermediate steps for functions in P belong to the defined hypothesis class.The section states that these lemmas establish the intermediate steps are simple enough for Theorem 2; it also derives the required polynomial using Lagrange basis functions.Source: B SUB TASKS LEARNABILITY PROOFS · pp. 14–15
C EXTENSION FOR SGD WITH FINITE PRECISION
This section extends the theorem’s proof to finite-precision stochastic gradient descent. It retains the initialization-based complexity argument and adapts the generalization analysis to the finite-precision algorithm.
C EXTENSION FOR SGD WITH FINITE PRECISION
- The proof extends Theorem 2 to Algorithm 2, a finite-precision variant of stochastic gradient descent.It modifies the analysis of the full-precision algorithm to handle finite-precision gradients.Source: C EXTENSION FOR SGD WITH FINITE… · p. 16
- The proof retains the learned-hypothesis complexity argument because it is independent of the gradients, and focuses the adaptation on a generalization bound.The remaining generalization argument from Wang et al. (2021) is stated to remain unchanged.Source: C EXTENSION FOR SGD WITH FINITE… · p. 16
- The analysis bounds binary cross-entropy loss by C · T^14 for some constant C > 0.This bound is used in the subsequent martingale-based generalization argument.Source: C EXTENSION FOR SGD WITH FINITE… · p. 18
D EXTENSION FOR GD WITH FINITE PRECISION
The paper extends its finite-precision learning result to gradient descent and supplies supporting RNN learnability bounds. The discussion of those bounds identifies polynomially bounded gradients as consistent with the negative-results setting, attributing efficient learning there to intermediate supervision.
D EXTENSION FOR GD WITH FINITE PRECISION
- Theorem 2 still holds with gradient-descent-based Algorithm 3 rather than stochastic-gradient-descent-based Algorithm 2, establishing efficient parity learning with sub-task decomposition supervision in the same setting where learning is impossible without intermediate supervision.The proof follows section C while accounting for the full non-stochastic gradients.Source: D EXTENSION FOR GD WITH FINITE PRECISION · p. 19
E LEARNABILITY OF RNNS
- The RNN used for the positive results satisfies the polynomially bounded-gradients requirement from the parity negative results, supporting the paper’s identification of intermediate supervision as the enabling factor.The passage states this implication holds with high probability.Source: E LEARNABILITY OF RNNS · p. 23
- The supporting RNN analysis bounds the finite-width deviation from its infinite-width neural tangent kernel.The proof relies on a result from Wang et al. (2021).Source: E LEARNABILITY OF RNNS · p. 23
F BIT-SUBSET PARITY END-TO-END NEGATIVE RESULT PROOFS
The section proves an end-to-end negative result for bit-subset parity: polynomially many gradient updates do not achieve non-negligible accuracy under the stated conditions. The proof combines exponentially low gradient variance with high random-guessing error.
F BIT-SUBSET PARITY END-TO-END NEGATIVE RESULT PROOFS
- The negative result uses the same setup as the positive result, including polynomially bounded gradients and finite-precision gradient descent.The text says the positive and negative proofs therefore use the exact same setup.Source: F BIT-SUBSET PARITY END-TO-END… · p. 24
- The proof establishes exponentially low gradient variance by showing that distinct parity functions are uncorrelated.The argument uses the parity hypothesis class and a gradient-variance bound for uncorrelated functions.Source: F BIT-SUBSET PARITY END-TO-END… · p. 25
- Low gradient variance alone is insufficient; the proof also establishes non-vanishing random-guessing error, Θ(1) in dimension d.The error is linear in the fraction of the hypothesis class from which guessing occurs.Source: F BIT-SUBSET PARITY END-TO-END… · pp. 24, 26
G BIT-SUBSET PARITY EXPERIMENTS FULL DETAILS
The experiments compare Transformer learning of bit-subset parity with and without intermediate supervision. Without supervision, models do not learn tasks above 18 bits; with it, accuracy is almost perfect even at 128 bits.
G BIT-SUBSET PARITY EXPERIMENTS FULL DETAILS
- The experiments use a BERT-Base-size Transformer decoder trained on randomly sampled, non-overlapping training, validation, and test sets.The implementation uses the standard GPT-2 code and a vocabulary size of 2.Source: G BIT-SUBSET PARITY EXPERIMENTS FULL… · pp. 27–28
- With intermediate supervision, test accuracy is almost perfect even at 128 bits, while unsupervised models cannot learn beyond random guessing above 18 bits.Results average three dataset seeds; the best hyperparameters are selected by validation binary cross-entropy loss.Source: G BIT-SUBSET PARITY EXPERIMENTS FULL… · p. 28
- Table 1 reports the first iteration above the random-accuracy threshold and test accuracy after 100k iterations; grokking makes the exact cutoff relatively insensitive.The paper reports validation accuracy becoming almost flawless (>95%) soon after it rises above random level.Source: G BIT-SUBSET PARITY EXPERIMENTS FULL… · p. 28, Figure 4 · p. 28
Related papers
This paper cites
- Attention Is All You Need2017
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding2018
- Language Models are Few-Shot Learners2020
- Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer2019
- Neural Tangent Kernel: Convergence and Generalization in Neural Networks2018
- Measuring and Narrowing the Compositionality Gap in Language Models2022