summaries: The paper proposes an original idea to use distillation to speed up modern distributed RL settings, when data collection is done on CPUs with the learning happening on accelerated hardware, e.g. GPU. More specifically, the authors propose to use a transformer for the learner and distil the policy into LSTM actors. With this, they achieve much faster wallclock time compared to the Transformer for Actors setup, however, losing in sample-efficiency.

summaries: The paper proposes a method for "actor-latency constrained" settings: Recently, transformers have been shown to be powerful models in RL which, in particular, exhibited better sample complexity in settings in which long-term credit assignment in partial observability was required (e.g. the T-maze). However, they are computationally expensive. Consequently, the authors propose to train transformers on the learner, supported by hardware acceleration, but also train a smaller LSTM agent which can be efficiently executed on the actors.

summaries: The paper proposes a solution to actor-latency constrained settings in RL by using policy distillation to compress a large “learner model” towards a more tractable “actor model”. In particular, it proposes to exploit the superior sample efficiency of transformer models while utilising an LSTM-based actor during execution. The proposed procedure, called Actor-Learner Distillation (ALD), provides comparable performance to transformers in terms of sample efficiency, yet produces a wall-clock run-time that's on par with LSTM agents.

abstractText: Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on unaccelerated hardware such as CPUs, which likewise restricts model size to prevent intractable experiment run times. These “actor-latency” constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during acting, we develop an “Actor-Learner Distillation” (ALD) procedure that leverages a continual form of distillation that transfers learning progress from a large capacity learner model to a small capacity actor model. As a case study, we develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, at the cost of significantly higher computational complexity. With transformer models as the learner and LSTMs as the actor, we demonstrate in several challenging memory environments that using Actor-Learner Distillation recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.

authors:
Abbas Abdolmaleki
Jost Tobias Springenberg
Yuval Tassa
Remi Munos
Nicolas Heess
Martin Riedmiller

title: Maximum a posteriori policy optimisation

venue: In International Conference on Learning Representations,

authors:
Samira Abnar
Mostafa Dehghani
Willem Zuidema

title: Transferring inductive biases through knowledge distillation

venue: arXiv preprint arXiv:2006.00555,

authors:
Jimmy Ba
Rich Caruana

title: Do deep nets really need to be deep? In Advances in neural information processing

authors:
Glen Berseth
Cheng Xie
Paul Cernek
Michiel Van de Panne

title: Progressive reinforcement learning with distillation for multi-skilled motion control

authors:
Tom B Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell

title: Language models are few-shot learners

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

venue: arXiv preprint arXiv:1904.10509,

authors:
Karl Cobbe
Jacob Hilton
Oleg Klimov
John Schulman

title: Phasic policy gradient

venue: arXiv preprint arXiv:2009.04416,

authors:
Wojciech Czarnecki
Siddhant Jayakumar
Max Jaderberg
Leonard Hasenclever
Yee Whye Teh
Nicolas Heess
Simon Osindero
Razvan Pascanu

title: Mix & match agent curricula for reinforcement learning

authors:
Zihang Dai
Guokun Lai
Yiming Yang
Quoc V Le

title: Funnel-transformer: Filtering out sequential redundancy for efficient language processing

venue: arXiv preprint arXiv:2006.03236,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova. Bert

title: Pre-training of deep bidirectional transformers for language understanding

authors:
Lasse Espeholt
Hubert Soyer
Rémi Munos
Karen Simonyan
Volodymyr Mnih
Tom Ward
Yotam Doron
Vlad Firoiu
Tim Harley
Iain Dunning
Shane Legg
Koray Kavukcuoglu

title: Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures

venue: In ICML,

authors:
Tommaso Furlanello
Zachary Lipton
Michael Tschannen
Laurent Itti
Anima Anandkumar

title: Born again neural networks

authors:
Matteo Hessel
Hubert Soyer
Lasse Espeholt
Wojciech Czarnecki
Simon Schmitt
Hado van Hasselt

title: Multi-task deep reinforcement learning with popart

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Geoffrey Hinton
Oriol Vinyals
Jeff Dean

title: Distilling the knowledge in a neural network

venue: arXiv preprint arXiv:1503.02531,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

venue: Neural computation,

authors:
Jan Humplik
Alexandre Galashov
Leonard Hasenclever
Pedro A Ortega
Yee Whye Teh
Nicolas Heess

title: Meta reinforcement learning as task inference

authors:
Max Jaderberg
Valentin Dalibard
Simon Osindero
Wojciech M Czarnecki
Jeff Donahue
Ali Razavi
Oriol Vinyals
Tim Green
Iain Dunning
Karen Simonyan

title: Population based training of neural networks

venue: arXiv preprint arXiv:1711.09846,

authors:
Adhiguna Kuncoro
Chris Dyer
Laura Rimell
Stephen Clark
Phil Blunsom

title: Scalable syntaxaware language models using knowledge distillation

venue: In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics,

authors:
Guillaume Lample
Alexandre Sablayrolles
Marc’Aurelio Ranzato
Ludovic Denoyer
Hervé Jégou

title: Large memory layers with product keys

venue: In Advances in Neural Information Processing Systems,

authors:
Michael Luo
Jiahao Yao
Richard Liaw
Eric Liang
Ion Stoica

title: Impact: Importance weighted asynchronous architectures with clipped target networks

authors:
Volodymyr Mnih
Adria Puigdomenech Badia
Mehdi Mirza
Alex Graves
Timothy Lillicrap
Tim Harley
David Silver
Koray Kavukcuoglu

title: Asynchronous methods for deep reinforcement learning

venue: In International conference on machine learning,

authors:
Emilio Parisotto
Jimmy Lei Ba
Ruslan Salakhutdinov

title: Actor-mimic: Deep multitask and transfer reinforcement learning

authors:
Emilio Parisotto
H Francis Song
Jack W Rae
Razvan Pascanu
Caglar Gulcehre
Siddhant M Jayakumar
Max Jaderberg
Raphael Lopez Kaufman
Aidan Clark
Seb Noury

title: Stabilizing transformers for reinforcement learning

venue: arXiv preprint arXiv:1910.06764,

authors:
Lerrel Pinto
Marcin Andrychowicz
Peter Welinder
Wojciech Zaremba
Pieter Abbeel

title: Asymmetric actor critic for image-based robot learning

venue: arXiv preprint arXiv:1710.06542,

authors:
Alec Radford
Jeff Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners

authors:
Jack W. Rae
Anna Potapenko
Siddhant M. Jayakumar
Chloe Hillier
Timothy P. Lillicrap

title: Compressive transformers for long-range sequence modelling

authors:
Benjamin Recht
Christopher Re
Stephen Wright
Feng Niu

title: Hogwild: A lock-free approach to parallelizing stochastic gradient descent

venue: In Advances in neural information processing systems,

authors:
A.A. Rusu
S.G. Colmenarejo
C. Gulcehre
G. Desjardins
J. Kirkpatrick
R. Pascanu
V. Mnih
K. Kavukcuoglu
R. Hadsell

title: Policy distillation

venue: In Proceedings of the International Conference on Learning Representations, ICLR”,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
H. Francis Song
Abbas Abdolmaleki
Jost Tobias Springenberg
Aidan Clark
Hubert Soyer
Jack W. Rae
Seb Noury
Arun Ahuja
Siqi Liu
Dhruva Tirumala
Nicolas Heess
Dan Belov
Martin Riedmiller
Matthew M. Botvinick

title: V-mpo: On-policy maximum a posteriori policy optimization for discrete and continuous control

authors:
R. Sutton
A. Barto

title: Reinforcement Learning: an Introduction

authors:
Yi Tay
Mostafa Dehghani
Dara Bahri
Donald Metzler

title: Efficient transformers: A survey

venue: arXiv preprint arXiv:2009.06732,

authors:
Yee Teh
Victor Bapst
Wojciech M Czarnecki
John Quan
James Kirkpatrick
Raia Hadsell
Nicolas Heess
Razvan Pascanu

title: Distral: Robust multitask reinforcement learning

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N. Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: Advances in Neural Information Processing Systems,

heading: 1 INTRODUCTION

text:
Compared to standard supervised learning domains, reinforcement learning presents unique challenges in that the agent must act while it is learning. In certain application areas, which we term actor-latency-constrained settings, there exists maximum latency constraints on the acting policy which limit its model size. These constraints on latency preclude typical solutions to reducing the computational cost of high capacity models, such as model compression or off-policy reinforcement learning, as it strictly requires a low computational-complexity model to be acting during learning. Here, the major constraint is that the acting policy model must execute a single inference step within a fixed budget of time, which we denote by Tactor – the amount of compute or resources used during learning is in contrast not highly constrained. This setting is ubiquitous within real-world application areas: for example, in the context of learning policies for robotic platforms, because of inherent limitations in compute ability due to power and weight considerations it is unlikely a large model could run fast enough to provide actions at the control frequency of the robot’s motors.
However, for many of these strictly actor-latency constrained settings there are orthogonal challenges involved which prevent ease of experimentation, such as the requirement to own and maintenance real robot hardware. In order to develop a solution to actor-latency constrained settings without needing to deal with substantial externalities, we focus on the related area of distributed onpolicy reinforcement learning (Mnih et al., 2016; Schulman et al., 2017; Espeholt et al., 2018). Here a central learner process receives data from a series of parallel actor processes interacting with the environment. The actor processes run step-wise policy inference to collect trajectories of interaction to provide to the learner, and they can be situated adjacent to the accelerator or distributed on different machines. With a large model capacity, the bottleneck in experiment run-times for distributed learning quickly becomes actor inference, as actors are commonly run on CPUs or devices without significant hardware acceleration available. The simplest solution to this constraint is to increase
the number of parallel actors, often resulting in excessive CPU resource usage and limiting the total number of experiments that can be run on a compute cluster. Therefore, while not a hard constraint, experiment run-time in this setting is largely dominated by actor inference speed. The distributed RL setting therefore presents itself as a accessible test-bed for solutions to actor-latency constrained reinforcement learning.
Within the domain of distributed RL, an area where reduced actor-latency during learning could make a significant impact is in the use of large Transformers (Vaswani et al., 2017) to solve partiallyobservable environments (Parisotto et al., 2019). Transformers (Vaswani et al., 2017) have rapidly emerged as the state-of-the-art architecture across a wide variety of sequence modeling tasks (Brown et al., 2020; Radford et al., 2019; Devlin et al., 2019) owing to their ability to arbitrarily and instantly access information across time as well as their superior scaling properties compared to recurrent architectures. Recently, their application to reinforcement learning domains has shown results surpassing previous state-of-the-art architectures while matching standard LSTMs in robustness to hyperparameter settings (Parisotto et al., 2019). However, a downside to the Transformer compared to LSTM models is its significant computational cost.
In this paper, we present a solution to actor-latency constrained settings, “Actor-Learner Distillation” (ALD), which leverages a continual form of Policy Distillation (Rusu et al., 2015; Parisotto et al., 2015) to compress, online, a larger “learner model” towards a tractable “actor model”. In particular, we focus on the distributed RL setting applied to partially-observable environments, where we aim to be able to exploit the transformer model’s superior sample-efficiency while still having parity with the LSTM model’s computational-efficiency during acting. On challenging memory environments where the transformer has a clear advantage over the LSTM, we demonstrate our Actor-Learner Distillation procedure provides substantially improved sample efficiency while still having experiment run-time comparable to the smaller LSTM.

heading: 2 BACKGROUND

text:
A Markov Decision Process (MDP) (Sutton & Barto, 1998) is a tuple of (S,A, T , γ,R) where S is a finite set of states, A is a finite action space, T (s′|s, a) is the transition model, γ ∈ [0, 1] is a discount factor and R is the reward function. A stochastic policy π ∈ Π is a mapping from states to a probability distribution over actions. The value function V π(s) of a policy π for a particular state s is defined as the expected future discounted return of starting at state s and executing π: V π(s) = ∑∞ t=0 γ
trt. The optimal policy π∗ is defined as the policy with the maximum value at every state, i.e. ∀s ∈ S, π ∈ Π we have V π∗(s) > V π(s). The optimal policy is guaranteed to exist (Sutton & Barto, 1998). In this work, we focus on Partially-Observable MDPs (POMDPs) which define environments which cannot observe the state directly and instead must reason over observations. POMDPs require the agent to reason over histories of observations in order to make an informed decision over which action to choose. This motivates the use of memory models such as LSTMs (Hochreiter & Schmidhuber, 1997) or transformers (Vaswani et al., 2017).
In this work, we use V-MPO (Song et al., 2020) as the main reinforcement learning algorithm. V-MPO, an on-policy value-based extension of the Maximum a Posteriori Policy Optimisation algorithm (Abdolmaleki et al., 2018), used an EM-style policy optimization update along with regularizing KL constraints to obtain state-of-the-art results across a wide variety of environments. Similar to Song et al. (2020), we use the IMPALA (Espeholt et al., 2018) distributed RL framework to parallelize acting and learning. Actor processes step through the environment and collect trajectories of data to send to the learner process, which batches actor trajectories and optimizes reinforcement learning objectives to update the policy parameters. The updated parameters are then communicated back to actor processes in order to maintain on-policyness.
In the following, we first refer to a single observation and associated statistics (reward, etc.) as a step. We refer to an environment step as the acquisition of a step from the environment after an action is taken. The total number of environment steps is a measure of the sample complexity of an RL algorithm. We refer to a step processed by an RL algorithm as an agent step. The total number of agent steps is a measure of the computational complexity of the RL algorithm. We use steps per second (SPS) to measure the speed of an algorithm. We refer to Learner SPS as the total number of agent steps processed by the learning algorithm per second. We refer to Actor SPS as the total number of environment steps acquired by a single actor per second.

heading: 3 ACTOR-LEARNER DISTILLATION

text:
Within the setting of distributed RL, a major challenge in applying transformers to reinforcement learning is their significant computational cost owing to their high actor latency. As an example, in Fig. 1 we present some results comparing a transformer to an LSTM on the I-Maze memory environment (see Sec. 5.1). In Fig. 1, Left, a 4-layer Gated Transformer-XL (GTrXL) (Parisotto et al., 2019) agent significantly surpasses a 32-dim. LSTM in terms of data efficiency, with the xaxis as environment steps. However, in the center diagram, when the x-axis is switched to wall-clock time the LSTM becomes the clearly more efficient model.
In order to work towards an ideal model with both high sample efficiency and low experiment run time, we analyzed which parts of the distributed actor-critic system led to the transformer’s main computational bottlenecks. As shown in the table on the right-hand-side of Fig. 1, we found the percentage of time that the asynchronous learner process was spending waiting for new trajectory data was substantially higher for the transformer than for the LSTM. Further looking at actor SPS between models revealed that the cause of this idling is mainly due to the transformer’s actor SPS being over 15 times slower than the LSTM. This substantial decrease in SPS could even render the transformer not viable in actor-latency-constrained settings, e.g. robotic control environments, which entails a hard constraint on actor inference speed.
These findings suggest that considerable improvements in training speed could be attainable if the actors produced data faster. While de-coupling actor and learner speeds using an off-policy RL algorithm could be a suitable solution for many applications, we desired a solution that respected strict constraints on actor latency, which precludes ever running a large model during inference. As an alternative solution which respected this constraint, we designed a model compression procedure termed “Actor-Learner Distillation” (ALD) which continually performs distillation between a largecapacity “Learner” model, which is trained using RL but never run directly in the environment, and a fast “Actor” model, which is trained using distillation from the learner and is used to collect data.
As we use an actor-critic training paradigm in this work, each of these models effectively comprises two functions, the policy and value function. We denote the actor model byMA = (πA, V πA ) and learner model byML = (πL, V πL ). The actor model is chosen with the appropriate model capacity to compute an inference step on the target hardware as fast as possible, i.e. within time Tactor. In contrast, there are far fewer constraints on learner modelML, as it is typically running on a central process with accelerated hardware. While there are no restrictions besides latency concerns on the model class of actor and learner in the context of ALD, in this work we choose to focus on the case of the actor model being an LSTM and the learner model being a transformer. As there is a clear data efficiency advantage to the transformer and a clear computational advantage to the LSTM, the success of ALD at extracting the benefits of both models can be more clearly recognized.
Actor-Learner Distillation proceeds with the learner modelML being trained using a standard reinforcement learning algorithm, with the exception of the environment data being generated by the
actor model. The actor model is trained using a policy distillation loss:
LπALD = Es∼πA [DKL(πA(·|s)||πL(·|s))] = Es∼πA [∑ a∈A πA(a|s) log πL(a|s) πA(a|s) ] (1)
Similar to previous distillation work in the multitask setting (Teh et al., 2017), we employ this loss bidirectionally: the actor is trained towards the learner policy and the learner policy is regularized towards the actor policy. This regularization loss on the learner was seen to enable smoother optimization. As actor and learner policy will naturally be different at some points during training, an off-policy RL algorithm for the learner could be thought to be required if the divergence is large enough. We found leveraging off-policy data to be effective for the actor, and sampled trajectory data from a replay buffer when computing Eq. 1 during actor model optimization.
Beyond a distillation loss, we experimented with a value distillation loss: LVALD = Es∼πA [ 1
2 (V πL (s)− V πA (s))2
] (2)
This loss functions to encourage actor model representations to model task reward structure. The use of value distillation in order to improve representations in a policy network has also been explored in concurrent work (Cobbe et al., 2020). Unlike the policy distillation loss, this loss is only used to train the actor model. The final Actor-Learner Distillation loss is:
LALD = απL π ALD + αV L V ALD (3)
where απ and αV are mixing coefficients to control the contribution of each loss.

heading: 4 DISTRIBUTED ACTOR-LEARNER DISTILLATION

text:
In order to make more efficient use of computational resources, we develop a distributed system for Actor-Learner Distillation based on IMPALA (Espeholt et al., 2018). An overview of this system is shown in Fig. 2, and we detail the function of each distinct process below, in the order of data flow:
Actor: There are NA parallel actors, each executing on single-threaded processes with only CPU resources available. Each Actor process steps through environment steps sequentially until a trajectory of TU time steps is gathered. Actions are sampled from a local copy of the actor policy πA. Once a completed trajectory is acquired, it is communicated to a Queue process.
Queue: The Queue process receives trajectories asynchronously across actors and accumulates them into batches. The batched trajectories are then passed to the Learner Runner.
Learner Runner: The Learner Runner process runsML in inference mode on the incoming batches of data using a hardware accelerator. This is done to provide (1) learning targets to the Distill processes downstream and (2) to make sure learner model initial memory states are updated for the Learner. Specifically whenML is a transformer model, we additionally batch over the time dimension to process data even more rapidly. The Learner Runner then passes the batched trajectories it received along with computed learner model outputs to two separate parallel processing streams, the Learner process and the Replay process.
Learner: The Learner process receives batched trajectories from the Learner Runner and computes, using a hardware accelerator, all relevant RL objectives (V-MPO in our case (Song et al., 2020)),
along with the loss in Eq. 1 to regularize the learner model towards the actor model. In the Learner process, only the learner model parameters are updated. Similarly to recent work on deep actor-critic algorithms (Song et al., 2020; Luo et al., 2020), we use a “target network” where we only communicate updated learner parameters to the Learner Runner every KL optimization steps. Although not used here, the Learner process can additionally receive data from the Replay process.
Replay: The Replay process manages a replay buffer containing a large store of previously collected batched trajectories. Incoming batches of trajectories from the Learner Runner are archived in a large first-in-first-out queue. The Learner and Distill processes can then request batches of trajectories which are uniformly sampled from this queue.
Distill: Distill processes request data from the Replay buffer and use the retrieved trajectories to compute the distillation loss in Eq. 3, which is then used to update the actor model parameters. Similar to the Learner process, we utilize a target network scheme which updates model parameters on the Actor processes every KA optimization steps.

heading: 4.1 IMPROVING DISTILL / RL STEP RATIO

text:
An observation found early in the development of ALD was the significance of the ”distillation steps per RL steps” (DpRL) ratio. The DpRL ratio measures how many agent steps are taken per second on the actor model (which takes ”distillation steps”) comparatively with the learner model’s agent steps per second (which takes ”reinforcement learning steps”). This is exemplified in the left side of Fig. 3, where in the Meta-Fetch environment with 3 objects (see Sec. 5.2), we ran ALD but set a fixed number of Actor agent steps for every Learner agent step. The graph shows a greatly enhanced sample efficiency when the number of actor agent steps is increased in relation to learner agent steps, until saturating at around 10 actor agent steps per learner agent step. Out of all hyperparameters we tested, we found that a high DpRL ratio was consistently the most critical parameter in increasing the sample efficiency of ALD. However, increasing the DpRL ratio naturally increased the total run time of the system, as now the distillation process became the main system bottleneck, and this encouraged investigation into parallelized ways of improving actor agent steps.
To avoid a complete system slowdown while maintaining a high DpRL ratio, we leveraged the parallelized training procedure HOGWILD! (Recht et al., 2011) as a replacement to the singleDistill-process system originally tested. In this setting, there are now ND parallel Distill processes, all of which asynchronously sample batched trajectories from the replay buffer to update actor model parameters. As there are typically more than 1 available accelerators on a single machine (see e.g. reference machines in App. B), the parallel Distill processes can be evenly distributed between available accelerators to make the best use of the resources at hand. In the right side of Figure 3, the HOGWILD! Actor-Learner Distillation variant consistently achieved better sample complexity with equivalent time complexity, across a wide sweep of hyper-parameter settings.

heading: 5 EXPERIMENTS

text:
Algorithm Details: For experiments, we use the V-MPO algorithm (Song et al., 2020) as the RL algorithm underlying each of the procedures we test. For ALD, we use V-MPO with V-trace correc-
tions (Espeholt et al., 2018) which worked slightly better in preliminary experiments. For baseline models V-MPO without V-trace worked better, as in Song et al. (2020). For all experiments, we use a single-layer LSTM of varying dimension depending on the environment as the actor model. For transformers, we use the Gated Transformer-XL (Parisotto et al., 2019), which we initially observed had better sample efficiency and optimization stability than standard transformers. We vary the number of transformer layers dependent on the difficulty of the environment. More details on the experimental procedure can be found in App. A.
Baselines: For each environment, we individually run the actor and learner model architectures used in ALD as baselines. As an additional baseline besides the learner and actor models in isolation, we introduce a model which maintains a distinct policy and value network. Here the policy network has the same architecture as the actor model in ALD (i.e. a small LSTM) and the value network has the same architecture as the learner model (i.e. a transformer). As the value function does not need to be run on actor processes, only the policy, this baseline achieves similar run time as ALD without requiring any extra insights. This baseline is an instance of Asymmetric Actor-Critic (Asymm. AC) (Pinto et al., 2017), but where instead of the value function having privileged access to state information it has privileged access to computational resources. Comparison of ALD to Asymm. AC will reveal whether the introduction of distillation losses has any positive effect over the simpler method of creating independent policy/value functions.

heading: 5.1 I-MAZE

text:
The I-Maze environment has the agent start at the top-left corner of a grid-world with the shape of an “I” (Fig. 7, Left). The agent must travel from the top-left corner to one of the bottom corners of the I, where the episode is ended and the agent receives reward. The particular corner the agent must travel to is revealed in the top-right corner of the I, which contains an “indicator” tile that is either 0 or 1 depending on whether the left or right bottom corner contains the reward. If the agent does not enter a terminating corner state in H time steps, the episode ends without a reward. The agent entering the correct bottom corner goal based on the indicator receives a reward of 1. Entering the incorrect corner results in the episode terminating with no reward. As memory is the main concern of these results instead of exploration, for larger maze results (i.e. a maze of width 15), we add a shaping reward of 0.01 whenever the agent visits a state along the main corridor of the “I”. Once the
agent has traversed a particular state, it cannot gain another shaping reward for returning there until the next episode. The agent has an orientation and observes every pixel directly in front of it until it reaches an occluding wall.
The environment provides a clear test on an agent’s ability to remember events over long horizons, as the only way to reliably terminate an episode successfully is by remembering the indicator observed near the start of the maze. In Figure 4, we provide complete reward curves for a maze of dimension 9× 9. We plot curves with two different x-axes: environment steps, measuring sample complexity, and wall-clock time. All curves were obtained on Reference Machine A (see Appendix B) under identical operating conditions, meaning wall-clock time is a fair comparison. We can first observe the clear sample efficiency gains of the transformer over the LSTM as well it’s poorer performance when considering wall-clock time. Furthermore, we can observe that Actor-Learner Distillation achieves significantly improved sample complexity over the stand-alone LSTM. This confirms the ability of ALD to recover the sample-efficiency gains of the learner model. In terms of wall-clock time, the Actor-Learner Distillation procedure is unmatched by any other procedure, demonstrating its clear practical advantage.
In contrast to Actor-Learner Distillation, the Asymm. AC baseline does not seem to perform as well despite equivalent model complexity and equivalent time complexity, as it contains both identical actor and learner models. To gain insight into this result, we plotted the per-seed curves of ActorLearner Distillation and Asymm. AC on the 15× 15 I-Maze task in Figure 6. We can clearly see an underlying pattern to the seed curves, all models have a local optimum at around 0.5 success where each model class spends a varying amount of time. A return of 0.5 suggests that the model has learnt to enter a goal, but not yet learnt to use the indicator to correctly determine which goal contains the positive reward. We measured how many environment steps each model class spent in this local minima out of the 10 million total, and found that the LSTM and Asymm. AC spend a majority of their time there. While it can be expected that the transformer can easily exit this optima due to its ability to directly look back in time, interestingly ALD seems to recover the same efficiency at escaping this minima as the stand-alone transformer. This suggests that ALD can successfully impart the learner model’s inductive biases to the actor model during learning in a way that is substantially more effective than just using the transformer as a value function (Asymm. AC). Finally, in the lefthand-side of Fig. 6 we ran an ablation on 9 × 9 I-Maze to determine whether the value distillation loss provided benefit to learning. We performed a hyperparameter sweep for each setting of value distillation enabled or disabled. The results demonstrated that using value distillation provides a slight but significant improvement.

heading: 5.2 META-FETCH

text:
The “Meta-Fetch” environment, shown in Figure 7, requires an agent to fetch a number of objects distributed randomly on an empty gridworld, with each object required to be fetched in a particular hidden sequence. Crucially, in Meta-Fetch (1) the agent can not sense the unique identity of each object, meaning each object looks the same to the agent, (2) the agent does not know beforehand the sequence of objects it must obtain, and (3) observations consist of local views of pixels in front of the agent’s current orientation (see Fig. 7). When the agent collects an object, it can either receive a positive reward if this object is the next correct one in the sequence, or receive no reward and
the sequence of objects is reset. To prevent cheating (i.e. the agent obtaining a reward and then immediately resetting the sequence, and repeating), the agent receives a positive reward only on the first instance of it collecting the next correct object in the sequence. Once all objects in the sequence have been fetched, the objects and rewards are reset and the environment proceeds as it did at the start of the episode (except the agent is now in the position of the last object fetched). Meta-Fetch presents a highly challenging problem for memory architectures because the agent must simultaneously learn the identity of objects by mapping their spatial relations using only local egocentric views, as well as learn a discrete search strategy to memorize the correct sequence of objects to hit. Meta-Fetch was inspired by the meta-learning environment “Number Pad” used in previous work (Humplik et al., 2019; Parisotto et al., 2019).
Results are reported in Figure 5 for a Meta-Fetch environment with 4 randomly located objects. We present plots of average return using both environment steps and wall-clock time as x-axes (left plots), along with an additional plot which measures “success” (right plots). Success in MetaFetch is defined as the agent completing at least 1 full object fetch sequence (agents can accomplish multiple full fetch sequences for additional reward). As in I-Maze, we re-confirm the observation that transformers achieve a clearly superior sample efficiency over LSTMs, while at the cost of a much slower wall-clock performance. Additionally, unlike I-Maze where Asymm. AC had a more positive effect on sample efficiency, there is a less substantial difference between LSTM and Asymm. AC.
However, we can see a major difference between the baselines and Actor-Learner Distillation, which early on achieves close to the transformer’s sample efficiency with a much smaller model. Additionally, it seems that in this much more challenging environment, which requires both memory and compositional reasoning, the LSTM actor model in ALD suffers in performance later on when the GTrXL achieves a higher rate of increase in reward. Examining the success rate reveals an interesting trend where the ALD-trained LSTM achieves better generalization than the LSTM and Asymm. AC models. In particular, we can see that although the LSTM and Asymm. AC have average return that is increasing, their success rate is relatively stagnant (Asymm. AC) or even decreasing (LSTM). This suggests the increase in return these baselines are observing mainly stem from them learning how to search through a fraction of possible object layouts more efficiently. In contrast, the ALDtrained LSTM’s success rate is correlated with its return, suggesting it is learning to succeed in a larger variety of object layouts.

heading: 6 RELATED WORK

text:
Whether a deep neural network could be compressed into a simpler form was examined shortly after deep networks were shown to have widespread success (Ba & Caruana, 2014; Hinton et al., 2015). Distillation objectives revealed that not only could high-capacity teacher networks be transferred to low-capacity student networks (Ba & Caruana, 2014), but that training using these imitation objectives produced a superior low-capacity model compared to training from scratch (Hinton et al., 2015), with this effect generalizing to the case where student and teacher share equivalent capacity (Furlanello et al., 2018). Similar to this work’s desired goal to leverage the rapid memorization capabilities of transformers while training an LSTM, other work showed that distillation between different architecture classes could enable the transfer of inductive biases (Kuncoro et al., 2019; Abnar et al., 2020). Within RL, distillation has most often been used as a method for stabilizing multitask learning (Rusu et al., 2015; Parisotto et al., 2015; Teh et al., 2017; Berseth et al., 2018).
The method most closely related to ALD is Mix&Match (Czarnecki et al., 2018), where an ensemble of policies is defined such that each successive model in the ensemble has increasing capacity. All models in the ensemble are distilled between each other, and a mixture of the ensemble is used for sampling trajectories, with the mixture coefficients learnt through population-based training (Jaderberg et al., 2017). There are however significant difference from ALD: the distillation procedure was meant to transfer simple, quick-learning skills from a low-capacity model to a high-capacity model (the opposite intended direction of the benefit of distillation in ALD) and a mixture policy was used to sample trajectories which did not consider actor-latency constraints (and preferred sampling from the larger model if it performed better).
Similar to ALD, various works have looked at exploiting asymmetries between acting and learning in reinforcement learning (Pinto et al., 2017; Humplik et al., 2019). In (Pinto et al., 2017), since the value function is ultimately not needed during deployment, it was given privileged access to state information during training (whereas the policy must operate directly from pixels). Our baseline
Asymm. AC was derived from this insight that value networks are solely needed during learning, but our variant had a value function that utilized privileged compute resources instead of state access during training. Another closely related method is the training paradigm of Humplik et al. (2019), where a belief network was trained to predict hidden state information in partially-observable metalearning environments. The hidden representations from the belief network were then used as an auxiliary input to an agent network trained to solve the task, enabling representation learning to use privileged information during training.
As an alternative to model compression, architecture development can reduce the computational cost of certain model classes. Transformers in particular have had a large amount of attention towards designing more efficient alternatives (Tay et al., 2020), as gains from increased model scale have yet to saturate (Brown et al., 2020). This architectural development has taken many forms: sparsified attention (Child et al., 2019), compressed attention (Rae et al., 2020; Dai et al., 2020), use of kNN instead of soft-attention (Lample et al., 2019), etc. (See Tay et al. (2020) for a comprehensive review). However, architectural development is largely orthogonal to the ALD procedure, as it means we can further scale actor and learner models by a corresponding degree.

heading: 7 CONCLUSION

text: In conclusion, towards the development of efficient learning methods in actor-latency-constrained settings, we developed the Actor-Learner Distillation (ALD) procedure, which leverages a continual form of model compression between separate actor and learner models. In the context of learning in partially-observable environments using a distributed actor-critic system, ALD successfully demonstrated the ability to largely capture the sample efficiency gains of a larger transformer learner model while still maintaining the reduced computational cost and lower experiment run time of the LSTM actor model. As supervised learning demonstrates ever increasing gains in performance from increasing model capacity, we believe the development of effective methods of model compression for RL will become a more prominent area of study in the future. In future work, we wish to investigate whether integrating the ALD procedure with batched inference for actors would still maintain the same performance increases we demonstrated in our results, while at the same time enabling larger actor models to be used and correspondingly larger learners.

heading: ACKNOWLEDGMENTS

text: This work was supported in part by the NSF IIS1763562 and ONR Grant N000141812861. We would also like to acknowledge NVIDIA’s GPU support.

heading: A EXPERIMENT DETAILS

text:
For all models, we sweep over the V-MPO target network update frequency KL ∈ {1, 10, 100}. In initial experiments, we also sweeped the “Initial α” setting over values {0.1, 0.5, 1.0, 5.0}. All experiment runs have 3 unique seeds. For each model, we choose the hyperparameter setting that achieved highest mean return over all seeds. Additionally we use PopArt (Hessel et al., 2019) for the value output.
A.1 I-MAZE EXPERIMENTS
Observations, Actions and Metrics: Observations in I-Maze are the single row of pixels starting at the agent’s current position and extending in the direction the agent is facing. The row of pixels extends as far as the maze dimension. If the agent has a wall in its field of view (represented as black pixels in Fig. 7), pixels further away from that point are occluded. For the 9 × 9 maze, a reward is only given when the agent enters the correct goal as decided by the indicator. For the 15 × 15 maze, it is the same but there is an additional shaping reward where whenever the agent enters one of the states along the central column of the “I” for the first time within its current episode, it receives a small reward of 0.01. This reward is meant to encourage exploration and prevent the very large number of environment steps otherwise necessary for the agent to end up at a goal when
taking actions randomly. The agent has 4 actions, move forward, turn left, turn right and do nothing. Moving forward into a wall causes no change in state. 9 × 9 I-Maze episodes last for a horizon of 150 time steps while 15 × 15 I-Maze episodes last for 350 steps. Observations have 3 channel dimensions, one for whether the pixel is free space or a wall, one for the green goal indicator and one for the red goal indicator.
LSTM: We use a single-layer LSTM with a hidden dimension of 32 for all experiments in this domain.
GTrXL: We use a 2-layer GTrXL for the experiments in 9× 9 and a 4-layer GTrXL for the experiments in 9×9. We use an embedding size of 256, 8 attention heads, a head dimension of 32, a gated initialization bias of 2 and a memory length of 64. Other details were followed from Parisotto et al. (2019).
ALD: We use the corresponding LSTMs and GTrXL described above depending on the environment. During hyperparameter sweeps, we tested KA ∈ {1, 10, 100}. We set απ = 1 and sweep αV ∈ {0, 0.1, 1}. Asymm. AC: We use the corresponding LSTMs and GTrXL described above depending on the environment.
A.2 META-FETCH
Observations, Actions and Metrics: In Meta-Fetch, an agent acts in an empty 2D gridworld of size 7 × 7 given only local observations. At each step the agent can choose to either move forward or turn left or right. Observations are the 15 pixels in front of the agent pointing in it’s orientation, along with the 2 pixels on either side of this ray, for a complete observation size of C × 15 × 3 pixels (see Fig. 7). Each observation has C channels which determine whether that pixel represents free space, a wall, an object, or a collected object. Every episode has the agent located in a new configuration of object locations. We set a maximum episode length of 300 steps for Meta-Fetch.
The goal of the agent in Meta-Fetch is to collect, by colliding with, each of the randomly located objects in a specific sequence. This sequence is not known beforehand and is randomized each episode. The only clue to discovering this sequence is that when the agent hits the next right object in the sequence, it gains a reward of 1 and the object is sensed as “collected”. Otherwise, if it hit the wrong object in the sequence, then all previously collected objects are reset and the agent must restart the fetch sequence from the beginning once again. Once all objects are collected, they are all reset and the agent can restart the same sequence to once again gain reward for each collected reward. Objects which have been reset before the complete fetch sequence is achieved do not give further reward upon re-collection, to prevent the agent from falling into an local optimum of repeatedly collecting a correct object and then immediately collecting an incorrect object to reset the correct object.
LSTM: We use a single-layer LSTM with a hidden dimension of 128.
GTrXL: We use a 4-layer GTrXL, an embedding size of 256, 8 attention heads, a head dimension of 32, a gated initialization bias of 2 and a memory length of 64. Other details were followed from Parisotto et al. (2019).
ALD: We use the corresponding LSTMs and GTrXL described above depending on the environment. During hyperparameter sweeps, we tested KA ∈ {10, 100}. We set απ = 1 and sweep αV ∈ {0, 0.1, 1}. Asymm. AC: We use the corresponding LSTMs and GTrXL described above depending on the environment.

heading: B COMPUTE DETAILS

text:
Reference Machine A: Reference Machine A has a 36-thread Intel(R) Core(TM) i9-7980XE CPU @ 2.60GHz, 64GB of RAM, and 2 GPUs: a GeForce GTX 1080 Ti and a TITAN V.
Reference Machine B: Reference Machine B has a 40-thread Intel(R) Xeon(R) CPU E5-2630 v4 @ 2.20GHz, 256GB of RAM and 2 GPUs: a Tesla P40 and a Tesla V100-PCIE-16GB.

heading: C DESCRIPTION OF HOGWILD!

text: In order to speed-up certain learning aspects of our proposed method, we use the HOGWILD! (Recht et al., 2011) algorithm. HOGWILD! enables lock-free stochastic gradient optimization to be performed by several parallel learning processes. Each process has shared access to the parameter vector of the learning model, and can update it without requiring a mutually-exclusive lock. This allows parameter updates to occur in parallel as soon as they are available, but potentially leads to race conditions where the parameter vector is read during it being updated by a different process. This means a fraction of the parameter vector can be stale when a learning process calculates the parameter gradients. However, while this parallel overwriting of the parameter vector introduces noise via this staleness, it usually does not significantly inhibit performance and is typically faster than alternatives which use explicit locking to prevent race conditions.


summaries: The paper addresses the problem of multi-domain few-shot image classification (where unseen classes and examples come from diverse data sources), and proposes a Universal Representation Transformer (URT) layer, which learns to transform a universal representation into task-adapted representations. The method proposed builds on top of SUR [Dvornik et al 2020], where a universal representation is extracted from the outputs of a collection of pre-trained and domain-specific backbones and a selection procedure infers how to weight each backbone for a given task at hand. While SUR inferred those weights by optimising a loss on the support set (the few examples provided in a task), the authors in this paper introduce an attention-based layer (inspired by Vaswani et al Transformer) that learns to weight the appropriate backbones for each task. This layer has the main advantage that it can be learned across few-shot tasks from many domains so it can support transfer across these tasks.

summaries: Few-shot learning on meta-dataset is challenging due to the domain gap between train and validation. In order to bridge this gap, the authors present a model that learns to combine domain-specific representations to generalize to new domains. This combination is done with a transformer model that pays attention to the features extracted from domain-specific backbones. The authors demonstrate empirically that their model attains comparable performance to previous state-of-the-art at higher efficiency and include ablation results to test their model components.

summaries: The paper presents a method for tackling multi-domain few-shot image classification problem where it obtains a task-adapted representation by weighing representations from pretrained domain-specific backbones according to the support set at hand. The desirable property of this framework is that the model can leverage information from other domains to make predictions. The effectiveness of Universal Representations have been discussed in the past work - SUR [1], and this work builds on top of it and introduces a learnable component (self-attention), and showed the improvement both quantitatively and qualitatively.

abstractText: Few-shot classification aims to recognize unseen classes when presented with only a small number of samples. We consider the problem of multi-domain few-shot image classification, where unseen classes and examples come from diverse data sources. This problem has seen growing interest and has inspired the development of benchmarks such as Meta-Dataset. A key challenge in this multi-domain setting is to effectively integrate the feature representations from the diverse set of training domains. Here, we propose a Universal Representation Transformer (URT) layer, that meta-learns to leverage universal features for few-shot classification by dynamically re-weighting and composing the most appropriate domain-specific representations. In experiments, we show that URT sets a new state-of-the-art result on Meta-Dataset. Specifically, it achieves top-performance on the highest number of data sources compared to competing methods. We analyze variants of URT and present a visualization of the attention score heatmaps that sheds light on how the model performs cross-domain generalization. Our code is available at https://github.com/liulu112601/URT.

authors:
Ferran Alet
Tomás Lozano-Pérez
Leslie P Kaelbling

title: Modular meta-learning

venue: In Conference on Robot Learning (CoRL),

authors:
Peyman Bateni
Raghav Goyal
Vaden Masrani
Frank Wood
Leonid Sigal

title: Improved few-shot visual classification

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Hakan Bilen
Andrea Vedaldi

title: Universal representations: The missing link between faces, text, planktons, and cat breeds

venue: arXiv preprint arXiv:1701.07275,

authors:
Wei-Yu Chen
Yen-Cheng Liu
Zsolt Liu
Yu-Chiang Frank Wang
Jia-Bin Huang

title: A closer look at few-shot classification

venue: In International Conference on Learning Representations (ICLR),

authors:
Yinbo Chen
Xiaolong Wang
Zhuang Liu
Huijuan Xu
Trevor Darrell

title: A new meta-baseline for few-shot learning

venue: arXiv preprint arXiv:2003.04390,

authors:
Guneet Singh Dhillon
Pratik Chaudhari
Avinash Ravichandran
Stefano Soatto

title: A baseline for few-shot image classification

authors:
Nikita Dvornik
Cordelia Schmid
Julien Mairal

title: Selecting relevant features from a universal representation for few-shot classification

venue: arXiv preprint arXiv:2003.09338,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: In The International Conference on Machine Learning (ICML),

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

venue: Technical report, University of Toronto,

authors:
Brenden M Lake
Ruslan Salakhutdinov
Joshua B Tenenbaum

title: Human-level concept learning through probabilistic program induction

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE,

authors:
Hongyang Li
David Eigen
Samuel Dodge
Matthew Zeiler
Xiaogang Wang

title: Finding taskrelevant features for few-shot learning by category traversal

authors:
Zhouhan Lin
Minwei Feng
Cicero Nogueira dos Santos
Mo Yu
Bing Xiang
Bowen Zhou
Yoshua Bengio

title: A structured self-attentive sentence embedding

venue: arXiv preprint arXiv:1703.03130,

authors:
Lu Liu
Tianyi Zhou
Guodong Long
Jing Jiang
Lina Yao
Chengqi Zhang

title: Prototype propagation networks (PPN) for weakly-supervised few-shot learning on category graph

venue: In International Joint Conferences on Artificial Intelligence (IJCAI),

authors:
Lu Liu
Tianyi Zhou
Guodong Long
Jing Jiang
Chengqi Zhang

title: Learning to propagate for graph meta-learning

venue: In Neural Information Processing Systems (NeurIPS),

title: Attribute propagation network for graph zero-shot learning

venue: In AAAI Conference on Artificial Intelligence (AAAI),

authors:
Yang Liu
Chengjie Sun
Lei Lin
Xiaolong Wang

title: Learning natural language inference using bidirectional lstm model and inner-attention

venue: arXiv preprint arXiv:1605.09090,

authors:
Boris Oreshkin
Pau Rodrı́guez López
Alexandre Lacoste

title: Tadam: Task dependent adaptive metric for improved few-shot learning

venue: In The Conference on Neural Information Processing Systems (NeurIPS),

authors:
Ethan Perez
Florian Strub
Harm de Vries
Vincent Dumoulin
Aaron C. Courville

title: Film: Visual reasoning with a general conditioning layer

authors:
Sylvestre-Alvise Rebuffi
Hakan Bilen
Andrea Vedaldi

title: Efficient parametrization of multidomain deep neural networks

authors:
Mengye Ren
Eleni Triantafillou
Sachin Ravi
Jake Snell
Kevin Swersky
Joshua B Tenenbaum
Hugo Larochelle
Richard S Zemel

title: Meta-learning for semi-supervised few-shot classification

authors:
James Requeima
Jonathan Gordon
John Bronskill
Sebastian Nowozin
Richard E Turner

title: Fast and flexible multi-task classification using conditional neural adaptive processes

authors:
Tonmoy Saikia
Thomas Brox
Cordelia Schmid

title: Optimized generic feature learning for fewshot classification across domains

venue: arXiv preprint arXiv:2001.07926,

authors:
Jake Snell
Kevin Swersky
Richard Zemel

title: Prototypical networks for few-shot learning

authors:
Eleni Triantafillou
Tyler Zhu
Vincent Dumoulin
Pascal Lamblin
Utku Evci
Kelvin Xu
Ross Goroshin
Carles Gelada
Kevin Swersky
Pierre-Antoine Manzagol

title: Meta-dataset: A dataset of datasets for learning to learn from few examples

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Oriol Vinyals
Charles Blundell
Tim Lillicrap
Daan Wierstra

title: Matching networks for one shot learning

authors:
2016. Han-Jia Ye
Hexiang Hu
De-Chuan Zhan
Fei Sha

title: Few-shot learning via embedding adaptation

authors:
2020. Sung Whan Yoon
Jun Seo
Jaekyun Moon

title: Tapnet: Neural network augmented with task

authors:
2020 Triantafillou et al
2020 Dhillon et al
Saikia et al
2020). Specifically
Saikia

title: 2020) have shown that competitive performance can be reached using a strong hyper-parameter optimization method applied on a carefully designed validation metric appropriate for few-shot learning. C STRUCTURES AND TRAINING STRATEGY OF BACKBONES

heading: 1 INTRODUCTION

text:
Learning tasks from small data remains a challenge for machine learning systems, which show a noticeable gap compared to the ability of humans to understand new concepts from few examples. A promising direction to address this challenge is developing methods that are capable of performing transfer learning across the collective data of many tasks. Since machine learning systems generally improve with the availability of more data, a natural assumption is that few-shot learning systems should benefit from leveraging data across many different tasks and domains—even if each individual task has limited training data available.
This research direction is well captured by the problem of multi-domain few-shot classification. In this setting, training and test data spans a number of different domains, each represented by a different source dataset. A successful approach in this multi-domain setting must not only address the regular challenge of few-shot classification—i.e., the challenge of having only a handful of examples per class. It must also discover how to leverage (or ignore) what is learned from different domains, achieving generalization and avoiding cross-domain interference.
Recently, Triantafillou et al. (2020) proposed a benchmark for multi-domain few-shot classification, Meta-Dataset, and highlighted some of the challenges that current methods face when training data is heterogeneous. Crucially, they found that methods which trained on all available domains would normally obtain improved performance on some domains at the expense of others. Following on their work, progress has been made, which includes the design of adapted hyper-parameter optimization strategies (Saikia et al., 2020) and more flexible meta-learning algorithms (Requeima et al., 2019). Most notable is SUR (Selecting Universal Representation) (Dvornik et al., 2020), a method that relies on a so-called universal representation, extracting from a collection of pre-trained and domain-specific neural network backbones. SUR prescribes a hand-crafted feature-selection procedure to infer how to weight each backbone for each task at hand, and produces an adapted representation for each task. This was shown to lead to some of the best performances on Meta-Dataset. ∗This work was done while Lu Liu was a research intern with Mila. †Canada CIFAR AI Chair
In SUR, the classification procedure for each task is fixed and not learned. Thus, except for the underlying universal representation, there is no transfer learning performed with regards to how classification rules are inferred across tasks and domains. Yet, cross-domain generalization might be beneficial in that area as well, in particular when tasks have only few examples per class.
Present work. To explore this question, we propose a Universal Representation Transformer (URT) layer, which can effectively learn to transform a universal representation into task-adapted representations. The URT layer is inspired from Transformer (Vaswani et al., 2017) and uses an attention mechanism to learn to retrieve or blend the appropriate backbones to use for each task. By training this layer across few-shot tasks from many domains, it can support transfer across these tasks.
We show that our URT layer on top of a universal representation’s pre-trained backbones sets a new state-of-the-art performance on Meta-Dataset. It succeeds at outperforming SUR on 4 dataset sources without impairing accuracy on the others. This leads to top performance on 7 dataset sources when comparing to a set of competing methods. To interpret the strategy that URT learns to weigh the backbones from different domains, we visualize the attention scores for both seen and unseen domains and find that our model generates meaningful weights for the pre-trained domains. A comprehensive analysis on variants and ablations of the URT layer is provided to show the importance of various components of URT, notably the number of attention heads.

heading: 2 FEW-SHOT CLASSIFICATION

heading: 2.1 PROBLEM SETTING

text:
In this section, we will introduce the problem setting for few-shot classification and the formulation of meta-learning for few-shot classification. Few-shot classification aims to classify samples where only few examples are available for each class. We describe a few-shot learning classification task as the pair of examples, comprising of a support set S to define the classification task and the query set Q of samples to be classified.
Meta-learning is a technique that aims to model the problem of few-shot classification as learning to learn from instances of few-shot classification tasks. The most popular way to train a meta-learning model is with episodic training. Here, tasks T = (Q,S) are sampled from a larger dataset by taking subsets of the dataset to build a support set S and a query set Q for the task. A common approach is to sample N -way-K-shot tasks, each time selecting a random subset of N classes from the original dataset and choosing only K examples for each class to add to the support set S.
The meta-learning problem can then be formulated by the following optimization:
min Θ
E(S,Q)∼p(T ) [L(S,Q,Θ)] , L(S,Q,Θ) = 1 |Q| ∑
(x,y)∼Q
− log p(y|x, S; Θ) + λΩ(Θ), (1)
where p(T ) is the distribution of tasks, Θ are the parameters of the model and p(y|x, S; Θ) is the probability assigned by the model to label y of query example x (given the support set S), and Ω(Θ) is an optional regularization term on the model parameters with factor λ.
Conventional few-shot classification targets the setting of N -way-K-shot, where the number of classes and examples are fixed in each episode. Popular benchmarks following this approach include Omniglot (Lake et al., 2015) or benchmarks made of subsets of ImageNet, such as miniImageNet (Vinyals et al., 2016) and tieredImageNet (Ren et al., 2018). In such benchmarks, the tasks for training cover a set of classes that is disjoint from the classes in the test set of tasks. However, with the training and test sets tasks coming from a single dataset/domain, the distribution of tasks found in either sets is similar and lacks variability, which may be unrealistic in practice.
It is in this context that Triantafillou et al. (2020) proposed Meta-Dataset, as a further step towards large-scale, multi-domain few shot classification. Meta-Dataset includes ten datasets (domains), with eight of them available for training. Additionally, each task sampled in the benchmark varies in the number of classes N , with each class also varying in the number of shots K. As in all few-shot learning benchmarks, the classes used for training and testing do not overlap.

heading: 2.2 BACKGROUND AND RELATED WORK

text:
Meta-Learning A promising approach for few-shot classification is to use meta-learning to more directly train a model to learn to perform few-shot classification, in an end-to-end way. The two most popular methods are Prototypical Networks (Snell et al., 2017) and Model Agnostic Meta-Learning (MAML) (Finn et al., 2017). Triantafillou et al. (2020) showed that prototypical networks and MAML could be combined by leveraging prototypes for the initialization of the output weights value in the inner loop. Requeima et al. (2019) also proposed Conditional Neural Adaptive Processes (CNAPs) for few-shot classification, which can be seen as extending prototypical networks with a more sophisticated architecture that allows for improved task adaptation. This architecture was later improved further by Bateni et al. (2020) with Simple CNAPS, leading to one of the current best methods on Meta-Dataset. Another line of work which leverages the idea of “transfer by fine-tuning” can be found in Appendix B.
Universal Representations In contrast, our work instead builds on that of Dvornik et al. (2020) and their method SUR (Selecting from Universal Representations). Bilen & Vedaldi (2017) introduced the term universal representation to refer to a representation that supports good performance in multiple domains. One proposal towards such a representation is to train different neural networks backbones separately on the data of each available domain, then simply to concatenate the representation learned by each. Another is to introduce some parameter sharing between the backbones, by having a single network conditioned on the domain of the provenance of each batch of training data (Rebuffi et al., 2018), e.g. using Feature-wise Linear Modulate (FiLM) (Perez et al., 2018). SUR proposes to leverage a universal representation in few-shot learning tasks with a feature selection procedure that assigns different weights to each of the domain-specific subvectors of the universal representation. The objective is to assign high weights only to the domain-specific representations that are specifically useful for each few-shot task at hand. The weights are inferred by optimizing a loss on the support set that encourages high accuracy of a nearest-centroid classifier. As such, the method does not involve any meta-learning—a choice motivated by the concern that meta-learning may struggle in generalizing to domains that are dissimilar to the training domains. SUR achieved some of the best performances on Meta-Dataset. However, a contribution of our work is to provide evidence that meta-learning can actually be used to replace SUR’s hand-designed inference procedure and improve performance further.
Task Adaptive Representations Another line of work tries to retrieve task adaptive representations for each task. Task specific representations can be conditioned on a representation of the current task (Oreshkin et al., 2018; Wang et al., 2019), projected to another space (Yoon et al., 2019), or masked based on inter-class commonality and inter-class uniqueness (Li et al., 2019). While the representation extracted from URT is also task adaptive, it is adaptive to a set of pretrained backbones and can be applied to more complicated multi-domain scenarios. Wang & Hebert (2016) proposed to improve a CNN by adding extra layers and train it using unsupervised data while our contribution mainly lies in composing representations instead of an improved CNN. Alet et al. (2018) introduced a modular meta-learning method, which learns a repertoire of modules that serves as nodes to construct a tree structure to solve a new robotic-related task. Comparatively, URT is a one-for-all layer which doesn’t need to construct different module structures for each task.
Transformer Networks Our meta-learning approach to leverage universal representations is inspired directly from Transformer networks (Vaswani et al., 2017). Our model structure is inspired by the structure of the dot-product self-attention in the Transformer, which we adapted here to multidomain few-shot learning by designing appropriate parametrizations for queries, keys and values. Self-attention was explored in the single-domain training regime by Ye et al. (2020); Liu et al. (2019b;a; 2020), however for a different purpose, where each representation of individual examples in a task support set is influenced by all other examples. Rather than using self-attention between individual examples in the support set, our model uses self-attention to select between different domain-specific backbones.

heading: 3 UNIVERSAL REPRESENTATION TRANSFORMER LAYER

text:
In this section, we describe our proposed URT layer, which uses meta-learning episodic training to learn how to combine the domain-specific backbones of a universal representation for any given fewshot learning classification task. URT layer can be built on top of any set of pretrained backbones
without further costly fine-tuning of the backbones. More details on how to train multiple domainspecific backbones can be found in Appendix C.
Conceptually, the proposed model views the support set S of a task as providing information on how to query and retrieve from the set {ri} of m pre-trained backbones the most appropriate backbone to build an adapted representation φ for the task.
We would like the model to support a variety of strategies on how to retrieve backbones. For example, it might be beneficial for the model to retrieve a single backbone from the set, especially if the domain of the given task matches perfectly that of a domain found in the training set. Alternatively, if some of the training domains benefit from much more training data than others, a better strategy might be to attempt some cross-domain generalization towards the few-shot learning task by blending many backbones together, even if none matches the domain of the task perfectly.
This motivates us to use dot-product self-attention, inspired by layers of Transformer networks (Vaswani et al., 2017). For this reason, we refer to our model as a Universal Representation Transformer (URT) layer. Additionally, since each class of the support set might require a different strategy, we perform attention separately for each class and their support set Sc = {x|(x, y) ∈ S and y = c}.

heading: 3.1 SINGLE-HEAD URT LAYER

text:
We start by describing an URT layer consisting of a single attention head. An illustration of a singlehead URT layer is shown in Figure 1. Let ri(x) be the output vector of the backbone for domain i. We then write the universal representation as
r(x) = concat(r1(x), . . . , rm(x)). (2)
This representation provides a natural starting point to obtain a representation of a support set class. Specifically, we will note
r(Sc) = 1 |Sc| ∑ x∈Sc r(x) (3)
as the representation for the set Sc. From this, we can describe the URT layer by defining the queries1, keys, the attention mechanism and output of the layer:
Queries qc: For each class c, we obtain a query through qc = Wqr(Sc) + bq , where we have a learnable query linear transformation represented by matrix Wq and bias bq .
Keys ki,c: For each domain i and class c, we define keys as ki,c = Wkri(Sc)+bk, using a learnable linear transformation Wk and bk and where ri(Sc) = 1/|Sc| ∑ x∈Sc ri(x), using a similar notation as for r(Sc). 1Unable to avoid the unfortunate double usage of the term ”query” due to conflicting conventions, we highlight the difference between the query sets Q of few-shot tasks and the queries qc of an attention mechanism.
Algorithm 1 Training of URT layer Input: Number of tasks τtotal, m pre-trained backbones ;
1: for τ ∈ {1, · · · , τtotal} do 2: Sample a few-shot task T with support set S and query set Q; 3: # Infer adapted representation for task from S 4: For each class, obtain representation using m pre-trained backbones as in Eq. (3); 5: Obtain attention scores using Eq. (4,5) for each head using support set S; 6: # Use adapted representation to predict labels in Q from support set S 7: Compute adapted representation of examples in S and Q as in Eq. (6,7); 8: Compute probabilities of label of examples in Q using Prototypical Network as in Eq. (9); 9: Compute loss as in Eq. (1,8) and perform gradient descent step on URT parameters Θ;
10: end for
Attention scores αi: as for regular Transformer layers, we use scaled dot-product attention
αi,c = exp(βi,c)∑ i′ exp(βi′,c) , βi,c = qc >ki,c√ l , (4)
where l is the dimensionality of the keys and queries. Then, these per-class scores are aggregated to obtain scores for the full support set by averaging
αi = ∑ c αi,c N . (5)
Equipped with these attention scores, the URT layer can now produce an adapted representation for the task (for the support and query set examples) by computing
φ(x) = ∑ i αiri(x) . (6)
As we can see, this approach has the flexibility of either selecting a single domain-specific backbone (by assigning αi = 1 for a single domain) or blending different domains together (by having αi >> 0 for multiple backbones).

heading: 3.2 MULTI-HEAD URT LAYER

text:
The URT layer described so far can only learn to retrieve a single backbone (or blending of backbones). Yet, it might be beneficial to retrieve multiple different (blended) backbones, especially for a few-shot task that would include many classes of varying complexity.
Thus, to achieve such diversity in the adapted representation, we also consider URT layers with multiple heads, i.e. where each head corresponds to the calculation of Equation 6 and each head has its own set of parameters (Wq,bq,Wk,bk). Denoting each head now as φh, a multi-head URT layer then produces as its output the concatenation of all of its heads:
φ(x) = concat(φ1(x), . . . , φH(x)). (7)
Empirically we found that the randomness in the initialization of head weights alone did not lead to uniqueness and being complimentary between the heads, so inspired by Lin et al. (2017), we add a regularizer to avoid duplication of the attention scores:
Ω(Θ) = ‖(AA> − I)‖F 2 , (8)
where ‖ · ‖F is the Frobenius norm of a matrix and A ∈ Rn×m is the matrix for attention scores, with Ah being the vector of all scores αi for head h. The identity matrix I regularizes each set of attention scores to be more focused so that multiple heads can attend to different domain-specific backbones.

heading: 3.3 TRAINING STRATEGY

text:
We train representations produced by the URT layer by following the approach of Prototypical Networks (Snell et al., 2017), where the probability of a label y for a query example x given the
support set of a task is modeled as:
p(y = c|x, S; Θ) = exp(−d(φ(x)− pc))∑N c′=1 exp(−d(φ(x)− pc′)) , (9)
where d is a distance metric and pc = 1/|Sc| ∑
x∈Sc φ(x) corresponds to the centroid of class c, referred to as its prototype. We use (negative) cosine similarity as the distance. The full training algorithm is presented in Algorithm 1.

heading: 4 EXPERIMENTS

text:
In this section, we seek to answer three key experimental questions:
Q1 How does URT compare with previous state-of-the-art on Meta-Dataset for multi-domain fewshot classification?
Q2 Do the URT attention heads generate interpretable and meaningful attention scores? Q3 Does the URT layer provide consistent benefits, even when pre-trained backbones are trained in
different ways?
In addition, we investigate architectural choices made, such as our models for keys/queries and their regularization, and study their contribution to achieving strong performance with URT.

heading: 4.1 DATASETS AND SETUP

text:
We test our methods on the large-scale few-shot learning benchmark Meta-Dataset (Triantafillou et al., 2020). It consists of ten datasets with various data distributions across different domains, including natural images (Birds, Fungi, VGG Flower), hand-written characters (Omniglot, Quick Draw), and human created objects (Traffic Signs, Aircraft). Among the ten datasets, eight provide data that can be used during either training, validation and testing (with each class assigned to only one of those sets), while two datasets are solely used for testing. Following Bateni et al. (2020); Requeima et al. (2019), we also report results on MNIST (LeCun et al., 1998), CIFAR10 and CIFAR100 (Krizhevsky et al., 2009) as additional unseen test datasets. Following Triantafillou et al. (2020), few-shot tasks are sampled with varying number of classesN , varying number of shots K and class imbalance. The performance is reported as the average accuracy over 600 sampled tasks. More details of Meta-Dataset can be found in Triantafillou et al. (2020).
The domain-specific backbones are pre-trained following the setup in (Dvornik et al., 2020). Then, we freeze the backbone and train the URT layer for 10,000 episodes, with an initial learning rate of 0.01 and a cosine learning rate scheduler. Following Chen et al. (2020), the training episodes have 50% probability coming from the ImageNet data source. Since different pre-trained backbones may produce representations with different vector norms, we normalize the outputs of the backbones as in Dvornik et al. (2020). URT is trained with parameter weight decay of 1e-5 and with a regularization factor λ = 0.1. The number of heads (H in Equation 7), is set to 2 and the dimension of the keys and
queries (l in Equation 4) is set to 1024. We choose the hyper-parameters based on the performance of the validation set. Details of the hyper-parameter selection and how the performance is influenced by them are outlined in Section 4.5. We find that the bottleneck of training URT is extracting features from CNN. Since we freeze the CNN when training the URT, we find dumping the extracted feature episodes can significantly speed up the training procedure from days to around 2 hours.

heading: 4.2 COMPARISON WITH PREVIOUS APPROACHES

text:
Table 1 presents a comparison of URT with SUR, as well as other baselines based on transfer learning by fine-tuning (Saikia et al., 2020) or meta-learning (Prototypical Networks (Snell et al., 2017), first-order MAML (Finn et al., 2017), ProtoMAML (Triantafillou et al., 2020), CNAPs (Requeima et al., 2019)) and Simple CNAPS(Bateni et al., 2020).
We observe in Table 1 that URT establishes a new state-of-the-art on Meta-Dataset, by achieving the top performance on 8 out of the 10 dataset sources. When comparing to its predecessor, URT outperforms SUR on 4 datasets without compromising performance on others, which is challenging to achieve in the multi-domain setting. Of note, the average inference time for URT is 0.04 second per task, compared to 0.43 for SUR, on a single V100. Thus, getting rid of the optimization procedure for every episode with our meta-trained URT layer also significantly increases the latency, by more than 10×. More results on additional datasets can be found in Appendix A.

heading: 4.3 INTERPRETING AND VISUALIZING ATTENTION BY URT

text:
To better understand how the URT model of Section 4.2 uses its two heads to build adapted representations, we visualize the attention scores produced on the test tasks of Meta-Dataset in Figure 2.
The blue (first head) and orange (second head) heatmaps summarize the values of the attention scores (Equation 5), averaged across several tasks for each test domain. Specifically, the element on row t and column i is the averaged attention scores αi computed on test set domain t for the backbone from domain i. Note that the last two rows are the two unseen domain datasets. We found that for datasets from the seen domains, i.e. the first eight rows, one head (right, orange) consistently puts most of its weight on the backbone pre-trained on the same domain, while the other head (left, blue) learns relatively smoother weight distributions that blend other related domains. For unseen datasets, the right head puts half of its weight on ImageNet and the left head learned to blend the representations from four backbones.

heading: 4.4 URT USING FILM MODULATED BACKBONES

text:
As additional evidence of the benefit of URT on universal representations, we also present experiments based on a different set of backbone architectures. Following SUR (Dvornik et al., 2020), we consider the backbones from a parametric network family, obtained by training a base backbone on one dataset (ILSVRC) and then learning separate FiLM layers (Perez et al., 2018) for each other dataset, to modulate the backbone so it is adapted to the other domains. These backbones collectively have only 0.5% more parameters than a single backbone. More details of the backbones can be found in Appendix C.
A comparison between SUR and URT using these backbones (referred to as SUR-pf and URT-pf) is presented in Table 2. Once again, URT can improve the performance on VGG Flower without sacrificing performance on others.

heading: 4.5 HYPER-PARAMETER AND ABLATION STUDIES

text:
We analyze the importance of the various components of URT’s attention mechanism structure and training strategy in Table 3. First we analyze the importance of using the support set to model queries and/or keys. To this end, we consider setting the matrices Wq / Wk of the query / key linear transformation to 0, which only leaves the bias term. We found that the support set representation is most crucial for building the keys (row w/o Wk in the table) and has minor benefits for queries (row w/o Wq) in the table. This observation is possibly related to the success of attention-based models with learnable constant queries (Liu et al., 2016; Lin et al., 2017). We also found that adding a regularizer Ω(Θ) as in Equation 8 is important for some datasets, specifically VGG Flower and Birds.
An important hyper-parameter in URT is the number of heads H . We chose this hyper-parameter based on the performance on validation set of tasks in Meta-Dataset. In Table 4, we show the validation performance of URT for varying number of heads. As suggested by Triantafillou et al. (2020), we considered looking at the rank of the performance achieved by each choice of H for each validation domains, and taking the average across domains as a validation metric. However,
since the performances when using two to four heads are similar and yield the same average rank, we instead simply consider the average accuracy as the selection criteria.
In general, we observe a large jump in performance when using multiple heads instead of just one. However, since the number of heads controls the capacity, predictably we also observe that having too many heads leads to overfitting.

heading: 5 CONCLUSION

text: We proposed the URT layer to effectively integrate representations from multiple domains and demonstrated improved performance in multi-domain few-shot classification. Notably, our URT approach was able to set a new state-of-the-art on Meta-Dataset, and never performs worse than its predecessor (SUR) while also being 10× more efficient at inference. This work suggests that combining meta-learning with pre-trained universal representations is a promising direction for new few-shot learning methods. Specifically, we hope that future work can investigate the design of richer forms of universal representations that go beyond simply pre-training a single backbone for each domain, and developing meta-learners adapted to those settings.

heading: A EXPERIMENTS ON MORE DATASETS

text: We also report performances on the MNIST, CIFAR-10 and CIFAR-100 dataset sources in Table 5, and compare with the subset of methods that have reported on these datasets. There, URT neither improves nor gets worse performance than SUR, yeilding top performance on the MNIST domain but not on the CIFAR-10/CIFAR-100 domain, on which Simple CNAPS has the best performance.

heading: B MORE RELATED WORKS

text: Transfer by fine-tuning A simple and effective method for few-shot classification is to perform transfer learning by first learning a neural network classifier on all data available for training and using its representation to initialize and then fine-tune neural networks on the few-shot classification tasks found at test time (Chen et al., 2019; Triantafillou et al., 2020; Dhillon et al., 2020; Saikia et al., 2020). Specifically, Saikia et al. (2020) have shown that competitive performance can be reached using a strong hyper-parameter optimization method applied on a carefully designed validation metric appropriate for few-shot learning.

heading: C STRUCTURES AND TRAINING STRATEGY OF BACKBONES

text:
Our URT layer can be built on top of a set of pretrained backbones. The structures of the backbones follow the approach in Dvornik et al. (2020). For Table 1, we use the ResNet18 architecture (He et al., 2016) for the backbones, where a separate backbone is pretrained on each domain separately using the corresponding training data in meta-dataset. The training domains include ImageNet, Omniglot, Aircraft, CU-Birds, Textures, Quick Draw, Fungi and VGG-Flower. For the parametric network family in Table 2, a base ResNet18 is first trained on ImageNet. Then a small number of modulating parameters are trained on the other domains, using their domain-specific training data, while the rest of the base backbone’s weights stay fixed. Specifically, FiLM feature modulation (Perez et al., 2018) is used. This type of parametric network family thus allows for a much reduced number of learnable parameters, by reusing the weights from the base network. In experiments, we use the pretrained backbones released by Dvornik et al. (2020) for both cases, without
any further finetuning. End-to-end training of a set of backbones and the URT layers requires unaffordable computational cost, so we fix the pretrained backbones and only train the URT layer.
Implementation details for training backbones The training details of the backbones come from Dvornik et al. (2020). For optimization, SGD with momentum was used, using cosine learning rate annealing. Since the datasets come from different domains, the starting learning rate, the maximum number of training iterations and annealing frequency are set individually for each dataset. Data augmentation is applied and a constant weight decay of 7×10−4 is set. For each dataset, a grid search over batch size in [8, 16, 32, 64] was run and the one that maximizes accuracy on the validation set was picked. For the parametric network family, the base ResNet18 trained on ImageNet is the same. For other backbones, cosine annealing as learning rate policy is also used, weight decay and data augmentation employed as above. Please refer to Table 4 and Table 5 in the original SUR paper (Dvornik et al., 2020) for the specific values of the hyperparameters for the individual feature networks and the parametric network family, respectively.

heading: D RESULTS ON TRAFFIC SIGNS

text: The shuffle buffer bug described in meta-dataset issue #54 (https://github.com/googleresearch/meta-dataset/issues/54) has been propagated to previous works such as CNAPs (Requeima et al., 2019). The results for the Traffic Signs dataset are considerably worse after fixing this bug. For instance, URT degrades from 69.4±0.8 to 51.1±1.1. Please visit the official meta-dataset GitHub repo for more details. The main paper shows the corrected results (for URT and competing approaches). For completeness, we also provide the tables as they were in the initial version of this paper for your reference in this Appendix. Both sets of results overall support the advantageous performance of URT over previous work.

title: A UNIVERSAL REPRESENTATION TRANSFORMER LAYER FOR FEW-SHOT IMAGE CLASSIFICATION