Source-linked AI summary
Adaptive Scheduling for Multi-Task Learning
Sébastien Jean, Orhan Firat, Melvin Johnson
TL;DR
Multilingual NMT fixed schedules offer limited control over performance trade-offs across tasks. The paper studies adaptive schedules that reweight tasks through sampling, gradient scaling, or learning-rate scaling based on validation performance relative to baselines. These approaches improve low-resource performance while preserving high-resource performance, though validation-driven schedules can oscillate when scores change faster than they are computed.
Problem
Fixed task-mixing schedules offer little control over performance trade-offs between low-resource and high-resource language pairs.
Method
The paper evaluates explicit schedules that alter sampling probabilities and implicit schedules that scale task gradients or learning rates using validation performance.
Results
Adaptive approaches achieve satisfactory performance on both tasks, improving low-resource En-De while keeping high-resource En-Fr comparable to its baseline.
Takeaways & Limitations
Adaptive scheduling is promising for multilingual models with many tasks, where searching over fixed schedules may become cumbersome or prohibitively expensive.
Takeaways & Limitations
Adaptive schedules may behave poorly when validation performance changes faster than it is recomputed, causing task probabilities to oscillate and one task to be catastrophically forgotten.
Abstract
from arXiv · showhide
To train neural machine translation models simultaneously on multiple tasks (languages), it is common to sample each task uniformly or in proportion to dataset sizes. As these methods offer little control over performance trade-offs, we explore different task scheduling approaches. We first consider existing non-adaptive techniques, then move on to adaptive schedules that over-sample tasks with poorer results compared to their respective baseline. As explicit schedules can be inefficient, especially if one task is highly over-sampled, we also consider implicit schedules, learning to scale learning rates or gradients of individual tasks instead. These techniques allow training multilingual models that perform better for low-resource language pairs (tasks with small amount of data), while minimizing negative effects on high-resource tasks.
1 Introduction
Multi-task learning can leverage additional data and reduce memory use, but limited capacity and poor scheduling can cause under-performance, over-training, or forgetting. In multilingual NMT, adaptive schedules target these trade-offs by reallocating task importance according to validation performance.
- Multi-task training can exploit multiple data sources while using less memory than separate models.
- Limited model capacity can make multi-task models under-perform, while simpler tasks may be over-trained or forgotten under poor schedules.
- Multilingual NMT commonly uses fixed uniform or dataset-size-based sampling, which offers little control over individual task performance.
- Adaptive schedules improve low-resource accuracy while keeping high-resource accuracy intact within the same multi-task model.
2 Explicit schedules
Predefined task schedules provide limited control over performance, so explicit adaptive scheduling adjusts sampling probabilities using validation performance relative to task baselines.
- Fixed schedules sample tasks uniformly, by dataset size, or with predefined time-varying preferences, limiting control over all task performances.
- The method assumes each task’s individual baseline validation performance is known and continually records multi-task validation scores.
- Tasks performing poorly relative to their baselines are over-sampled, while well-performing tasks are sampled less often.The aggressiveness parameter α controls oversampling, and ε prevents numerical errors while smoothing the distribution.
3 Implicit schedules
Implicit schedules avoid directly changing sampling frequencies by scaling task gradients or learning rates, while using validation performance to favor underperforming tasks and protect against forgetting. Their behavior depends on optimizer design, especially with Adam.
- Implicit scheduling achieves effects similar to explicit schedules by manipulating task gradients or learning rates instead of sampling probabilities.This can be useful when explicit schedules over-sample one task or when many tasks make direct scheduling restrictive.
- Validation performance guides gradient scaling because training loss is not always a good proxy for validation or test performance.
- Task weights begin near uniform, favor tasks with low relative validation performance, and include a minimum weight intended to avoid catastrophic forgetting.
- Scaling gradients or learning rates is equivalent under standard SGD but not under Adam, whose adaptive updates make optimizer state and accumulator sharing consequential.Separate optimizers with scaled gradients are ineffective, while shared optimizers with scaled learning rates blur tasks through momentum.
- Baseline scores can come from trained single-task models or from an expected performance target for the multi-task model.
4 Experiments
Experiments compare fixed and adaptive schedules on highly imbalanced En-Fr and En-De translation tasks. Adaptive methods improve low-resource performance while generally avoiding the high-resource degradation caused by uniform sampling, though tuned constant schedules perform slightly better overall.
- 4.1 Data: The experiments restrict En-De to 1 million sentences versus almost 40 million for full En-Fr, creating a pronounced resource imbalance.Both tasks use WMT’14 data and a joint 32K-subword vocabulary.
- 4.2 Models: Models use Transformer-base configurations with shared encoders after fully shared models performed multiple BLEU points below the baselines.
- 4.3 Results: Uniform sampling improves En-De by more than 1 BLEU point but significantly degrades En-Fr, whereas 75% En-Fr sampling makes En-Fr comparable to its baseline.
- 4.3 Results: Explicit adaptive scheduling behaves similarly on En-De and somewhat trails the En-Fr baseline.
- 4.3 Results: Implicit validation-based scheduling improves over the En-De baseline less than other approaches while matching the En-Fr baseline and fixed 75% En-Fr sampling.
- 4.3 Results: Adaptive approaches provide satisfactory performance on both tasks, while hyper-parameter search over constant schedules achieves slightly better results overall.The comparison evaluates multilingual scheduling methods against single-task baselines on En-Fr and En-De.
5 Discussion and other related work
The discussion places adaptive scheduling alongside dynamic weighting, Pareto-based optimization, and learning-rate methods. It also identifies a stability boundary: rapidly changing validation performance can make adaptive weights oscillate and cause catastrophic forgetting.
- 5 Discussion and other related work: Dynamic weight averaging upweights straggling tasks based on recent loss progress, contrasting with curriculum learning, which prefers tasks progressing faster.
- 5 Discussion and other related work: Pareto-based multi-task optimization targets Pareto-optimal training losses but does not guarantee Pareto optimality for the evaluation metric.
- 5 Discussion and other related work: Implicit scheduling remains constrained by the underlying predefined learning-rate schedule, although hypergradient descent suggests a possible multi-task extension.
- 5 Discussion and other related work: Adaptive schedules may oscillate when validation performance changes faster than it is measured, causing one task to be favored while the other is catastrophically forgotten.
6 Conclusion
The paper concludes that validation-controlled adaptive schedules can improve low-resource translation while maintaining comparable high-resource performance. It identifies scaling to many tasks as an important direction for increasing their utility.
- 6 Conclusion: Adaptive schedules control task weights using validation BLEU, either by changing sampling directly or by adjusting optimization.
- 6 Conclusion: Compared with single-task baselines, adaptive scheduling improves low-resource En-De performance while remaining comparable on high-resource En-Fr.
- 6 Conclusion: Applying adaptive schedulers to many simultaneous tasks could increase their utility, where searching over fixed schedules becomes cumbersome.
A Impact of hyper-parameters
Hyper-parameter experiments compare explicit constant and dev-based schedules with implicit GradNorm and dev-based schedules. For explicit constant scheduling, favoring the high-resource task produced more balanced results than uniform sampling.
- Overview: The appendix examines hyper-parameter effects for explicit constant, explicit dev-based, implicit GradNorm, and implicit dev-based schedules.Figures 1–4 correspond to these schedule families.
- Explicit constant schedules: A 75%–25% split favoring the high-resource task produced more balanced BLEU results than uniform 50%–50% sampling.The two tasks had a 1/40 dataset-size ratio, and oversampling the high-resource task improved overall performance for both tasks.
- Explicit constant schedules: Figure 1 compares BLEU scores across explicit constant sampling ratios from 10% to 90%, with darker colors indicating higher scores.Development and test BLEU were evaluated using each fixed ratio throughout training.
- Explicit dev-based schedules: Explicit dev-based schedules vary α and ϵ to control oversampling and forgetting.The supplied passages identify the explored hyper-parameters but do not report their comparative outcomes.
B Implicit validation-based scheduling progress
Implicit validation-based scheduling adapts task weights and learning rates during training using validation performance. The procedure tracks relative scores and validation BLEU against task-specific baselines.
- Implicit validation-based scheduling progress: The implicit schedule uses α = 16, β = 0.1, γ = 0.05, baselines of 24 and 35 for En-De and En-Fr, and an inverse-square-root learning-rate schedule with learning rate 1.5 and 40K warm-up steps.These settings describe the reported best-performing model configuration.
- Implicit validation-based scheduling progress: Task weights are changed during training, and predicted weights adjust each task’s learning rate.The schedule also computes relative task scores and tracks validation BLEU against corresponding baseline values.
- Implicit validation-based scheduling progress: Figure 5 shows task weights, predicted learning-rate adjustments, relative scores, and validation BLEU progression against baselines.The four views occupy the figure’s top-left, top-right, bottom-left, and bottom-right panels, respectively.
C Possible training instabilities
An appendix experiment exhibits wild oscillations in BLEU scores and task probabilities under an explicit schedule. The instability involves alternating task preference and catastrophic forgetting.
- Possible training instabilities: The failed experiment used an explicit schedule with tied encoder parameters, the first four decoder layers, and the softmax.The appendix presents this as a setting exhibiting wildly varying oscillations.