Generating Auxiliary Tasks with Reinforcement Learning
Judah Goldfeder, Matthew So, Hod Lipson
TL;DR
This paper tackles the cost and complexity of designing auxiliary tasks for Auxiliary Learning by proposing RL-AUX, a reinforcement-learning framework that dynamically assigns auxiliary labels per training example and can optionally learn per-sample weights. The main network training is embedded in an RL environment where the agent’s actions specify auxiliary labels and rewards reflect improvements on the primary task, allowing avoidance of bi-level optimization. Experiments on CIFAR-100 with a 20-superclass structure show RL-AUX matching or slightly surpassing MAXL while beating human-labeled auxiliaries, with weight-aware extensions (WA-MAXL, WA-RL) delivering further gains on CIFAR-100 and SVHN. The results indicate reinforcement learning is a viable, scalable path for automatic auxiliary task generation in image classification and can be combined with sample-level weighting to further enhance performance.
Abstract
Auxiliary Learning (AL) is a form of multi-task learning in which a model trains on auxiliary tasks to boost performance on a primary objective. While AL has improved generalization across domains such as navigation, image classification, and NLP, it often depends on human-labeled auxiliary tasks that are costly to design and require domain expertise. Meta-learning approaches mitigate this by learning to generate auxiliary tasks, but typically rely on gradient based bi-level optimization, adding substantial computational and implementation overhead. We propose RL-AUX, a reinforcement-learning (RL) framework that dynamically creates auxiliary tasks by assigning auxiliary labels to each training example, rewarding the agent whenever its selections improve the performance on the primary task. We also explore learning per-example weights for the auxiliary loss. On CIFAR-100 grouped into 20 superclasses, our RL method outperforms human-labeled auxiliary tasks and matches the performance of a prominent bi-level optimization baseline. We present similarly strong results on other classification datasets. These results suggest RL is a viable path to generating effective auxiliary tasks.