PaperlayerView exact revision on arXiv

Source-linked AI summary

Routing Networks: Adaptive Selection of Non-linear Functions for Multi-Task Learning

Clemens Rosenbaum, Tim Klinger, Matthew Riemer

arXiv:1711.01239v2cs.LGcs.CVcs.NE

TL;DR

Multi-task learning seeks to exploit task commonalities while limiting interference, but existing approaches often require manually designed sharing. Routing networks address this with recursively selected function blocks trained through collaborative multi-agent reinforcement learning, achieving higher accuracy, sharper convergence, and substantially lower training time than the compared baselines.

  • Problem

    Multi-task learning must leverage commonalities across tasks while minimizing interference that can reduce transfer benefits.

  • Method

    A routing network uses a router to recursively select function blocks for each input, with collaborative multi-agent reinforcement learning training the router and blocks.

  • Results

    Routing networks consistently outperform cross-stitch networks and baselines across the evaluated tasks; on CIFAR-MTL, they beat cross-stitch networks by 7%.

  • Takeaways & Limitations

    Routing networks provide dynamic task-dependent composition with nearly constant training computation as function blocks are added, unlike cross-stitch networks’ linear scaling.

  • Takeaways & Limitations

    The routing state space grows exponentially with routing depth, making it challenging to scale the approach to deeper networks in its entirety.

Abstract

from arXiv · show

Multi-task learning (MTL) with neural networks leverages commonalities in tasks to improve performance, but often suffers from task interference which reduces the benefits of transfer. To address this issue we introduce the routing network paradigm, a novel neural network and training algorithm. A routing network is a kind of self-organizing neural network consisting of two components: a router and a set of one or more function blocks. A function block may be any neural network - for example a fully-connected or a convolutional layer. Given an input the router makes a routing decision, choosing a function block to apply and passing the output back to the router recursively, terminating when a fixed recursion depth is reached. In this way the routing network dynamically composes different function blocks for each input. We employ a collaborative multi-agent reinforcement learning (MARL) approach to jointly train the router and function blocks. We evaluate our model against cross-stitch networks and shared-layer baselines on multi-task settings of the MNIST, mini-imagenet, and CIFAR-100 datasets. Our experiments demonstrate a significant improvement in accuracy, with sharper convergence. In addition, routing networks have nearly constant per-task training cost while cross-stitch networks scale linearly with the number of tasks. On CIFAR-100 (20 tasks) we obtain cross-stitch performance levels with an 85% reduction in training time.

1 INTRODUCTION

The paper introduces routing networks for multi-task learning, combining a router with composable function blocks that are selected recursively for each input. Reinforcement learning trains the non-differentiable routing decisions, and experiments report improved accuracy and faster convergence than established baselines.

  • Routing networks dynamically select and recursively compose function blocks for each input, allowing task-dependent sharing or separation.The router chooses a block, applies it, and passes the result back to the router up to a fixed depth.
  • Routers can condition decisions on activations, task labels, recursion depth, and historical decisions, while function blocks may be arbitrary composable neural-network components.Blocks can include fully connected layers, convolutional layers, or larger network parts.
  • Reinforcement learning trains the router because its sequential hard routing decisions are not differentiable.The paper models routing as an RL problem and discusses multiple training formulations.
  • The evaluation uses multi-task MNIST, Mini-imagenet, and CIFAR-100, comparing routing networks with cross-stitch and shared-layer baselines.CIFAR-100 treats its 20 label superclasses as separate tasks.
  • Routing networks achieve significant accuracy improvements, sharper convergence, and often orders-of-magnitude shorter training times than cross-stitch networks.These results are reported across the evaluated multi-task settings.

2 RELATED WORK

The related work spans manually designed sharing, selective transfer, mixtures of experts, dynamic weight generation, architecture search, conditional routing, and continual learning. The paper positions routing networks as dynamic multi-task architecture search using recursive hard decisions and multi-agent reinforcement learning.

  • Routing networks replace manually designed shared and task-specific architectures with fully dynamic compositional structures that adjust differently for each task.This targets the need to choose which parameters are shared and which remain task-specific.
  • Unlike related selective-transfer and gating methods, the paper evaluates up to 20 tasks and compares directly with cross-stitch networks.The cited related experiments used only two tasks at a time.
  • Mixtures of experts typically use differentiable soft weighting over expert outputs, whereas routing networks use hard routing decisions.The comparison distinguishes the routing mechanism rather than the broader use of specialized components.
  • Dynamic weight-generation methods can struggle to scale to deep models with many parameters, while routing networks route layers of architectures such as convnets.The paper presents layer routing as an alternative dynamic-architecture construction.
  • The work applies automated architecture-search ideas to multi-task learning, using reinforcement learning to reduce manual architecture design.The paper describes this application as novel to the authors’ knowledge.
  • Compared with single-task conditional-routing work, routing networks introduce multi-task learning, multi-agent reinforcement learning, and recursive decision processes.These are identified as the paper’s distinguishing elements.
  • Continual-learning research presents tasks sequentially over potentially long periods, providing a related setting distinct from the paper’s multi-task focus.The cited related work may also use routes but employs evolutionary algorithms instead of reinforcement learning.

3 ROUTING NETWORKS

Routing networks recursively select and compose function blocks for each input, framing routing as sequential decisions trained with reinforcement learning. The router and selected blocks are trained jointly, with task-based multi-agent variants addressing changing interactions among agents and blocks.

  • Architecture: A routing network combines a router with function blocks and recursively selects blocks to compose an output representation.The router receives the representation, task identifier, and routing depth, then selects successive blocks until the configured depth is reached.
  • Architecture: Routing networks can represent ordinary neural networks and expose N^d possible trainable paths when selecting among N blocks for depth d.Layer-specific block groups can constrain choices to reproduce conventional layered architectures.
  • Architecture: Each routing step is an RL action over function-block choices or PASS, updating the representation and increasing the depth.Non-PASS actions transform the representation; PASS preserves it while advancing the routing process.
  • Router training: The training procedure backpropagates through the selected function blocks and uses the recorded routing trace to train the router with RL.The trace contains visited states, actions, immediate rewards, and a final reward based on prediction correctness.
  • RL algorithms: Task-based routing can assign one routing agent per task, while routers may also condition decisions on activations, task labels, depth, or prior instance decisions.These choices support dynamic specialization and possible reuse of function blocks across instances.
  • RL algorithms: Joint training creates a non-stationary environment, motivating multi-agent methods such as WPL because function blocks and router policies change together.The paper reports that single-agent policy-gradient methods are less adapted to this changing environment, while WPL is designed to dampen oscillation and accelerate convergence.

4 QUANTITATIVE RESULTS

Across MNIST-MTL, MIN-MTL, and CIFAR-MTL, routing networks outperform cross-stitch and task-specific baselines, while hard routing substantially reduces training computation. Results also show that WPL and per-task agents are effective routing choices.

  • Datasets: 20 CIFAR-100 superclasses are treated as separate tasks, with 2.5k training and 500 testing instances per task.Results are averaged over 3 runs.
  • RL algorithm comparison: WPL outperforms tabular Q-Learning by about 4% on CIFAR-MTL, while multiple agents outperform a single agent.Tabular versions outperform approximation versions in this experiment.
  • Routing architecture comparison: The fully routed WPL architecture outperforms alternative routing architectures, including soft-mixture, single-agent, and dispatched variants.The best performer is routing-all-fc, which routes all fully connected layers.
  • Cross-domain comparison: 7% on CIFAR-MTL and about 2% on MIN-MTL are the routing network’s margins over cross-stitch networks.The corresponding margins over the nearest task-specific-1-fc baseline are 11% and about 6%, respectively.
  • Cross-domain comparison: 9% on MNIST-MTL is the routing network’s margin over cross-stitch networks, while its margin over task-specific-all-fc is 3%.Cross-stitch networks and the soft version had difficulty training on this dataset.
  • Training cost: 85% is the reduction in CIFAR-MTL training time, from roughly 38 hours to 5.6 hours.Routing computation remains effectively unchanged as function blocks are added, whereas cross-stitch computation scales linearly with their number.

5 QUALITATIVE RESULTS

The qualitative analyses show routing agents converging from exploratory policies toward stable task-specific choices, while learned routing structures vary across depth rather than following a fixed shared-then-separated pattern.

  • Policy dynamics: Routing policies begin with high entropy, then converge toward nearly 100% probability for one function block per task.
  • Policy dynamics: Block 7 selection oscillates before the pink and green agents emerge as its high-probability users for their respective tasks.
  • Routing structure: The MNIST-MTL routing map represents tasks across the top and routing decisions across successive layers.
  • Routing structure: Instead of sharing early and diverging later, the learned MNIST-MTL structure uses 7 blocks, compresses to 4, then expands to 5.

6 FUTURE WORK

The paper presents routing networks and multi-agent training as broadly applicable, while identifying depth-related scaling and mixed online few-shot results as directions requiring further work.

  • Future directions: The proposed architecture and training algorithm outperform cross-stitch networks, baselines, and other single-agent approaches.
  • Future directions: Routing networks can be applied to state-of-the-art models so they learn to dynamically adjust their representations.
  • Limitations: The routing state space grows exponentially with routing depth, making scaling to entirely deeper networks challenging.
  • Future directions: The most successful experiments used one agent per task with tabular Weighted Policy Learner training, while neural approximators remain under investigation.
  • Future directions: Online few-shot experiments added agents for new tasks, but the results were mixed and the setting remains planned future work.

7 APPENDIX

The appendix documents routing variants, policy representations, reinforcement-learning choices, and figures examining collaboration, computation, task scaling, and recurrent application of blocks.

  • Additional analyses: The appendix includes figures on collaboration reward, normalized per-task training cost, and results for the first n CIFAR-MTL tasks.
  • Recurrent routing: Compatible function blocks can be applied repeatedly when their input and output dimensions match, making the routing network fully recurrent in the described case.
  • Implementation variants: The experiments tested 9 routing-architecture implementation variants summarized in Appendix Tables 3 and 4.
  • Implementation variants: Router implementations vary in agent count, including one agent per task, a single agent, and #tasks + 1 agents.
  • Policy representations: Policies are represented either as tables indexed by routing layers and actions or as layer-specific MLP approximators with hidden dimension 64.
  • Policy representations: Policy inputs may use task and depth, with approximation variants additionally incorporating the representation vector.
  • Training algorithm: Weighted Policy Learner dampens policy oscillation by reducing an agent’s learning rate after policy-gradient changes.
Loading 1711.01239v2…