summaries: The work studies the auxiliary task selection in deep learning to resolve the burden of selecting relevant tasks for pre-training or the multitask learning. By decomposing the auxiliary updates, one can reweight separately the beneficial and harmful directions so that the net contribution to the update of the primary task is always positive. The efficient implementation is experimented in text classification, image classification, and medical imaging transfer tasks.

summaries: Leveraging the power of the data-rich related tasks have been studied (e.g., pre-training and multitask learning). This paper points out that careful utilization of auxiliary task is required to gain enhanced performance in primary tasks. In order to prevent harming the performance of primary tasks, they suggest the method to decompose auxiliary updates into three directions which have positive, negative and neutral impact on the primary task.

summaries: The authors present a general formulation of different settings in multitask learning (including pretraining regimes), in a setting where the goal is to get best performance for a pre-specified primary task and additional auxiliary tasks. The main idea is to divide the gradients on the auxiliary task into 2 subspaces: a subspace where the gradients influence performance of the primary task and a subspace where they only influence the auxiliary task without changing the loss on the primary task. Within the subspace that does have influence on the primary task, it is easy to compute directions that have a positive or negative effect on the primary task, which allows to create different learning schemes given the gradients that point toward: i) auxiliary influence only, ii) positive influence on auxiliary tass, iii) negative influence on primary task. Experimental results show improvements over previously identified meta learning methods on 2 natural language datasets and 3 image datasets.

abstractText: While deep learning has been very beneficial in data-rich settings, tasks with smaller training set often resort to pre-training or multitask learning to leverage data from other tasks. In this case, careful consideration is needed to select tasks and model parameterizations such that updates from the auxiliary tasks actually help the primary task. We seek to alleviate this burden by formulating a modelagnostic framework that performs fine-grained manipulation of the auxiliary task gradients. We propose to decompose auxiliary updates into directions which help, damage or leave the primary task loss unchanged. This allows weighting the update directions differently depending on their impact on the problem of interest. We present a novel and efficient algorithm for that purpose and show its advantage in practice. Our method leverages efficient automatic differentiation procedures and randomized singular value decomposition for scalability. We show that our framework is generic and encompasses some prior work as particular cases. Our approach consistently outperforms strong and widely used baselines when leveraging out-of-distribution data for Text and Image classification tasks.

authors:
Amittai Axelrod
Xiaodong He
Jianfeng Gao

title: Domain adaptation via pseudo in-domain data selection

venue: In EMNLP,

authors:
Atilim Gunes Baydin
Barak A. Pearlmutter
Alexey Andreyevich Radul
Jeffrey Mark Siskind

title: Automatic differentiation in machine learning: a survey

authors:
Mathilde Caron
Piotr Bojanowski
Armand Joulin
Matthijs Douze

title: Deep clustering for unsupervised learning of visual features

venue: In Proceedings of the European Conference on Computer Vision (ECCV),

authors:
Mathilde Caron
Piotr Bojanowski
Julien Mairal
Armand Joulin

title: Unsupervised pre-training of image features on non-curated data

venue: In Proceedings of the IEEE International Conference on Computer Vision,

title: Multitask learning

venue: Machine Learning,

authors:
Arslan Chaudhry
Marc’Aurelio Ranzato
Marcus Rohrbach
Mohamed Elhoseiny

title: Efficient lifelong learning with a-gem

authors:
Zhao Chen
Vijay Badrinarayanan
Chen-Yu Lee
Andrew Rabinovich

title: Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks

venue: In International Conference on Machine Learning,

authors:
Ronan Collobert
Jason Weston

title: A unified architecture for natural language processing: Deep neural networks with multitask learning

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: IEEE conference on computer vision and pattern recognition,

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

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

venue: arXiv preprint arXiv:1810.04805,

authors:
Yunshu Du
Wojciech M. Czarnecki
Siddhant M. Jayakumar
Razvan Pascanu
Balaji Lakshminarayanan

title: Adapting auxiliary losses using gradient similarity

authors:
Mehrdad Farajtabar
Navid Azizan
Alex Mott
Ang Li

title: Orthogonal gradient descent for continual learning

venue: arXiv preprint arXiv:1910.07104,

authors:
Suchin Gururangan
Ana Marasović
Swabha Swayamdipta
Kyle Lo
Iz Beltagy
Doug Downey
Noah A. Smith

title: Don’t stop pretraining: Adapt language models to domains and tasks

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

authors:
Nathan Halko
Per-Gunnar Martinsson
Joel A Tropp

title: Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions

venue: SIAM review,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

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:
Jeremy Irvin
Pranav Rajpurkar
Michael Ko
Yifan Yu
Silviana Ciurea-Ilcus
Chris Chute
Henrik Marklund
Behzad Haghgoo
Robyn Ball
Katie Shpanskaya

title: Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison

authors:
Amit Kumar Jaiswal
Prayag Tiwari
Sachin Kumar
Deepak Gupta
Ashish Khanna
Joel JPC Rodrigues

title: Identifying pneumonia in chest x-rays: A deep learning approach

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Simon Kornblith
Jonathon Shlens
Quoc V. Le

title: Do better imagenet models transfer better

venue: IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

authors:
Ya Le
Xuan Yang

title: Tiny imagenet visual recognition challenge

authors:
Xingyu Lin
Harjatin Baweja
George Kantor
David Held

title: Adaptive auxiliary task weighting for reinforcement learning

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

authors:
David Lopez-Paz
Marc’Aurelio Ranzato

title: Gradient episodic memory for continual learning

authors:
Andrew Maas
Raymond E Daly
Peter T Pham
Dan Huang
Andrew Y Ng
Christopher Potts

title: Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics

venue: Human language technologies,

authors:
Julian McAuley
Christopher Targett
Qinfeng Shi
Anton Van Den Hengel

title: Image-based recommendations on styles and substitutes

venue: In Proceedings of the 38th international ACM SIGIR conference on research and development in information retrieval,

authors:
Ishan Misra
Abhinav Shrivastava
Abhinav Gupta
Martial Hebert

title: Cross-stitch networks for multi-task learning

venue: In Conference on Computer Vision and Pattern Recognition,

authors:
Robert C Moore
William Lewis

title: Intelligent selection of language model training data

venue: In ACL,

title: Accuracy of singular vectors obtained by projection-based svd methods

venue: BIT Numerical Mathematics,

authors:
Hongseok Namkoong
John C Duchi

title: Variance-based regularization with convex objectives

venue: Advances in Neural Information Processing Systems

authors:
Jiquan Ngiam
Daiyi Peng
Vijay Vasudevan
Simon Kornblith
Quoc V. Le
Ruoming Pang

title: Domain adaptive transfer learning with specialist models

authors:
Razvan Pascanu
Tomas Mikolov
Yoshua Bengio

title: On the difficulty of training recurrent neural networks

venue: In International conference on machine learning,

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in pytorch

title: Fast exact multiplication by the hessian

venue: Neural computation,

authors:
Maithra Raghu
Chiyuan Zhang
Jon Kleinberg
Samy Bengio

title: Transfusion: Understanding transfer learning for medical imaging

authors:
Vladimir Rokhlin
Arthur Szlam
Mark Tygert

title: A randomized algorithm for principal component analysis

venue: SIAM Journal on Matrix Analysis and Applications,

authors:
Clemens Rosenbaum
Tim Klinger
Matthew Riemer

title: Routing networks: Adaptive selection of non-linear functions for multi-task learning

venue: arXiv preprint arXiv:1711.01239,

title: An overview of multi-task learning in deep neural networks

authors:
Ozan Sener
Vladlen Koltun

title: Multi-task learning as multi-objective optimization

authors:
Ayan Sinha
Zhao Chen
Vijay Badrinarayanan
Andrew Rabinovich

title: Gradient adversarial training of neural networks

venue: arXiv preprint arXiv:1806.08028,

authors:
Kaitao Song
Xu Tan
Tao Qin
Jianfeng Lu
Tie-Yan Liu

title: Mass: Masked sequence to sequence pre-training for language generation

authors:
Mihai Suteu
Yike Guo

title: Regularizing deep multi-task networks using orthogonal gradients

venue: arXiv preprint arXiv:1912.06844,

authors:
Simon Vandenhende
Stamatios Georgoulis
Marc Proesmans
Dengxin Dai
Luc Van Gool

title: Revisiting multi-task learning in the deep learning era, 2020

authors:
Wei Wang
Ye Tian
Jiquan Ngiam
Yinfei Yang
Isaac Caswell
Zarana Parekh

title: Learning a multi-domain curriculum for neural machine translation

authors:
Xinyi Wang
Hieu Pham
Paul Michel
Antonios Anastasopoulos
Jaime Carbonell
Graham Neubig

title: Optimizing data usage via differentiable rewards

authors:
Zhilin Yang
Zihang Dai
Yiming Yang
Jaime Carbonell
Russ R Salakhutdinov
Quoc V Le

title: Xlnet: Generalized autoregressive pretraining for language understanding

authors:
Tianhe Yu
Saurabh Kumar
Abhishek Gupta
Sergey Levine
Karol Hausman
Chelsea Finn

title: Gradient surgery for multi-task learning

venue: arXiv preprint arXiv:2001.06782,

authors:
Sergey Zagoruyko
Nikos Komodakis

title: Wide residual networks

venue: arXiv preprint arXiv:1605.07146,

authors:
Zhanpeng Zhang
Ping Luo
Chen Change Loy
Xiaoou Tang

title: Facial landmark detection by deep multi-task learning

venue: In European conference on computer vision. Springer,

heading: 1 INTRODUCTION

text:
Multitask learning (Caruana, 1997) and pretraining (Devlin et al., 2018; Caron et al., 2019) have transformed machine learning by allowing downstream tasks with small training sets to benefit from statistical regularities from data-rich related tasks (Collobert & Weston, 2008; Zhang et al., 2014; Liu et al., 2019; Kornblith et al., 2019). Despite these advances, leveraging the mixing of tasks is still an art left to the practitioner. When one is interested in a primary task, it is unclear how to select helpful auxiliary tasks, an appropriate parameter sharing architecture and a good way to filter out auxiliary data which might be detrimental to the primary tasks. Without careful choices, pretraining might hurt end-task performance (Gururangan et al., 2020) or have limited impact (Raghu et al., 2019).
Prior work has examined these problems and proposed solutions, either to choose auxiliary tasks depending on their impact on the primary task (Du et al., 2018; Lin et al., 2019) or to equalize the impact of updates across tasks (Sener & Koltun, 2018; Chen et al., 2018; Hessel et al., 2019). Recently, several approaches (Sinha et al., 2018; Suteu & Guo, 2019; Yu et al., 2020) have been proposed that attempt to minimize interference between the updates across tasks. Our work builds on this direction, but unlike these previous approaches, we do not consider a symmetric view of multi-task learning in the sense that our goal is not to train a model performing well on all tasks. Instead, we focus on improving generalization for a single task, the primary task, and the other tasks, the auxiliary tasks are considered only through their impact on the problem of interest.
For that purpose, we introduce a framework which decomposes the gradient updates from the auxiliary tasks according to their impact on the primary task. We analyze the auxiliary task gradients in the subspace spanned by the primary task per-example gradients. This allows us to decompose auxiliary gradients into into three components : components that help, interfere or have no impact on the primary task according to the Taylor expansion of the expected primary loss. This decomposition allows us to re-weight each component differently prior to the update. Our framework enables us to treat each auxiliary update differently depending on its impact on the task of interest and it
encompasses prior methods such as classical multitask learning (Caruana, 1997) or more novel gradient surgery techniques (Yu et al., 2020). To achieve a tractable approach, we introduce an efficient, robust algorithm (ATTITTUD, Auxiliary Task Training with Influence from Target Task Update Direction) to estimate the subspace spanned by the primary task gradients in an online manner and decompose the auxiliary updates appropriately. As a result, we can integrate our approach with the stochastic training of large neural networks in various contexts.
The contribution of our work is four-fold. To our knowledge, this paper proposes the first approach to adapt auxiliary gradients using a decomposition built from the span of the primary task Jacobian. In order to scale this approach to deep neural nets, we contribute a tractable and efficient algorithm called ATTITTUD that leverages insights from randomized linear algebra and automatic differentiation such as the R-operator (Pearlmutter, 1994). As our third contribution, we show that the fine-grained manipulation of the auxiliary task gradients under ATTITTUD, represents a unified framework that encompasses several previous approaches to asymmetrical task learning as special cases. Finally, we demonstrate the efficacy of our approach in both data-rich and data-starved primary tasks, over both images and textual data.

heading: 2 RELATED WORK

text:
Methods to leverage data outside of the task of interest have been popular in machine learning since the inception of multitask learning (Caruana, 1997; Ruder, 2017; Vandenhende et al., 2020). These methods address multiple task simultaneously and have been successful in various application domains (Collobert & Weston, 2008; Zhang et al., 2014; Misra et al., 2016). The optimization problem induced by multitask learning is difficult and solutions have been proposed for the various difficulties, including dealing with task gradients of different magnitude (Sener & Koltun, 2018; Chen et al., 2018; Hessel et al., 2019), or interfering with each others (Sinha et al., 2018; Suteu & Guo, 2019; Yu et al., 2020). The specific problem of interference has been studied extensively in the context of continual learning. Continual learning visits task in sequence and update interference is particularly problematic as it yields newer tasks to damage previously mastered tasks. In particular, a family of methods to project the gradient of the new tasks to be orthogonal to the gradient of the previous tasks has been proposed (Lopez-Paz & Ranzato, 2017; Chaudhry et al., 2018; Farajtabar et al., 2019).
Different from many previous approaches, we are not interested in addressing multiple tasks per se. In our setting, only the primary task matters and the other auxiliary task have the sole role of improving generalization on the primary task. This is the setting considered by Du et al. (2018); Lin et al. (2019), who favor auxiliary tasks whose gradient directions are helpful to the primary task. Unlike these works that use coarse properties like the cosine similarity between averaged gradients, our approach allows fine-grained gradient manipulation within a subspace. Also, in our case, we do not distinguish between the different auxiliary tasks. Instead, we aim at correcting every auxiliary gradient in the same manner to improve the loss on the primary task. This type of gradient correction is related to Yu et al. (2020), which considers projecting multi-task gradients such that the directions of disagreement are removed. This method is actually a special case of our framework.
Our work also shares some similarities with data selection and domain adaptation approaches. In this case, the training data comes from a single task but its distribution is different from the validation/test distribution (Moore & Lewis, 2010; Axelrod et al., 2011; Ngiam et al., 2018). This classical problem has recently been addressed by sampling training points whose gradient aligns well with the expected validation gradient (Wang et al., 2020b;a). Instead of sampling individual points based on an estimated distribution of how helpful they will be to the primary task, our work avoids the use (and inherent challenges) of this reinforcement learning approach by operating on batch gradients of groups of points.
Our primary task/auxiliary task setting is also related to the pre-training then fine-tuning paradigm in which the auxiliary tasks are visited first (pre-training) to give an initialization for training on the primary task (fine-tuning). These methods have been very successful in settings where primary task data are rare. In particular, it is common to first rely on an unsupervised task over very large datasets prior to fine tuning over a supervised task (Devlin et al., 2018; Liu et al., 2019; Kornblith et al., 2019; Yang et al., 2019; Song et al., 2019; Caron et al., 2018).

heading: 3 AUXILIARY TASK UPDATE DECOMPOSITION

text:
This section introduces a new method to improve generalization on a primary task T ∗ using training data from auxiliary tasks T = {T1, . . . , Tn}, where θ ∈ RD denote the parameters shared by all tasks. Our approach leverages gradient updates from the auxiliary tasks, but unlike the traditional approach, we decompose these gradients to maximize their usefulness to T ∗. Precisely, we decompose the auxiliary task gradients into directions which decrease a first-order approximation of the primary task loss, increase it or have no effect. This decomposition allows weighting these three directions differently when learning from the auxiliary tasks.
In order to decompose the auxiliary gradient, we must collect more fine-grained statistics about the primary task. At each training step, we collect the gradient of the loss with respect to θ for individual examples from the primary task, {∇θLprimi ,∀i }. The span of these vectors,
S = Span{∇θLprimi ,∀i}
defines a subspace in which any linear combination of primary task gradients lies, including the gradient of the expected primary task loss, i.e. gprim = E(∇θLprimi ) ∈ S. We denote the size of the subspace, |S| = K. This is upper-bounded by the number of examples m, used to construct S. If we define the orthogonal complement of S as S⊥, any vector v ∈ S⊥, is therefore orthogonal to gprim, i.e. v · gprim = 0. This means that adding such a vector to the parameters has no impact on the expected primary task loss, according the order-1 Taylor expansion of Lprim, i.e.
Lprim(θ + v) ' Lprim(θ) + v · gprim = Lprim(θ).
We propose to project auxiliary task gradients onto S and S⊥. This allow to distinguish between the directions of the auxiliary task updates which impact the primary task loss and those which do not. If we denote the averaged auxiliary task gradient as gaux = E(∇θLauxi ), we can decompose this gradient as gaux = g±aux + g ⊥ aux. where g ± aux ∈ S is the portion of the gradient that lies in the span of the primary task example gradients and g⊥aux ∈ S⊥ is the portion that lies outside of it. Since g⊥aux ∈ S⊥, it is orthogonal to the average primary task gradient and parameter updates along the direction of g⊥aux are expected to have limited impact on the primary task loss. On the other hand, updates along the direction of g±aux can potentially improve or damage the averaged primary task loss. This component deserves a more careful treatment.
For that purpose, we introduce {ui, i = 1, . . . ,K} an orthonormal basis of S. In this basis, we can measure if the components of g±aux agree or disagree with gprim. We say that the two gradients agree along ui iif sign(g±aux · ui) = sign(gprim · ui). This means that we can decompose g±aux = g+aux + g−aux
where g+aux refers to the projection of g ± aux onto the basis vectors where g ± aux and gprim agree. By this decomposition, g+aux helps the primary task, g + aux · gprim > 0, while g−aux interfere with the primary task, g−aux · gprim < 0. Guided by the primary task, we can therefore decompose the auxiliary task gradient as
gaux = g ⊥ aux + g + aux + g − aux (1)
which is described on Fig 1. Our approach proposes to re-weight differently the components of gaux, i.e.
g̃aux = η⊥g ⊥ aux + η+g + aux + η−g − aux (2) where ηaux = ( η⊥, η+, η− ) are hyper-parameters adjusting the auxiliary gradient according to the impact on the main task. If we also wish to include the primary task gradient in descent, as with multitasking, we can introduce ηprim as a scalar control variable to modulate its weighting.
A consequence of introducing ηaux is that specific configurations lead us to gradient updates that are guaranteed to do no harm to both tasks. This is captured by Theorem 1 below. Theorem 1. Let Laux(θt) and Lprim(θt) represent the full batch losses of the auxiliary tasks and primary task respectively at step t. We assume the gradients of Laux and Lprim are Lipschitz continuous with constant L > 0. Following the update rule : θt+1 = θt − α · g̃aux, where α ≤ 1L is the learning rate, we are guaranteed :
Laux(θt+1) ≤ Laux(θt) Lprim(θt+1) ≤ Lprim(θt)

heading: If η− = 0 and η⊥, η+ ≥ 0

text:
Proof. See Appendix A
This theorem focuses on a single update and guarantees progress on both auxiliary and primary tasks. However, our asymmetric scenario is not interested in improving the auxiliary tasks per se and is amenable to more aggressive settings. Ideally we want gradient updates during pre-training with T to not only do-no-harm to T ∗ when applied downstream but also to be along descent directions that are maximally beneficial to T ∗. We can consider η− < 0 as in Fig 1. Reversing the direction of g−aux by setting η− < 0 preserves the descent guarantee on Lprim(θt+1) but no longer ensures descent on Laux(θt+1). There are other interesting settings for our control parameters. One can recover the original gradient gaux with η⊥ = η− = η+ = 1.0. One can choose to drop gradients orthogonal to the primary task gradient span with η⊥ = 0.0, or ignore those which conflict with the main task by setting η− = 0.0.
Relationships to other approaches Our framework is generic and encompasses other approaches as a particular case. One can train solely on the primary task by selecting ηaux =( 0.0, 0.0, 0.0) and ηprim = 1.0. Classical multitasking corresponds to ηaux = ( 1.0, 1.0, 1.0) and ηprim > 0.0, while classical pre-training corresponds to performing a first phase with ηaux = ( 1.0, 1.0, 1.0) and ηprim = 0.0. Interestingly, our formulation introduces novel variants of pre-training, for instance, one can consider pre-training with only auxiliary gradients helpful to the primary task, ηaux = ( 0.0, 1.0, 0.0) and ηprim = 0.0, followed by fine-tuning with
ηaux = ( 0.0, 0.0, 0.0) and ηprim = 1.0.
Our approach also instantiates PCGrad (Yu et al., 2020) as a particular case. This method was introduced to address the issue of conflicting gradients in multitask settings. PCGrad orthogonalizes the gradients of each task and removes conflicting gradients. To recover PCGrad under our approach, note that it is equivalent to a specific choice of our decomposition in the 1-D subspace spanned by the gprim. PCGrad then removes components of gaux that conflict with gprim which is equivalent to ηaux = ( αaux, αaux, 0.0 ) and ηaux = αprim.

heading: 4 IMPLEMENTATION

text:
Equation 2 requires selecting a basis for the span of primary task gradients. Multiple choices are possible to define the basis {ui}, to represent the span at each optimization time-step. This choice
is important since the components of g±aux are labeled positive or negative depending on how they agree with the projection of the averaged primary task gradient onto the same basis. A natural choice is to select the basis as the singular vectors of the matrix of primary task per-example gradients J∗ ∈ Rm×D, also know as the Jacobian. To improve efficiency and prevent over-fitting on a few examples, we consider the span defined by the, k < |S|, largest principal vectors of J∗. Using the principal vectors as directions of descent instead of the mean induces a more robust algorithm since the mini-batch average gradient is susceptible to outliers and skew from replicated data-points. To the best of our knowledge, we are the first to propose using the singular vectors of J∗ as directions of descent. We leave the theoretical implications of this algorithm to future work but note that its variance reduction properties may induce generalization benefits (Namkoong & Duchi, 2017).
We also consider alternative choices of bases as baselines, including the canonical parameter basis. This choice will examine the sign of every parameter update to verify whether it agrees with gprim. Whilst Theorem 1 holds irrespective of the choice of basis, its proof reveals that the amount of progress made on each loss depends on the choice of basis. Specifically, the reduction in Lprim(θt+1),Laux(θt+1) after a gradient step along g̃aux is proportional to the fraction of the norms of gprim and gaux captured by the subspace spanned by our choice of basis. To justify our use of the top singular values of J∗, we evaluate this fraction for different choice of basis in our experiments (see Appendix C).
We are interested in applying our approach to the training of large neural networks and must consider a scalable algorithmic solution. As stochastic optimization is prevalent in this setting, we construct subspace S from a mini-batch of primary task data. Similarly, the expected gradients gprim and gaux are defined over a mini-batch. Instead of computing the singular value decomposition (SVD) of {∇θLprimi ,∀i} exactly, we rely on a randomized approximation (Halko et al., 2011; Rokhlin et al., 2010; Nakatsukasa, 2017). This method does not require instantiating the vectors {∇θLprimi ,∀i} and only needs a low dimensional projection onto a random subspace. This is advantageous for high dimensional cases, i.e. when the number of model parameters is large. In our case, this method also allows us to benefit from memory-efficient computation of Jacobian Vector product using the R-operator (Pearlmutter, 1994) offered by automatic differentiation packages (Baydin et al., 2015) like Pytorch (Paszke et al., 2017). This means that we can compute SVD with a limited computational and memory burden, albeit without sacrificing approximation accuracy (Nakatsukasa, 2017). Additionally, we do not recompute the basis at every optimization step but at every n steps, which is efficient when training with small updates, e.g. when small learning rates and gradient clipping are used (Pascanu et al., 2013) (see Appendix C for more details about n).
We study the impact of these choices in practice in Section 6. Putting it all together results in the ATTITTUD algorithm, Auxiliary Task Training with Influence from Target Task Update Direction, shown as Algorithm 1. The sub-procedure randomized lowrank approx is detailed in Appendix B as Algorithm 2 .
Algorithm 1: ATTITTUD : Construct Auxiliary Task Surrogate Gradient Require : gaux, J∗ : Auxiliary task average gradient, primary task Jacobian Require : ηaux = ( η⊥, η+, η− ) : Auxiliary task control parameters Require : k : Size of subspace gprim = 1 m ∑m i=1 J ∗ i,:
V ← randomized lowrank approx(J∗, k) pprim, paux = Vt ( gprim )T , Vt ( gaux )T // ◦ is the hadamard product operator
p+aux,p − aux =
( 1[ pprim◦paux ≥ 0 ]) ◦ paux, (1[ pprim◦paux < 0 ]) ◦ paux
// Calculate the decomposition components
g+aux, g − aux = ( p+aux )T V , ( p−aux )T V // Calculate the out of span component
g⊥aux = gaux − ( g+aux + g − aux ) g̃aux = ( η⊥ · g⊥aux ) + ( η+ · g+aux ) + ( η− · g−aux
) Return : g̃aux : Auxiliary task surrogate gradient

heading: 5 EXPERIMENTAL SETUP

text:
We compare ATTITTUD with previous methods on a variety of tasks and domains. We rely on both text and image classification tasks to conduct our analysis. We also present ablation experiments to explain the impact of hyper-parameter selection. We make code for ATTITTUD and related experiments available on github. 1
Text Classification. We apply our method on binary sentiment classification. We consider the Amazon Helpfulness (McAuley et al., 2015) and Imdb Movie Review (Maas et al., 2011) tasks. The Amazon Helpfulness task splits text reviews into 115k/5k/25k documents for train-validation-test split whilst the Imdb Review dataset has a 20k/5k/25k split. The Imdb Review task also has 50k unlabeled reviews as extra data which we utilize.
For our models we build on top of Gururangan et al. (2020)’s work where they introduce TaskAdaptive Pre-training (TAPT). TAPT further pre-trains a generic model, Roberta (Liu et al., 2019), by performing Masked Language Modelling, MLM, (Devlin et al., 2018) on the task specific data (ignoring the labels) before doing supervised learning with the same data. We replicate Gururangan et al. (2020)’s experimental setup and re-use their hyper-parameters for our experiments. We use the TAPT task as our auxiliary task. We extend TAPT to use our method by modifying the TAPT gradient with guidance from the supervised-learning task gradients. As baselines, we compare against TAPT and cross-TAPT: where we swap the masked language modelling pre-training data for the two tasks. Cross-TAPT is a setting where one uses out-of-distribution data for pre-training.
Image Classification. We apply our method to both high-resource and limited-data image classification tasks. We use the Cifar100 dataset (Krizhevsky et al., 2009) to explore the high-resource setting. We follow Rosenbaum et al. (2017) and treat each of the 20 super-classes / coarse labels of Cifar100 as a separate task. In our asymmetrical task setting, each of the 20 tasks is treated as a primary task, whilst the remaining 95 classes are grouped into a single auxiliary task. Thus, for each coarse label, we have an auxiliary 95-way classification task and a 5-way primary classification task. Moving forward, we refer to this setting as MultiCifar100.
We use a down-sampled version of Cifar10 (Krizhevsky et al., 2009) as a low-resource setting. Specifically, we rely on Cat-vs-Dog for the primary task and use the remaining 8 classes for the auxiliary task. Our auxiliary task is therefore an 8-way classification task where each class has 5,000 examples. We restrict the Cat and Dog classes to only 50 training examples from each class. We use the low-resource setting to compare against other methods and for our ablation study.
For these vision experiments, we use a WideResNet-22 architecture (Zagoruyko & Komodakis, 2016) with a depth of k = 4. We compare our method to 4 different baselines : no pre-training, vanilla pre-training, multitasking and PCGrad (Yu et al., 2020). Our architecture is more standard and allows gradient descent optimization unlike the routing network of Rosenbaum et al. (2017) and (Yu et al., 2020), which requires reinforcement learning for training.
Medical Imaging Transfer. We apply our method to cross-domain transfer for low-resource medical image classification. Specifically, we use 5k training examples from the ChexPert Dataset (Irvin et al., 2019) as our primary task and seek to identify 5 different thoracic pathologies: atelectasis, cardiomegaly, consolidation, edema and pleural effusion. This setup has been used in several crossdomain pretraining studies (Raghu et al., 2019; Jaiswal et al., 2019). Note that since we do not have access to the test set for this task, we use the validation set (231 images) as a proxy test set, and sample 100 images from the training data as a new validation set. We rely on generic photographs (Imagenet) as an auxiliary task (Deng et al., 2009). We use Tiny Imagenet Dataset (Le & Yang, 2015), a subset of Imagenet which consists of 500 examples each from 200 classes, instead of training on full Imagenet. All approaches are applied to the Resnet18 model (He et al., 2016) trained with Adam (Kingma & Ba, 2014).
For all our experiments, we select the auxiliary task control parameters ηaux within {(1.0, 1.0,−1.0), (1.0, 1.0, 0.0), (1.0, 0.0,−1.0), (1.0, 0.0, 0.0)} for ease of interpretability. For settings where we compare against multi-tasking, we select ηprim within a small subset of the settings that worked best with multitasking baseline experiments. These choices limit the overhead of hyper-parameter search but still allow us to show the empirical advantage of our method. In
1Code available here https://github.com/ldery/ATTITTUD
all our experiments, we provide all methods with similar hyper-parameter search budgets, e.g. for Cifar10-Cat-vs-Dog, we ran a grid search with 16 configurations for regular pretraining, 16 configurations for PCGrad and 12 configurations for ATTITUD. More experimental details are available in Appendix C

heading: 6 RESULTS AND DISCUSSION

text:
Text Classification. Table 1 shows the results for text classification. When the same data is used both for the auxiliary task of MLM and the primary classification task, TAPT and ATTITTUD both bring a similar improvement over Roberta (Imdb, Amazon columns). For the Cross-TAPT setting where different data is used for the auxiliary task and the primary task (Imdb + Amazon MLM, Amazon + Imdb MLM columns), TAPT does not perform as well as ATTITTUD. This highlights the advantage of ATTITTUD when the auxiliary task data distribution differ from the primary task distribution.
Image Classification. Our results are presented in Table 2. Both for MultiCifar100 (high resource setting) and Cifar10-Cat-vs-Dog (low resource setting), ATTITUD shows a strong improvement over baselines. In general, we find that primary-task aware pre-training (Multitasking, PCGrad, Ours) is better than vanilla pre-training which also performs better than having no pre-training at all. For MultiCifar100, we find that using ηaux = (1.0, 1.0,−1.0),ηprim = 0.1 worked best for 11 out of the 20 Cifar100 super-classes tasks. Note that ηaux = (1.0, 1.0,−1.0) is an aggressive but novel configuration we introduce. Multitask learning and PCGrad produce better models on 6 and 3 tasks respectively. In the low-resource Cat-vs-Dog, setting ATTITUD produces a bigger boost in performance compared to baselines, with the best performing configuration being ηaux = (1.0, 0.0, 0.0),ηprim = 0.01. We posit that this configuration is successful because removal of the in-span components makes overfitting less likely. Applying the out-of-span components means the model learns features that do not harm the loss of the current mini-batch but could be useful later. Note that our best performing configurations are all novel and never an instantiation of PCGrad.
Medical Imaging Transfer. Table 3 shows our results on the ChexPert multi-label classification task. Per-pathology breakdowns are in Appendix C. Doing no pre-training at all performs worst.
Our method outperforms using a pre-trained Resnet18 model over Imagenet. We apply the endtask-aware ATTITUD over 100k ImageNet images after the initial pretraining and we reach 83.3% AUC, an improvement over 81.4%.
Ablation Study. Our approach relies on the top-k singular vectors from randomized svd to define the basis to identify the positive and negative component of the auxiliary task gradient, see Section 4. This method is more accurate than several alternatives; see Table 4. Namely, we compare our choice to random, the basis spanned by k randomly chosen orthogonal vectors in RD, unit avg grad, the basis spanned by the average primary task gradient, and canonical, the per-parameter basis. This ablation was performed under a more limited tuning budget (we cross-validated on configurations (1, 1, 0) and (1, 1,−1) only) than the full Cat-vs-Dog experiments from Table 2.
We also examine the number of samples to estimate the principal directions of the per-example primary task gradient. Larger sample sizes involve more computation but have limited benefit on average accuracy. Large sample sizes however reduce variance, as shown in Figure 2 (left). This is as expected since using more samples gives a higher fidelity estimate of the top-k singular vectors.
Another parameter of our algorithm is the size of our subspace, k. In general, we observe that in low-resource settings, it is better to operate on the auxiliary task gradient in a smaller dimensional subspace. The opposite holds for high-resource settings. This can be seen in Figure 2 (right). Whilst using a larger dimensional subspace captures a richer description of the J∗, it also creates the risk of over-fitting especially in a limited data setting. This trade-off therefore has to be validated on a per-task basis.

heading: 7 CONCLUSIONS

text:
In this work, we propose a new approach to training a model with additional help from an auxiliary task. Our method decomposes the gradients of the auxiliary task according to three directions, with positive, negative and neutral impact on the primary task. This decomposition allows a flexible re-weighting of the auxiliary task components and give rise to a family of training strategies,
which encompasses novel and existing approaches. We leverage insights from randomized linear algebra and automatic differentiation to scale the approach to large deep networks. Experiments in multitasking, pretraining and domain transfer over vision and text classification task demonstrate the empirical benefit of our framework.

heading: A PROOF OF THEOREM 1

text:
Theorem. Let Laux(θt) and Lprim(θt) represent the full batch losses of the auxiliary tasks and primary task respectively at step t. We assume the gradients ofLaux andLprim are Lipschitz continuous with constant L > 0. Following the update rule : θt+1 = θt −α · g̃aux, where α ≤ 1L is the learning rate, we are guaranteed :
Laux(θt+1) ≤ Laux(θt) Lprim(θt+1) ≤ Lprim(θt)

heading: If η− = 0 and η⊥, η+ ≥ 0

text:
Proof. Let Vt ∈ RK×D be the orthonormal matrix whose rows span the per-example primary task gradients J∗ at timestep t. The projections of the average primary task gradient gprim = 1m ∑m i=1 J ∗ i,: and average auxiliary task gradient gaux at iteration t are
pprim = Vt ( gprim )T paux = Vt ( gaux
)T pprim and paux will agree on some directions (same sign on those components). We use the operator [x]+ to mark these directions of agreement. This operator preserves components that agree and sets those that disagree to zero. As an example given pprim = [1, 1,−1] and paux = [1, 3, 10], [pprim]+ = [1, 1, 0] and [paux]+ = [1, 3, 0]. For directions that disagree (different signs of the respective components), we introduce the operator [x]−. In the above example [pprim]− = [0, 0,−1] and [paux]− = [0, 0, 10]. Note that our operators are defined by comparing two vectors x1 and x2, Our operators have the following properties by definition :
x = [x]− + [x]+
and [x]+ ⊥ [x]−, [x1]± ⊥ [x2]∓
From Equation 2 : g̃aux = η+g + aux + η−g − aux + η⊥g ⊥ aux We can re-write this in terms of [x]± as :
g̃aux = η+[paux]+ + η−[paux]− + η⊥ ( gaux − paux ) We now proceed to show the effect of the gradient descent update below on Laux(θt+1) and Lprim(θt+1).
θt+1 = θt − α · g̃aux (3)
How does this update affect the loss on the primary task loss Lprim(θt+1)? Lprim(θt+1) = Laux(θt − α · g̃aux)
≈ Lprim(θt)− α ( g̃aux )T gprim (First order Taylor Expansion)
= Lprim(θt)− α ( η+[paux]+ + η−[paux]− + η⊥g ⊥ aux )T gprim
= Lprim(θt)− α ( η+[paux]+ + η−[paux]− + η⊥g ⊥ aux )T( [pprim]+ + [pprim]− ) = Lprim(θt)− α ( η+ ( [paux] T +[pprim]+ + [paux] T +[pprim]− ) + η− ( [paux] T −[pprim]+ + [paux] T −[pprim]−
)) = Lprim(θt)− α ( η+[paux] T +[pprim]+ + η−[paux] T −[pprim]−
) ≤ Lprim(θt) (if η− ≤ 0, η⊥, η+ ≥ 0)
Note that in going from line 3 to 4 in the proof above, we use the fact that ( g⊥aux )T gprim = 0 since g⊥aux lies outside the subspace and gprim lies inside it. For the last step of the proof, we use the observations below :
[paux]+[pprim]+ ≥ 0 since these directions agree in sign [paux]−[pprim]− ≤ 0 since these directions disagree in sign [paux]+[pprim]− = 0 by the property of the [x]± operator [paux]−[pprim]+ = 0 same motivation as above
How does Equation 3 affect the auxiliary task loss Laux(θt+1)?
Laux(θt+1) = Laux(θt − α · g̃aux) ≈ Laux(θt)− α ( g̃aux )T gaux (First order Taylor Expansion)
= Laux(θt)− α ( η⊥g ⊥ aux + η+g + aux + η−g − aux )T ( g⊥aux + g + aux + g − aux ) = Laux(θt)− α ( η⊥‖g⊥aux‖2 + η+‖g+aux‖2 + η−‖g−aux‖2 ) (Cross terms cancel due to orthogonality) ≤ Laux(θt) (If η−, η⊥, η+ ≥ 0)
Thus, choosing η− = 0 ensures that we are minimizing both Laux(θt) and Lprim(θt). We can combine this with the constraint on α ≤ 1L to derive convergence guarantees after some T steps as in optimization literature.

heading: B RANDOMIZED MATRIX THEORY

text:
Algorithm 2: randomized lowrank approx : Construct low rank approximation Require : J ∈ Rm×D : Input Matrix Require : k : Rank of subspace
Π ∼ N (0, I) ∈ Rk×m C = ΠJ V ← Gram Schmidt(C)
Return : V ∈ Rk×D : Low rank approximation of J
The Gram Schmidt procedure orthogonalizes the rows of an input matrix.

heading: C MORE EXPERIMENTAL DETAILS

text:
Image Classification For MultiCifar100, unlike Rosenbaum et al. (2017); Yu et al. (2020) who use a 500-100 train-test split for examples under each fine-grained CIFAR 100 label, we include a validation set and therefore opt for a 400-100-100 train-validation-test split. We test on all 1000 test examples per class.
For Cat-vs-Dog, we use 100 examples from the training set as validation and test on all 1000 test examples per-class.
For Image Classification experiments, we perform pre-training with a learning rate of 1e-4 for all experiments and finetuning learning rate of 5e-4. These values were selected after coarse hyperparameter search. In both pre-training and finetuning settings, we decay the learning rate by 0.5 if the validation loss has not improved over 4 epochs, up till a minimum learning rate of 1e-5. we use the Adam Optimizer (Kingma & Ba, 2014) with β = (0.9, 0.999). We clip all gradient norms to 1.0 before performing gradient descent. We cross-validated dropout rates within the set {0.05, 0.1, 0.2, 0.3} for both pre-training and finetuning steps. We cross validate ηprim based on the relative sizes of primary and auxilary task datasets. All experiments are averaged over 5 random seeds. For all our Vision experiments, we either recompute our subspace basis every n = 5 or n = 10 iterations. We find that n is not as important as the other hyper-parameters, with the two choices showing similar
performance when the other hyper-parameters (learning rate and gradient norm clipping) are fixed to reasonable values.
Due to the fact that Yue et al (PCGrad) treat all tasks symmetrically, which is different from our primary-auxiliary setting, we introduced an extra parameter, αprim, for PCGrad to account for weighting the primary task. We cross validated values of αprim ∈ {0.1, 0.05, 0.01, 0.001}. Medical Imaging Transfer Table 4 presents a more detailed breakdown of the ChexPert task. For 50k examples from Imagenet, our best performing configuration was ηaux = (1.0, 0.0,−1.0). We did not use the primary task gradient directly for pre-training so ηprim = 0.0 for all cases. For ATTITUD, we use the same learning rates as in the Image classification setup above. For the No-Pretraining and Vanilla pretraining we cross-validated the learning rates for both finetuning and pre-training from the set {1e-3, 1e-4}. We cross-validated the same list of dropout values above.
Text Classification For our NLP experiments, we tried limiting the number of layers we applied ATTITUD to. We achieved good performance without applying ATTITUD to the word embedding layers (these were updated with untouched auxiliary task gradients). We cross-validated ηprim = {0.01, 0.05, 0.0025}. For all our NLP experiments, we either recompute our subspace basis every n = 1 or n = 4 times
For all experiments involving ATTITUD, We cross-validate the following choices of the subspace size k ∈ {5, 10, 20} from J∗ ∈ Rm×D using m ∈ {32, 64}. We recompute the subspace every 10 steps for vision experiments and every 4 steps for NLP experiments. We run all experiments for a maximum of 150 pretraining epochs and 500 finetuning epochs. We performed early stopping for all experiments if no improvement after 10 consecutive epochs.
Ablation of Fraction of Norm within Subspace The left pane of Figure 3 reinforces our intuition and confirms that our choice of the top-k singular vectors (randomized svd) gives the best accuracy as averaged across 5 seeds. random is the basis spanned by k randomly chosen orthogonal vectors in RD, unit avg grad is the basis spanned by the average primary task gradient whilst canonical uses the per-parameter basis. Note that k = 5 for random and randomized svd whilst for unit avg grad and canonical, k = 1 and k = D respectively. We use the fraction of the norm of sample gradients within a subspace as indicators of how semantically meaningful that choice of subspace is. We expect that a semantically meaningful choice of basis will achieve better generalization performance because it captures the essential parts of the gradient with k D . canonical trivially captures all the norm of the sampled gradient vectors but because k = D, it generalizes poorly. Notice that only small fractions of the norms of sample primary and auxiliary task average gradients lie in the subspace for random and unit avg grad, whilst significant fractions lie in randomized svd.


summaries: The authors use a 3D world to explore grounded language learning, in which an agent uses RL to combine novel word-learning with stably acquired meanings to successfully identify and manipulate objects.  They show that a novel, psychologically-inspired memory mechanism is more memory-efficient than Transformers (both of which outperform plain LSTMs) and that it exhibits surprisingly robust generalization to novel action-object pairs.  The results should be of interest to many working in grounded language / multimodal representation learning, and the experiments are thorough and well-motivated.

summaries: This paper presents experiments for acquiring words via fast-mapping in an embodied environment. The technical contribution is interesting and solid, but the experiments fail to address some important questions that are yet scoped by the claims of the paper (namely, that learning is being done -both- fast and slow, as per the title). Notably, the paper is really well-written and readable, and the experiments on novel category + novel instance recognition are really convincing specifically for fast-mapping (4.1).

summaries: An agent following instructions in a grounded world is a core task in AI. This paper studies agent that accomplish this using memory-based architecture. This paper presents an argument for a multi-modal memory-architecture called DCEM whose key/queries and values are dependent on language and vision modalities respectively (or vice versa). An argument is made that this will be helpful for generalizing to novel language at test-time. Results are presented in a simple 3D domain containing several objects randomly sampled each time from a set of 30 objects. Task contain two types of instructions: "pick up an object" and "place an object on another object". Interaction proceeds in episodes where each episode contains a discovery phase where the agent learns the phrase associated with each object, and an instruction phase where the agent solves a given instruction. The proposed DCEM model outperforms baselines on various metrics and ablation. Importantly, it is shown that the DCEM can generalize to novel object names.

abstractText: Recent work has shown that large text-based neural language models acquire a surprising propensity for one-shot learning. Here, we show that an agent situated in a simulated 3D world, and endowed with a novel dual-coding external memory, can exhibit similar one-shot word learning when trained with conventional RL algorithms. After a single introduction to a novel object via visual perception and language (“This is a dax”), the agent can manipulate the object as instructed (“Put the dax on the bed”), combining short-term, within-episode knowledge of the nonsense word with long-term lexical and motor knowledge. We find that, under certain training conditions and with a particular memory writing mechanism, the agent’s one-shot word-object binding generalizes to novel exemplars within the same ShapeNet category, and is effective in settings with unfamiliar numbers of objects. We further show how dual-coding memory can be exploited as a signal for intrinsic motivation, stimulating the agent to seek names for objects that may be useful later. Together, the results demonstrate that deep neural networks can exploit meta-learning, episodic memory and an explicitly multi-modal environment to account for fast-mapping, a fundamental pillar of human cognitive development and a potentially transformative capacity for artificial agents.

authors:
Adrià Puigdomènech Badia
P. Sprechmann
Alex Vitvitskyi
Daniel Guo
B. Piot
Steven Kapturowski
O. Tieleman
Martı́n Arjovsky
A. Pritzel
Andew Bolt
Charles Blundell

title: Never give up: Learning directed exploration strategies

authors:
Sven Bambach
David Crandall
Linda Smith
Chen Yu

title: Toddler-inspired visual object learning

venue: In Advances in neural information processing systems,

authors:
Douglas A Behrend
Jason Scofield
Erica E Kleinknecht

title: Beyond fast mapping: Young children’s extensions of novel words and novel facts

venue: Developmental Psychology,

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:
Susan Carey
Elsa Bartlett

title: Acquiring a single new word

venue: Papers and Reports on Child Language Development,

authors:
Angel X Chang
Thomas Funkhouser
Leonidas Guibas
Pat Hanrahan
Qixing Huang
Zimo Li
Silvio Savarese
Manolis Savva
Shuran Song
Hao Su

title: Shapenet: An information-rich 3D model repository

venue: arXiv preprint arXiv:1512.03012,

authors:
Devendra Singh Chaplot
Kanthashree Mysore Sathyendra
Rama Kumar Pasumarthi
Dheeraj Rajagopal
Ruslan Salakhutdinov

title: Gated-attention architectures for task-oriented language grounding

venue: In Thirty-Second AAAI Conference on Artificial Intelligence,

authors:
Zihang Dai
Zhilin Yang
Yiming Yang
Jaime G Carbonell
Quoc Le
Ruslan Salakhutdinov

title: Transformer-xl: Attentive language models beyond a fixed-length context

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

authors:
Lasse Espeholt
Hubert Soyer
Remi Munos
Karen Simonyan
Volodymir Mnih
Tom Ward
Yotam Doron
Vlad Firoiu
Tim Harley
Iain Dunning

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

venue: arXiv preprint arXiv:1802.01561,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

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

venue: arXiv preprint arXiv:1703.03400,

authors:
Meire Fortunato
Melissa Tan
Ryan Faulkner
Steven Hansen
Adrià Puigdomènech Badia
Gavin Buttimore
Charles Deck
Joel Z Leibo
Charles Blundell

title: Generalization of reinforcement learners with working and episodic memory

venue: In Advances in Neural Information Processing Systems,

authors:
Alex Graves
Greg Wayne
Malcolm Reynolds
Tim Harley
Ivo Danihelka
Agnieszka GrabskaBarwińska
Sergio Gómez Colmenarejo
Edward Grefenstette
Tiago Ramalho
John Agapiou

title: Hybrid computing using a neural network with dynamic external

venue: memory. Nature,

authors:
Karl Moritz Hermann
Felix Hill
Simon Green
Fumin Wang
Ryan Faulkner
Hubert Soyer
David Szepesvari
Wojciech Marian Czarnecki
Max Jaderberg
Denis Teplyashin

title: Grounded language learning in a simulated 3D world

venue: arXiv preprint arXiv:1706.06551,

authors:
Felix Hill
Stephen Clark
Karl Moritz Hermann
Phil Blunsom

title: Understanding early word learning in situated artificial agents

venue: Proceedings of CogSci,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: Long short-term memory

authors:
Timothy P Lillicrap
Jonathan J Hunt
Alexander Pritzel
Nicolas Heess
Tom Erez
Yuval Tassa
David Silver
Daan Wierstra

title: Continuous control with deep reinforcement learning

venue: In ICLR (Poster),

authors:
James L McClelland
Bruce L McNaughton
Randall C O’Reilly

title: Why there are complementary learning systems in the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory

venue: Psychological review,

authors:
James L. McClelland
Felix Hill
Maja Rudolph
Jason Baldridge
Hinrich Schütze

title: Extending machine language models toward human-level language understanding

venue: PNAS (to appear),

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

authors:
A. Emin Orhan
Vaibhav V. Gupta
Brenden M. Lake

title: Self-supervised learning through the eyes of a child, 2020

title: Mental imagery in associative learning and memory

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
Matthew M. Botvinick
Nicolas Heess
Raia Hadsell

title: Stabilizing transformers for reinforcement learning, 2019

authors:
Adam Santoro
Sergey Bartunov
Matthew Botvinick
Daan Wierstra
Timothy Lillicrap

title: Metalearning with memory-augmented neural networks

authors:
Jake Snell
Kevin Swersky
Richard Zemel

title: Prototypical networks for few-shot learning

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
Timothy Lillicrap
koray kavukcuoglu
Daan Wierstra

title: Matching networks for one shot learning

authors:
Haley Vlach
Catherine M Sandhofer

title: Fast mapping across time: Memory processes support children’s retention of learned words

venue: Frontiers in psychology,

authors:
Xin Wang
Qiuyuan Huang
Asli Celikyilmaz
Jianfeng Gao
Dinghan Shen
Yuan-Fang Wang
William Yang Wang
Lei Zhang

title: Reinforced cross-modal matching and self-supervised imitation learning for vision-language navigation

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

authors:
Sandra R Waxman
Amy E Booth

title: Principles that are invoked in the acquisition of words, but not facts

venue: Cognition, 77(2):B33–B43,

authors:
Greg Wayne
Chia-Chun Hung
David Amos
Mehdi Mirza
Arun Ahuja
Agnieszka GrabskaBarwinska
Jack Rae
Piotr Mirowski
Joel Z Leibo
Adam Santoro

title: Unsupervised predictive memory in a goal-directed agent

venue: arXiv preprint arXiv:1803.10760,

title: Meta-learning for few-shot natural language processing: A survey

venue: arXiv preprint arXiv:2007.09604,

heading: 1 INTRODUCTION

text:
Language models that exhibit one- or few-shot learning are of growing interest in machine learning applications because they can adapt their knowledge to new information (Brown et al., 2020; Yin, 2020). One-shot language learning in the physical world is also of interest to developmental psychologists; fast-mapping, the ability to bind a new word to an unfamiliar object after a single exposure, is a much studied facet of child language learning (Carey & Bartlett, 1978). Our goal is to enable an embodied learning system to perform fast-mapping, and we take a step towards this goal by developing an embodied agent situated in a 3D game environment that can learn the names of entirely unfamiliar objects in a single exposure, and immediately apply this knowledge to carry out instructions based on those objects. The agent observes the world via active perception of raw pixels, and learns to respond to linguistic stimuli by executing sequences of motor actions. It is trained by a combination of conventional RL and predictive (semi-supervised) learning.
We find that an agent architecture consisting of standard neural network components is sufficient to follow language instructions whose meaning is preserved across episodes. However, learning to fast-map novel names to novel objects in a single episode relies on semi-supervised prediction mechanisms and a novel form of external memory, inspired by the dual-coding theory of knowledge representation (Paivio, 1969). With these components, an agent can exhibit both slow word learning and fast-mapping. Moreover, the agent exhibits an emergent propensity to integrate both fast-mapped and slowly acquired word meanings in a single episode, successfully executing instructions such as “put the dax in the box” that depend on both slow-learned (“put”, “box”) and fast-mapped (“dax”) word meanings.
Via controlled generalization experiments, we find that the agent is reasonably robust to a degree of variation in the number of objects involved in a given fast-mapping task at test time. The agent also exhibits above-chance success when presented with the name for a particular object in the ShapeNet
ar X
iv :2
00 9.
01 71
9v 4
[ cs
.C L
] 1
4 O
ct 2
02 0
taxonomy (Chang et al., 2015) and then instructed (using that name) to interact with a different exemplar from the same object class, and this propensity can be further enhanced by specific metatraining. We find that both the number of unique objects observed by the agent during training and the temporal aspect of its perceptual experience of those objects contribute critically to its ability to generalize, particularly its ability to execute fast-mapping with entirely novel objects. Finally, we show that a dual-coding memory schema can provide a more effective basis to derive a signal for intrinsic motivation than a more conventional (unimodal) memory.

heading: 2 AN ENVIRONMENT FOR FAST WORD LEARNING

text:
We conduct experiments in a 3D room built with the Unity game engine. In a typical episode, the room contains a pre-specified number N of everyday 3D rendered objects from a global set G. In all training and evaluation episodes, the initial positions of the objects and agent are randomized. The objects include everyday household items such as kitchenware (cup, glass), toys (teddy bear, football), homeware (cushion, vase), and so on.
Episodes consist of two phases: a discovery phase, followed by an instruction phase (see Figure 1).1 In the discovery phase, the agent must explore the room and fixate on each of the objects in turn. When it fixates on an object, the environment returns a string with the name of the object (which is a nonsense word), for example “This is a dax” or “This is a blicket”. Once the environment has returned the name of each of the objects (or if a time limit of 30s is reached), the positions of all the objects and the agent are re-randomized and the instruction phase begins. The environment then emits an instruction, for example “Pick up a dax” or “Pick up a blicket”. To succeed, the agent must then lift up the specified object and hold it above 0.25m for 3 consecutive timesteps, at which point the episode ends, and a new episode begins with a discovery phase and a fresh sample of objects from the global set G. If the agent first lifts up an incorrect object, the episode also ends (so it is not possible to pick up more than one object in the instruction phase). To provide a signal for the agent to learn from, it receives a scalar reward of 1.0 if it picks up the correct object in the instruction phase. In the default training setting, to encourage the necessary information-seeking behaviour, a smaller shaping reward of 0.1 is provided for visiting each of the objects in the discovery phase.
Given this two-phase episode structure, two distinct learning challenges can be posed to the agent. In a slow-learning regime, the environment can assign the permanent name (e.g. “cup”, “chair”) to objects in the environment whenever they are sampled. By contrast, in the fast-mapping regime,
1Rendered images are higher resolution than those passed to the agent.
which is the principal focus of this work, the environment assigns a unique nonsense word to each of the objects in the room at random on a per-episode basis. The only way to consistently solve the task is to record the connections between words and objects in the discovery phase, and apply this (episode-specific) knowledge in the instruction phase to determine which object to pick up.

heading: 3 MEMORY ARCHITECTURES FOR AGENTS WITH VISION AND LANGUAGE

text:
The agents that we consider build on a standard architecture for reinforcement learning in multimodal (vision + language) environments (see e.g. (Chaplot et al., 2018; Hermann et al., 2017; Hill et al., 2020)). The visual input (raw pixels) is processed at every timestep by a convolutional network with residual connections (a ResNet). The language input is passed through an embedding lookup layer plus self-attention layer for processing. Finally, a core memory integrates the information from the two input sources over time. A fully-connected plus softmax layer maps the state of this core memory to a distribution over 46 actions, which are discretizations of a 9-DoF continuous agent avatar. A separate layer predicts a value function for computing a baseline for optimization according to the IMPALA algorithm (Espeholt et al., 2018).
We replicated previous studies by verifying that a baseline architecture with LSTM core memory (Hochreiter & Schmidhuber, 1997) could learn to follow language instructions when trained in the slow-learning regime. However, the failure of this architecture to reliably learn to perform abovechance in the fast-learning regime motivated investigation of architectures involving explicit external memory modules. Given the two observation channels from language and vision, there are various ways in which observations can be represented and retrieved in external memory.
Differentiable Neural Computer (DNC) In the DNC (Wayne et al., 2018), at each timestep t a latent vector et = w(ht−1, rt−1,xt), computed from the previous hidden state ht−1 of the agent’s core memory LSTM, the previous memory read-out rt−1, and the current inputs xt, is written to a slot-based external memory. In our setting, the input xt is a simple concatenation [vt, lt] of the output of the vision network and the embedding returned by the language network. Before writing to memory, the latent vector et is also passed to the core memory LSTM to produce the current state ht. The agent reads from memory by producing a query vector q(ht) and read strength β(ht), and computing the cosine similarity between the query and all embeddings currently stored in memory ei (i < t). The external memory returns only the k most similar entries in the memory (where k is a hyperparameter), and corresponding scalar similarities. The returned embeddings are then aggregated into a single vector r̂t by normalizing the similarities and taking a weighted average of the embeddings. This reading procedure is performed simultaneously by n independent read heads, and the results [r̂1t , . . . , r̂ n t ] are concatenated to form the current memory read-out rt. The vectors et and ht are output to the policy and value networks.
Dual-coding Episodic Memory (DCEM) We propose an alternative external key-value memory architecture inspired by the Dual-Coding theory of human memory (Paivio, 1969). The key idea is to allow different modalities (language and vision) to determine either the keys (and queries) or the values. In the present work, because of the structure of the tasks we consider, we align the keys and queries with language and the values with vision. However, for different problems (such as those requiring language production) the converse alignment could be made, or a single memory system could implement both alignments.
In our implementation, at each timestep the agent writes the current linguistic observation embeddings lt to the keys of the memory and the current visual embedding vt to its values. To read from the memory, a query q(vt, lt,ht−1) is computed and compared to the keys by cosine similarity. The k values whose keys are most similar to the query, [mj ]j≤k, are returned together with similarities [sj ]j≤k. To aggregate the returned memories into a single vector rt, the similarities are first normalized into a distribution {ŝj} and then applied to weight the memories m̂j = ŝjmj . These k weighted memories are then passed through a self-attention layer and summed elementwise to produce rt. As before this is repeated for n read heads, and the results concatenated to form the current memory read-out rt. rt is then concatenated with ht−1 and new inputs xt to compute a latent vector et = w(ht−1, rt,xt), which is passed to the core memory LSTM to produce the subsequent state ht, and finally et and ht are output to the policy and value networks.
Mean (S.D) accuracy Architecture 1e9 training steps
Gated Transformer (XL) We also consider an architecture where the agent’s core memory is a Transformer (Vaswani et al., 2017), including the gating mechanism from Parisotto et al. (2019). The only difference from Parisotto et al. (2019) is that we consider a multi-modal environment, where the observations xt passed to the core memory are the concatenation of visual and language embeddings. We use a 4-layer Transformer with a principal embedding size of 256 (8 parallel heads with query, key and value size of 32 per layer). These parameters are chosen to give a comparable number of total learnable parameters to the DCEM architecture.
Policy learning The agent’s policy is trained by minimizing the standard V-trace off-policy actorcritic loss (Espeholt et al., 2018). Gradients flow through the policy layer and the core LSTM to the memory’s query network and the embedding ResNet and self-attention language encoder. We also use a policy entropy loss as in (Mnih et al., 2016; Espeholt et al., 2018) to encourage random-action exploration. For more details and hyperparameters see Appendix A.4.
Observation reconstruction In order to provide a stronger representation-shaping signal, we make use of a reconstruction loss in addition to the standard V-trace setup. The latent vector et is passed to a ResNet g that is the transpose of the image encoder, and outputs a reconstruction of the image input dimt = g(et). The image reconstruction loss is the cross entropy between the input and reconstructed images: limt = −ximt logdimt − (1 − ximt ) log(1 − dimt ). The language decoder is a simple LSTM, which also takes the latent vector et as input and produces a sequence of output vectors that are projected and softmaxed into classifications over the vocabulary dlangt . The loss is the cross entropy between the classification produced and the one-hot vocabulary indices of the input words: llangt = −x lang t logd lang t − (1− x lang t ) log(1− d lang t ). For more details regarding the flow of information and gradients see Appendix A.4.

heading: 4 EXPERIMENTS

text:
We compared the different memory architectures with and without semi-supervised reconstruction loss on a version of the fast-mapping task involving three objects (N = 3) sampled from a global set of 30 (|G| = 30). As shown in Table 1, only the DCEM and Transformer architectures reliably solve the task after 1× 109 timesteps of training.
DCEM vs. TransformerXL Importantly, the Transformer and DCEM are the two architectures that can exploit the principle of dual-coding. Since the inputs to the Transformer are the concatenation of visual and language codes, this model can recover the dual-coding aspect of the DCEM by learning self-attention weights Wk and Wq that project the language code to keys and queries, and weights Wv to project the visual code to values. Learning in the DCEM was marginally more sample-efficient, but this is perhaps expected given it was designed with fast-mapping tasks in mind. In light of this, is it really worth pursing memory systems with explicit episodic memories?
To show one clear justification for external memory architectures, we conducted an additional comparison in which the memory windows of both the DCEM and the Transformer agents were limited to 100 timesteps (from 1024 in the original experiment), approximately the length of an episode if an agent is well-trained to the optimal policy. With a memory span of 100, the Transformer is forced to use the XL window-recurrence mechanism to pass information across context windows (Dai et al., 2019), while any capacity to retain episodic information beyond 100 timesteps in the DCEM must be managed by by the LSTM controller. In this setting we observed that the DCEM was substantially more effective (Table 1, left, bottom). While this imposed memory constraint may seem arbitrary, in real-world tasks working memory will always be at a premium. These results suggest that DCEM is more ‘working-memory-efficient’ than the Transformer agent. Indeed, by employing a simple heuristic by which the agent only writes to its external memory when the language observation changes from one timestep to the next, the DCEM agent with only 20 memory slots could solve the task with similar efficiency to a Transformer agent with a 1024-slot memory. See Appendix A.1 for these results and details of the selective writing heuristic.

heading: 4.1 GENERALIZATION

text:
To explore the generalization capabilities of our agents, we subjected trained agents to various behavioural probes, and measured performance across thousands of episodes without updating their weights. Unless stated otherwise, all experiments in this section involve the DCEM+Recons agent.
Number of objects We first probed the robustness of the agent to fast-mapping episodes with different numbers of objects. In all conditions, the same objects appear in both the discovery and instruction phases of the episode, and the objects are sampled from the same global setG (|G| = 30). As shown in Figures 2(b) and (c) (red curves), with the (default) meta-training setting involving three objects in each episode, performance on episodes involving five objects is approximately 70%, and with eight objects around 50%. This sub-optimal performance suggests that, with this metatraining regime, the agent does tend to overfit, to some degree, to the “three-ness” of its experience. Figure 2(b) shows, however, that the overfitting of the agent can be alleviated by increasing the number of objects during meta-training. Finally, Figure 2(a) confirms, perhaps unsurprisingly, that the agent has no problem generalizing to episodes with fewer objects than it was trained on.
Novel objects To probe the ability of the agents to quickly learn about any arbitrary new object, we instrumented trials with objects sampled from a global test set of novel objects H : H ∩ G = ∅, |H| = 10. As shown in Figure 3, we found that an agent meta-trained on 20 objects (i.e. |G| = 20) was almost perfectly robust to novel objects. As may be expected, this robustness degraded to some degree with decreasing |G|, which is symptomatic of the agent specializing (and overfitting) to the particular features and distinctions of the objects in its environment. However, we only observed a substantial reduction in robustness to new objects when |G|was reduced as low as three – i.e. a metatraining experience in which all episodes contain the same three objects (the first three elements of G alphabetically, i.e. a boat, a book and a bottle).
Fast category extension Children aged between three and four can acquire in one shot not only bindings between new words and specific unfamiliar objects, but also bindings between new words and categories (Behrend et al., 2001; Waxman & Booth, 2000; Vlach & Sandhofer, 2012). We conducted an analogous experiment by exploiting the category structure in ShapeNet (Chang et al., 2015). In a test trial, in the discovery phase the agent is presented with exemplars from three novel (held-out) ShapeNet categories (together with nonsense names). In the instruction phase, the agent must then pick up a different and unseen exemplar from one of these three new categories as instructed. As shown in Figure 4, when trained as described previously, the agent achieves around 55% accuracy on test trials, which is above chance (33%) but still a substantial error rate. However, this performance can be improved by requiring the agent to extend the training object categories as it learns. In this regime, three ShapeNet exemplars from distinct classes are encountered by the agent in the discovery phase of training episodes, and the instruction phase involves different exemplars from the same three classes. When trained in this way (which share similarities with matching networks (Vinyals et al., 2016)), performance on extending novel categories increases to 88%.
Role of temporal aspect Through ablations we found that both novel objects generalization and category extension relied on the agent reading multiple values from memory for each query. See A.2 for a discussion of these results, which suggest that the temporal aspect of the agent’s experience (and learning from multiple views of the same object) is an important driver of generalization.

heading: 4.2 INTRINSIC MOTIVATION

text:
The default version of the fast-mapping task includes a shaping reward to encourage the agent to visit all objects in the room. Without this reward, the credit assignment problem of a fast-mapping episode is too challenging. However, we found that the DCEM agent was able to solve the task
without shaping rewards by employing a memory-based algorithm for intrinsic motivation (NGU; Badia et al. (2020)). NGU computes a ‘surprise’ score for observations by computing its distance to other observations in the episodic memory, as described in Appendix A.4.3. The surprise score is applied as a reward signal rNGU which is added to the environment reward to encourage the agent to seek new experiences. We compared the effect of doing this in the DNC and the DCEM agents. For DCEM, the NGU computation can be applied to the memory’s keys (language) column, its values (vision) column, or both. In the former case, the agent seeks novelty in the language space rNGUlang , and in the latter, in the visual space. The final reward is r = rext + λlangrNGUlang + λimr NGU im . As shown in Figure 5, we found that the DCEM agent (with λlang = 10−3 and λim = 3 × 10−5) was able to solve the fast-mapping tasks without any shaping reward. This was not the case for the DNC agent, presumably because the required signal for ‘language-novelty’ is not approximated as well by the surprise score of the merged visual-language codes in the episodic memory.

heading: 4.3 INTEGRATING FAST AND SLOW LEARNING

text: To test whether our agents can integrate new information with existing lexical (and perceptual and motor) knowledge, we combined a fast-mapping task with a more conventional instruction-following task. In the discovery phase, the agent must explore to find the names of three unfamiliar objects, but in this case the room also contains a large box and a large bed, both of which are immovable. The positions of all objects and the agent are then re-randomized as before. In the instruction phase, the agent is then instructed to put one of the three movable objects (chosen at random) on either the bed or in the box (again chosen at random). As shown in Figure 6, if the training regime consisted of conventional lifting and putting tasks, together with a fast-mapping lifting task and a fast-mapping putting task, the agent learned to execute the evaluation trials with near-perfect accuracy. Notably, we also found that substantially-above-chance performance could be achieved on the evaluation trials without needing to train the agent on the evaluation task in any form. If we trained the agent on conventional lifting and putting tasks, and a fast-mapping task involving lifting only, the agent could recombine the knowledge acquired during this training to resolve the evaluation trials as a novel (zero-shot) task with less-than-perfect but substantially-above-chance accuracy.

heading: 4.4 RESULTS WITH ANOTHER ENVIRONMENT

text: To verify that the observed effects hold beyond our specific Unity environment, we added a new task to the DeepMind Lab suite (Beattie et al., 2016). Results for this task are given in Appendix A.3.

heading: 5 RELATED WORK

text:
Meta-learning, of the sort observed in our agent, has been applied to train matching networks: image classifiers that can assign the correct label to a novel image, given a small support set of (image, label) pairs that includes the correct target label (Vinyals et al., 2016). Our work is also inspired by
Snell et al. (2017), who propose a more efficient way to integrate a small support set of experience into a coherent space of image ‘concepts’ for improved fast learning, and Santoro et al. (2016), who show that the successful meta-training of image classifiers can benefit substantially from external memory architectures such as Memory Networks (Weston et al., 2014) or DNC (Graves et al., 2016).
In NLP, meta-learning has been used to train few-shot classifiers for various tasks (see Yin (2020) for a recent survey). Meta-learning has also previously been observed in reinforcement learning agents trained with conventional policy-gradient algorithms (Duan et al., 2016; Wang et al., 2019). In Model-Agnostic Meta Learning (Finn et al., 2017), models are (meta) trained to be easily tunable (by any gradient algorithm) given a small number of novel data points. When combined with policygradient algorithms, this technique yields fast learning on both 2D navigation and 3D locomotion tasks. In cognitive tasks where fast learning is not explicitly required, external memories have proven to help goal-directed agents (Fortunato et al., 2019), and can be particularly powerful when combined with an observation reconstruction loss (Wayne et al., 2018).
Recent work at the intersection of psychology and machine learning is also relevant in that it shows how the noisy, first-person perspective of a child can support the acquisition of robust visual categories in artificial neural networks (Bambach et al., 2018). When deep networks are trained on data recorded from children’s head cameras, unsupervised or semi-supervised learning objectives can substantially improve the quality of the resulting representations (Orhan et al., 2020).

heading: 6 CONCLUSION

text:
Our experiments have highlighted various benefits of having an explicitly multi-modal episodic memory system. First, mechanisms that allow the agent to query its memory in a modality-specific way (either within or across modalities) can better allow them to rapidly infer and exploit connections between perceptual experience and words, and therefore to realize fast-mapping, a notable aspect of human learning. Second, external (read-write) memories can achieve better performance for the same number of memory ‘slots’ than Transformer-based memories. This greater ‘memoryefficiency’ may be increasingly important as agents are applied to real-world tasks with very long episodic horizons. Third, in cases where it is useful to estimate the degree of novelty or “surprise” in the current state of the environment (for instance to derive a signal for intrinsic motivation), a more informative signal may be obtained by separately estimating novelty based on each modality and aggregating the result. Finally, an episodic memory system may ultimately be essential for fast knowledge consolidation. The potential for memory buffers and offline learning processes such as experience replay to support knowledge consolidation is not a new idea (McClelland et al., 1995; Mnih et al., 2016; Lillicrap et al., 2016; McClelland et al., 2020). For language learning agents, the need to both rapidly acquire and retain multi-modal knowledge may further motivate explicit
external memories. Retaining in memory visual experiences together with aligned (and hopefully pertinent) language (i.e. a dual-coding schema) may facilitate something akin to offline ‘supervised’ language learning. We leave this possibility for future investigations, which we will facilitate by releasing publicly the environments and tasks described in this paper.

heading: A APPENDICES

heading: A.1 COMPARING TRANSFORMERXL TO DCEM WHEN MEMORY IS LIMITED

text:
Both the TransformerXL and DCEM/DNC agents have a hyperparameter that determines the effective size of their explicit working memory. In a vanilla Transformer it determines the size of the window of timesteps that the network can be applied to for each forward (and backward) pass. To give models some chance of passing information beyond this hard constraint, TransformerXL architecture conditions each forward pass also on representations computed in the previous window, which establishes a form of recurrence over time from window to window. In our original experiments, this mechanism was not tested, because we set the window size to 1024 timesteps, which is longer than most episodes of the fast-mapping task, which are typically 80-120 timesteps for a well-trained agent.
To examine the performance of the TransformerXL in cases where it is required to pass information across context windows, we reduced the size of the window. For a fair comparison, we similarly reduced the equivalent parameter (the capacity in rows in the FIFO external memory) for the DCEM agent. As shown in Figure 7, for cases where the memory window size (or buffer) is reduced to (100) or below (50, 20) the normal episode length, the DCEM performs better than the TransformerXL agent. This suggests that the TransformerXL has difficulty making the necessary (visual and linguistic) information available to policy head when that information must be passed between context windows. Surprisingly, the DCEM was able to learn the tasks efficiently with a memory size of 50, which suggests that it must exploit its LSTM controller to retain sufficient information when its external memory begins to overflow. Both architectures fail when the memory size is reduced to 20, but in that case a DCEM agent can in fact learn the optimal policy if a simple heuristic for selective writing, described below, is employed. This highlights an advantage of explicitly read-write external memories; information can be managed via the reading or the writing process. It is not immediately obvious how the same strategy could be applied with a window-based memory architectecture like TransformerXL.

heading: A.1.1 SELECTIVE WRITING HEURISTIC

text:
One advantage of explicit external (read-write) memories is that the flow of information to the agent’s policy can be influenced by the writing process as well as the reading function. To verify this fact, we implemented a simple non-parametric heuristic writing condition in the DCEM architecture, whereby observations are written to the external memory when there is a change to the observation in the language channel. This heuristic aligns with the principle of dual-coding exploited elsewhere in the paper: while visual observations change continuously every timestep, changes to observed language are rare events that might signal some important change in the environment.
More formally, our heuristic relies on a window size parameter w that we set to 3 in all cases. For observation xt = {vt, lt}, a language change indicator It ∈ N is set as I0 = 0
It = { t, if lt = lt−1 It−1, otherwise.
Then, the content ct written to memory at t is
ct = { {vt, lt}, if t− It < w ∅, otherwise,
where, as before, vt, lt are the agent embeddings of the visual and linguistic observations. Thus, memories are written for w timesteps proceeding a change in the language observation.

heading: A.2 ROLE OF TEMPORAL ASPECT IN GENERALIZATION

text:
In seeking to understand the mechanisms that support the generalization effects reported in the main paper, we found that the parameter k was an important factor, where the top-k memories are returned to the agent policy head per memory read. As the agent explores during the discovery phase of episodes, it writes multiple perspectives of the same object to memory. As shown in Figure 8, both its robustness to entirely novel objects (left) and its ability to extend categories from novel exemplars (right), as well as its ability to solve the training task, are enhanced when k > 1; i.e. when it determines which object to visit in the instruction phase based on memories written from more than one view of each object.
A.3 VERIFICATION IN DEEPMIND LAB
At a high level, the design of an episode is very similar to the default fast-mapping task in the Unity environment. The agent must move down a corridor, bumping into (and collecting) three distinct objects. When an agent collects an object it is immediately presented with the (episode-specific) name for that object. After passing three objects, the corridor opens into a room containing two of the three objects found in the corridor. Upon entering the room, the agent is presented with the name corresponding to one of the two objects, and must bump into that object in order to receive a reward of 1. As before, a shaping reward of 0.1 is given as the agent collects each object in the corridor. Compared with the Unity environment, the DeepMind Lab action space is smaller (8 vs. 46 actions), the objects are larger, and the agent has no substantive way to interact with the objects (they disappear the moment the agent collides with them). Note also that the agent must choose between 2 (rather than 3) objects in the instruction phase, so an agent selecting objects at random would achieve 50% accuracy.
To provide some sense of the robustness and generality of the effects observed thus far, we applied the various agent architectures directly to this environment with no environment-specific tuning. As shown in Table 2, without any further tuning of the agent, we observe a similar pattern of results
in DeepMind Lab as in the Unity room. As in that case, the Transformer and DCEM architectures performed best, with three and two seeds out of five (respectively) mastering the training task. As in the Unity environment, we also observed above-chance ability to apply fast-mapping knowledge zero-shot to unseen objects at test time.

heading: A.4.1 ARCHITECTURE DIAGRAMS

heading: A.4 AGENT ARCHITECTURE DETAILS

text:
A.4.2 HYPERPARAMETERS
A.4.3 INTRINSIC MOTIVATION ALGORITHM
The intrinsic reward is based on the similarity (Euclidean distance) between the new embedding e and the nearest neighbors already present in memory {ei}. The average distance ρ̄ used below is a lifetime average of all ρi that is updated with every computation.
ρi = |e− ei|2
ρ̄+ c
ki =
max(ρi − ρmin, 0) +
s = (∑ i ki )1/2 + c
rNGU = { 1/s if s < smax 0 otherwise
The computation introduces the constants c, , ρmin, and smax, and the number of neighbours |{ei}| used for the similarity estimate. Table 4 lists the values we used.

heading: A.5 ENVIRONMENT DETAILS

heading: A.5.1 UNITY ACTION SPACE

text:
The following discrete actions (and strengths) are available to the agent in all experiments except for those in DeepMind Lab. The scalar strengths are translated into force and torque (for rotations) by the environment engine.
Movement without grip Fine grained movements without grip Movement with grip NOOP, MOVE RIGHT(0.05), GRAB, MOVE FORWARD(1), MOVE RIGHT(-0.05), GRAB + MOVE FORWARD(1), MOVE FORWARD(-1), LOOK DOWN(0.03), GRAB + MOVE FORWARD(-1), MOVE RIGHT(1), LOOK DOWN(-0.03), GRAB + MOVE RIGHT(1), MOVE RIGHT(-1), LOOK RIGHT(0.2), GRAB + MOVE RIGHT(-1), LOOK RIGHT(1), LOOK RIGHT(-0.2), GRAB + LOOK RIGHT(1), LOOK RIGHT(-1), LOOK RIGHT(0.05), GRAB + LOOK RIGHT(-1), LOOK DOWN(1), LOOK RIGHT(-0.05), GRAB + LOOK DOWN(1), LOOK DOWN(-1), GRAB + LOOK DOWN(-1),
Fine grained movments with grip Object manipulation Fine grained object manipulation

heading: GRAB + MOVE RIGHT(0.05), GRAB + SPIN RIGHT(1), GRAB + PULL(0.5),

text: GRAB + MOVE RIGHT(-0.05), GRAB + SPIN RIGHT(-1), GRAB + PULL(-0.5), GRAB + LOOK DOWN(0.03), GRAB + SPIN UP(1), PULL(0.5), GRAB + LOOK DOWN(-0.03), GRAB + SPIN UP(-1), PULL(-0.5), GRAB + LOOK RIGHT(0.2), GRAB + SPIN FORWARD(1), GRAB + LOOK RIGHT(-0.2), GRAB + SPIN FORWARD(-1), GRAB + LOOK RIGHT(0.05), GRAB + PULL(1), GRAB + LOOK RIGHT(-0.05), GRAB + PULL(-1),

heading: A.5.2 SHAPENET

text:
ShapeNet contains 3D models of objects with a wide range of complexity and quality. To guarantee that the models are recognizable and of a high quality, we manually filtered the ShapeNet Sem dataset, selecting a subset of everyday semantic classes, ensuring that the selected models had a reasonable number of vertices, and reasonable size and weight dimensions. The selected classes (number of models in each class) were as follows, with a total of 1,437 models across 31 different classes:
armoire (31), bag (11), bed (65), book (47), bookcase (13), bottle (26), box (37), bunk bed (9), chair (150), chest of drawers (133), coffee table (43), computer (6), floor lamp (68), glass (11), hammer (15), keyboard (5), lamp (100), loudspeaker (37), microwave (16), monitor (58), mug (12), piano (11), plant (31), printer (22), rug (36), soda can (20), sofa (145), stool (25), table (181), vase (66), wine bottle (7).

title: GROUNDED LANGUAGE LEARNING FAST AND SLOW