summaries: The paper proposes a method (Skyformer) to apply Nystrom approximation to the attention matrix. It embeds the attention matrix inside a larger PSD matrix (unlike Nystromformer), allowing Nystrom method to work well. Theoretical and empirical validation shows that the method is competitive with other forms of efficient attention (e.g. BigBird, Performer, Reformer).

summaries: The paper presents an alternative approach to attention in the Transformer. There are two modeling contributions:  1. Kernelized attention that omits the typical normalization term used for softmax attention. 2. A Nystrom method approach that allows efficiently approximating the output of the attention mechanism  The contributions also include theoretical results regarding the accuracy of the Nystrom approximation, and experimental results showing better convergence, higher accuracy, and medium efficiency gains from the approach.  Below I summarize my understanding/takeaways from the two key methods; please correct me in the author response if these are inaccurate: 1. As I understand it, "kernelized attention" mean that target keys don't compete with each other, the only thing that matters is their distance to the query. In standard softmax attention, suppose that all key vectors are orthogonal to a query vector -- in that case, attention would weigh all values uniformly. In kernelized attention, on the other hand, if key vectors are all far away from the query vector, then all value vectors would be attended to with near-zero weight. 2. The Nystrom approximation involves sampling some queries and keys, and then computing kernels between all pairs. But whereas Nystromformer computes similarity for query+key pairs only, the proposed approach also computes kernels for query+query and key+key pairs as well.

summaries: The author proposes a modified Nyström method named Skyformer to approximate the Kernelized Attention. They also provide a theoretical guarantee that the approximation error is small in terms of the spectral norm. Extensive experiments are also conducted to show that Skyformer achieves comparable performance to vanilla self-attention with fewer computational costs.

summaries: The paper proposes a new form of self-attention for use in transformer models using Gaussian Kernels. They show that the new attention method is comparable in accuracy to traditional softmax self-attention. They then go on to show that kernel based attention can be approximated by taking a Nystrom sketch of the kernel matrix. They give rigorous bounds on the spectral norm of the approximation (assuming that the kernel matrix has a high condition number).

summaries: The authors reformulate soft max structure of the self-attention mechanism in Transformers in terms of a Gaussian kernel evaluation. They claim that such a model is empirically stable during training, while allowing the use of various Kernel approximation strategies to improve the quadratic computational complexity.   In this work, approximation strategy used by the authors is the popular Nyström approximation. However, since the kernel matrix formed between the query and the key matrices is not symmetric (hence not PSD), the Nyström approximation is not directly applicable. The authors propose to form a PSD kernel matrix by first evaluating the Gaussian kernel on the row concatenated  query and key matrices which is PSD, performing the Nyström approximation and then choosing the appropriate block in the resulting matrix.  By adapting an existing theorem [Musco, Musco 17],  the authors claim that such a low-rank approximation of the kernelized attention score is bounded in terms of the relative spectral norm. This is claimed to be due to the observed fast decay of the eigenvalues in an empirical kernel matrix.   The matrix inversion required for the Nyström approximation, however is a potential issue in terms of numerical stability on a GPU. To circumvent the problem, instead of a standard conjugate gradient method, they use a Schulz type iteration which only uses matrix products, without any floating point divisions.    Empirical evaluation shows marginal improvement on most tasks on the LRA benchmark, as compared to the naive Nyström approximation [Xiong et al 2021]. However, this is at the cost of relatively larger time and memory usage.

abstractText: Transformers are expensive to train due to the quadratic time and space complexity in the self-attention mechanism. On the other hand, although kernel machines suffer from the same computation bottleneck in pairwise dot products, several approximation schemes have been successfully incorporated to considerably reduce their computational cost without sacrificing too much accuracy. In this work, we leverage the computation methods for kernel machines to alleviate the high computational cost and introduce Skyformer, which replaces the softmax structure with a Gaussian kernel to stabilize the model training and adapts the Nyström method to a non-positive semidefinite matrix to accelerate the computation. We further conduct theoretical analysis by showing that the matrix approximation error of our proposed method is small in the spectral norm. Experiments on Long Range Arena benchmark show that the proposed method is sufficient in getting comparable or even better performance than the full self-attention while requiring fewer computation resources.

authors:
Arturs Backurs
Piotr Indyk
Tal Wagner

title: Space and time efficient kernel density estimation in high dimensions

venue: Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems

authors:
Iz Beltagy
Matthew E. Peters
Arman Cohan

title: Longformer: The long-document transformer

authors:
Srinadh Bhojanapalli
Chulhee Yun
Ankit Singh Rawat
Sashank Reddi
Sanjiv Kumar

title: Lowrank bottleneck in multi-head attention models

venue: In International Conference on Machine Learning,

authors:
Moses Charikar
Paris Siminelakis

title: Hashing-based-estimators for kernel density in high dimensions

venue: IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS),

authors:
Yifan Chen
Yun Yang

title: Accumulations of projections—a unified framework for random sketches in kernel ridge regression

venue: In International Conference on Artificial Intelligence and Statistics,

title: Fast statistical leverage score approximation in kernel ridge regression

authors:
Rewon Child
Scott Gray
Alec Radford
Ilya Sutskever

title: Generating long sequences with sparse transformers

authors:
Krzysztof Choromanski
Valerii Likhosherstov
David Dohan
Xingyou Song
Andreea Gane
Tamás Sarlós
Peter Hawkins
Jared Davis
Afroz Mohiuddin
Lukasz Kaiser
David Belanger
Lucy Colwell
Adrian Weller

title: Rethinking attention with performers

venue: URL https://arxiv.org/abs/2009.14794

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

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

authors:
Yihe Dong
Jean-Baptiste Cordonnier
Andreas Loukas

title: Attention is not all you need: Pure attention loses rank doubly exponentially with depth

venue: arXiv preprint arXiv:2103.03404,

authors:
Petros Drineas
Michael W Mahoney
Nello Cristianini

title: On the nyström method for approximating a gram matrix for improved kernel-based learning

venue: journal of machine learning research,

authors:
Alex Gittens
Michael W Mahoney

title: Revisiting the nyström method for improved large-scale machine learning

venue: The Journal of Machine Learning Research,

authors:
Sariel Har-Peled
Piotr Indyk
Rajeev Motwani

title: Approximate nearest neighbor: Towards removing the curse of dimensionality

venue: Theory of computing,

authors:
Xiaoqi Jiao
Yichun Yin
Lifeng Shang
Xin Jiang
Xiao Chen
Linlin Li
Fang Wang
Qun Liu

title: Tinybert: Distilling BERT for natural language understanding

venue: Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings,

authors:
Angelos Katharopoulos
Apoorv Vyas
Nikolaos Pappas
François Fleuret

title: Transformers are rnns: Fast autoregressive transformers with linear attention

authors:
Nikita Kitaev
Lukasz Kaiser
Anselm Levskaya

title: Reformer: The efficient transformer

venue: In 8th International Conference on Learning Representations,

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

authors:
Sanjiv Kumar
Mehryar Mohri
Ameet Talwalkar

title: Sampling techniques for the nystrom method

venue: In Artificial Intelligence and Statistics,

authors:
Mike Lewis
Yinhan Liu
Naman Goyal
Marjan Ghazvininejad
Abdelrahman Mohamed
Omer Levy
Veselin Stoyanov
Luke Zettlemoyer

title: BART: denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension

authors:
Drew Linsley
Junkyung Kim
Vijay Veerabadran
Charles Windolf
Thomas Serre

title: Learning long-range spatial dependencies with horizontal gated recurrent units

venue: Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems

authors:
Liyuan Liu
Xiaodong Liu
Jianfeng Gao
Weizhu Chen
Jiawei Han

title: Understanding the difficulty of training transformers

venue: Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020,

authors:
Weijie Liu
Peng Zhou
Zhiruo Wang
Zhe Zhao
Haotang Deng
Qi Ju

title: Fastbert: a self-distilling BERT with adaptive inference time

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

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

venue: URL http://arxiv.org/abs/1907.11692

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

title: Learning word vectors for sentiment analysis

venue: The 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, Proceedings of the Conference,

authors:
Stephen Merity
Caiming Xiong
James Bradbury
Richard Socher

title: Pointer sentinel mixture models

venue: In 5th International Conference on Learning Representations,

authors:
Cameron Musco
Christopher Musco

title: Recursive sampling for the nystrom method

venue: Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017,

authors:
Nikita Nangia
Samuel R. Bowman

title: Listops: A diagnostic dataset for latent tree learning

venue: Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics,

authors:
Jiezhong Qiu
Hao Ma
Omer Levy
Wen-tau Yih
Sinong Wang
Jie Tang

title: Blockwise selfattention for long document understanding

authors:
Dragomir R. Radev
Pradeep Muthukrishnan
Vahed Qazvinian
Amjad Abu-Jbara

title: The ACL anthology network

venue: corpus. Lang. Resour. Evaluation,

authors:
Ali Rahimi
Benjamin Recht

title: Random features for large-scale kernel machines

venue: In NIPS,

authors:
M Kafaei Razavi
Asghar Kerayechian
Mortaza Gachpazan
Stanford Shateyi

title: A new iterative method for finding approximate inverses of complex matrices

venue: In Abstract and Applied Analysis,

authors:
Aurko Roy
Mohammad Saffar
Ashish Vaswani
David Grangier

title: Efficient content-based sparse attention with routing transformers

venue: Trans. Assoc. Comput. Linguistics,

authors:
Edward Snelson
Zoubin Ghahramani

title: Local and global sparse gaussian process approximations

authors:
Raphael Tang
Yao Lu
Linqing Liu
Lili Mou
Olga Vechtomova
Jimmy Lin

title: Distilling taskspecific knowledge from BERT into simple neural networks

authors:
Yi Tay
Dara Bahri
Donald Metzler
Da-Cheng Juan
Zhe Zhao
Che Zheng

title: Synthesizer: Rethinking self-attention in transformer models

authors:
Yi Tay
Mostafa Dehghani
Samira Abnar
Yikang Shen
Dara Bahri
Philip Pham
Jinfeng Rao
Liu Yang
Sebastian Ruder
Donald Metzler

title: Long range arena: A benchmark for efficient transformers

authors:
Yi Tay
Mostafa Dehghani
Dara Bahri
Donald Metzler

title: Efficient transformers: A survey

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

title: Attention is all you need

authors:
Elena Voita
David Talbot
Fedor Moiseev
Rico Sennrich
Ivan Titov

title: Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned

venue: Proceedings of the 57th Conference of the Association for Computational Linguistics,

title: High-dimensional statistics: A non-asymptotic viewpoint, volume 48

authors:
Sinong Wang
Belinda Z. Li
Madian Khabsa
Han Fang
Hao Ma

title: Linformer: Self-attention with linear complexity

authors:
Christopher Williams
Matthias Seeger

title: Using the nyström method to speed up kernel machines

venue: In Proceedings of the 14th annual conference on neural information processing systems, number CONF,

authors:
Thomas Wolf
Lysandre Debut
Victor Sanh
Julien Chaumond
Clement Delangue
Anthony Moi
Pierric Cistac
Tim Rault
Rémi Louf
Morgan Funtowicz
Jamie Brew

title: Huggingface’s transformers: State-of-the-art natural language processing

venue: URL http://arxiv.org/abs/1910.03771

authors:
Yunyang Xiong
Zhanpeng Zeng
Rudrasis Chakraborty
Mingxing Tan
Glenn Fung
Yin Li
Vikas Singh

title: Nyströmformer: A nyström-based algorithm for approximating self-attention

authors:
Yun Yang
Mert Pilanci
Martin J Wainwright

title: Randomized sketches for kernels: Fast and optimal nonparametric regression

venue: The Annals of Statistics,

authors:
Manzil Zaheer
Guru Guruganesh
Kumar Avinava Dubey
Joshua Ainslie
Chris Alberti
Santiago Ontañón
Philip Pham
Anirudh Ravula
Qifan Wang
Li Yang
Amr Ahmed

title: Big bird: Transformers for longer sequences

venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems

authors:
Haoyi Zhou
Shanghang Zhang
Jieqi Peng
Shuai Zhang
Jianxin Li
Hui Xiong
Wancai Zhang

title: Informer: Beyond efficient transformer for long sequence time-series forecasting

heading: 1 Introduction

text:
The cost of language model training increases exponentially. Among different models, Transformerbased language models [Vaswani et al., 2017, Devlin et al., 2019, Liu et al., 2019, Lewis et al., 2020] are shown to enjoy state-of-the-art (SOTA) performances on many Natural Language Processing (NLP) tasks despite their enormous training cost. One of the computation bottlenecks lies in the self-attention mechanism, which is known to be resource-intensive with quadratic time and space complexity (O(n) where n is the input sequence length). Consequently, Transformers cannot support long sequence processing and large batch size with limited resources.
The challenge of improving computational efficiency of Transformers has motivated several recent studies on attention acceleration, using either sparse attention pattern [Qiu et al., 2020, Child et al., 2019, Zaheer et al., 2020, Beltagy et al., 2020, Kitaev et al., 2020] or low-rank approximation [Choromanski et al., 2020, Wang et al., 2020]. However, there is usually a lack of theoretical analysis on the approximation error of these methods due to the complex softmax structure, which makes the theoretical comparison between the efficiency of each method infeasible. It is also unclear in theory how to set the hyper-parameters of those methods to attain a desired level of approximation accuracy.
Another issue of Transformers is the training instability that small perturbations in parameter updates tend to be amplified, resulting in significant disturbances in the model output [Liu et al., 2020a]. Transformers on some NLP tasks have shown to be sensitive to hyper-parameters, learning schedulers, or even random seeds, which usually demands a time-costly grid search for the best configuration in real-world applications. It has also been observed in our experiments that a slight change in the learning rate may cause the failure of convergence for some models. We conjecture that the instability
∗ Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
in Transformer training comes from the softmax structure, as the un-normalized attention score matrices before softmax tend to have extremely large condition numbers due to its fast singular value decay. To alleviate the instability issue, an extra factor of 1/ √ p in the softmax kernel SM is suggested by Vaswani et al. [2017] to restrain the scale variation; Liu et al. [2020a] proposes a new scheme to control the magnitude of output change and stabilize the training in early stages. In practice, we also need to consider the lower numerical precision of GPU implementation in model training, which further deteriorates the stability.
Kernel methods may be the answer to both challenges. As pointed out by Choromanski et al. [2020], the softmax structure is closely related to Gaussian kernels up to diagonal matrix multiplications, as the pairwise dot products naturally appear when expanding the squared `2 distance. We further notice some important connections between self-attention and Gaussian kernels. First, the un-normalized attention score matrix can be formed via basic matrix operations on an empirical Gaussian kernel matrix. Moreover, the form of Gaussian kernels has the natural interpretation of assigning “attention” to different tokens. Compared to the softmax function, Gaussian kernels automatically perform the normalization as softmax does (c.f. Section 4.1). These observations motivate us to replace the softmax structure with Gaussian kernels. As we demonstrated in this paper, the new attention model, Kernelized Attention, empirically stabilizes the model training while being comparable to self-attention in model accuracy.
To further improve the efficiency, we propose Skyformer (Symmetrization of Kernelized attention for NYström method) to accelerate kernelized attention. Skyformer adapts the Nyström method [Williams and Seeger, 2001, Drineas et al., 2005] to the non-PSD empirical Gaussian kernel matrix (as query matrices in general do not equal to key matrices), by instead lifting the kernelized attention score matrix into a large PSD matrix that contains the un-normalized attention score matrix as the off-diagonal block. We further conduct theoretical analysis by showing that Skyformer has a small matrix approximation error on kernelized attention in the spectral norm. Our experiments on the LRA benchmark show that Skyformer consistently uses less space and time while achieving better accuracy than other baseline methods.
In summary, our main contributions are:
(1) We revisit the intrinsic connection between self-attention and kernel methods, and explore a new kernel-based structure, kernelized attention, to stabilize the training of Transformers.
(2) We propose Skyformer, which approximates the kernelized attention via low dimensional randomized sketches by adapting the Nyström method to a non-PSD matrix. We provide the theoretical guarantee that the matrix multiplication error is small in term of spectral norm.
(3) Extensive experiments show that Skyformer achieves comparable performance to the original self-attention with fewer computational costs.2

heading: 2 Related Work

text:
Among all the transformer acceleration methods, including attention layer simplification by pruning redundant attention heads [Voita et al., 2019, Michel et al., 2019] and model size reduction with knowledge distillation [Jiao et al., 2020, Tang et al., 2019, Liu et al., 2020b], we focus on attention approximation models, which are closely related to kernel methods.
To reduce the time and space complexity by avoiding exhaustive computation over the attention metric, recent studies propose to apply sparse attention patterns to limit the numbers of elements participating in matrix multiplications [Qiu et al., 2020, Child et al., 2019, Zaheer et al., 2020, Beltagy et al., 2020]. Beyond limiting the attention to fixed patterns, some approaches learn the patterns by determining token assignments to relevant groups [Kitaev et al., 2020, Roy et al., 2021]. Those models utilize local and global information in the attention score matrix to perform approximation, which coincides with the attempt to accelerate the computation in Gaussian processes [Snelson and Ghahramani, 2007].
2Our code is released at https://github.com/pkuzengqi/Skyformer
The attention score matrix is known to exhibit a very fast rate of singular value decay [Bhojanapalli et al., 2020, Dong et al., 2021], similar to that of an empirical kernel matrix [Yang et al., 2017]. This near singular property motivates many low-rank attention approximation methods to skillfully leverage the computation techniques in kernel methods. Among them, Linformer [Wang et al., 2020] compresses the size of the key and value matrix with random projections based on the Johnson–Lindenstrauss transform, a common randomized sketching method in Gaussian processes [Yang et al., 2017]; Reformer [Kitaev et al., 2020] applies locality-sensitive hashing (LSH) [Har-Peled et al., 2012] to simplify the computation of the attention score matrix, which is widely used in kernel density estimation [Charikar and Siminelakis, 2017, Backurs et al., 2019]; Performer [Choromanski et al., 2020] projects both query and key matrix through random Fourier features [Rahimi et al., 2007], heavily exploiting Bochner Theorem for stationary kernels.
The most related papers to ours are linear attention [Katharopoulos et al., 2020], Synthesizer [Tay et al., 2020a], and Nyströmformer [Xiong et al., 2021]. Linear attention takes the softmax structure in self-attention as a measure of similarity and replaces it with the dot product of separately activated query and key matrices; Synthesizer aims to modify the original self-attention by replacing the dot product before softmax with Synthetic Attention, which generates the alignment matrix independent of token-token dependencies. Their attempts indicate that the softmax structure in self-attention is not the only feasible choice, and justify our usage of kernelized attention. Rather than remodeling self-attention, Nyströmformer applies the Nyström method [Williams and Seeger, 2001, Drineas et al., 2005], a powerful and effective method for large-scale kernel machines acceleration, to approximate the attention score matrix. However, Nyströmformer applies the Nyström method to a non-PSD matrix, and thus fails to utilize the full potential of the Nyström method. This issue is resolved in our proposed Skyformer by instead lifting the kernelized attention score matrix into a large PSD matrix which contains the target non-PSD matrix as its off-diagonal block. For more details on attention approximation methods, we refer readers to a survey paper on efficient transformers [Tay et al., 2020c].

heading: 3 Preliminaries and notations

heading: 3.1 Revisiting self-attention

text:
For a given input sequence X ∈ Rn×d0 of length n and embedding dimension d0, The dot-product attention for a single head in Transformer Vaswani et al. [2017] is defined as
Attention(Q,K,V ) = softmax ( QKT √ p ) V
where Q = XWQ, K = XWK , and V = XWV , and WQ, WK and WV are the query, key, and value weight metrics that linearly project the input X of d0 dimension to an output tensor of p dimensions.
To simplify the future analysis, the left softmax term can be rewritten into D−1A, where A := exp(QKT / √ p) is the un-normalized attention score matrix; D is a diagonal matrix whose diagonal is exp(QKT / √ p) · 1 (by convention 1 is a size-n vector with all elements being 1). Following the notation in Performer [Choromanski et al., 2020], we define SM(q,k) := exp(qTk/ √ p) as the softmax kernel function, and represent A by the notation SM(Q,K), which means the element aij from the i-th row and j-th column in A is equal to SM(qi,kj). Throughout this paper qi (resp. kj) means the i-th (resp. j-th) row in Q (resp. K).
We close this subsection with a short lemma to show SM(·, ·) is a positive semidefinite (PSD) kernel function [Wainwright, 2019, Definition 12.6] by relating it to Gaussian kernels. Lemma 1. SM(·, ·) is a PSD kernel function. Equivalently, for all integers n ≥ 1 and elements {qi}ni=1 ⊆ Rp, the n-by-n matrix C = SM(Q,Q) is PSD.
Proof. We first state an important equation to connect the softmax kernel and Gaussian kernels as follows:
SM(qi, qj) = exp ( qTi qj√ p ) = exp ( ‖qi‖2 2 √ p ) exp ( −‖qi − qj‖ 2 2 √ p ) exp ( ‖qj‖2 2 √ p ) .
The middle part exp ( ‖qi−qj‖2
2 √ p
) is exactly a Gaussian kernel with bandwidth p 1 4 . (Choromanski
et al. [2020] have more discussion on the findings.)
Through this equation, we can rewrite C as
C = D 1/2 Q · κ
( Q
p1/4 ,
Q
p1/4
) ·D1/2Q , (1)
where DQ is a diagonal matrix with elements (DQ)ii = exp ( ‖qi‖2√
p
) ,∀i ∈ [n], and κ(qi, qj) :=
exp ( −‖qi − qj‖2/2 ) is the standard Gaussian kernel function.
We prove the lemma by using the fact that κ is a PSD kernel and κ (
Q p1/4 , Q p1/4
) is a PSD matrix. ♦

heading: 3.2 Nyström method

text:
Due to the intrinsic low-rankness of an empirical kernel matrix B, the so-called Nyström method that replaces B with its low-rank approximation B̃, has been applied to accelerate kernel methods [Gittens and Mahoney, 2016, Kumar et al., 2009, Williams and Seeger, 2001]. Specifically, the Nyström approximation of B is the matrix B̃ = BS(STBS)†STB, where (·)† denotes the Moore-Penrose pseudoinverse of a matrix, and S ∈ Rn×d is a zero-one sub-sampling matrix whose columns are a subset of the columns in I , indicating which d observations have been selected. The formal definition of the uniform sub-sampling matrix is given as follows:
Definition 1 (Uniform sub-sampling matrix). For a random matrix S ∈ Rn×d, if S has i.i.d. columns and the j-th column S(j) can randomly be √ 1 dei with probability 1 n , where ei is the i-th column of the n-by-n identity matrix In, then S is called a uniform sub-sampling.
We close this subsection with a remark that it is not appropriate to directly extend the Nyström method from kernel method to self-attention due to a core requirement that B should be PSD with consideration of approximation performance improvement. We will show in the next section how to address this challenge and properly adapt Nyström method to non-PSD matrices.

heading: 3.3 Approximation evaluation

text:
Beyond the time and space complexity, attention acceleration methods have been mostly evaluated with empirical experiment results, such as the perplexity of pretrained language models and the finetuned performance on downstream natural language understanding tasks. Specifically, Long Range Arena benchmark [Tay et al., 2020b] has been proposed to systematically evaluate the performance of efficient transformers with ten NLP tasks in long-context scenarios. However, such empirical results are indirect for theoretical analysis. Therefore, we introduce a common criterion used in matrix approximation, spectral norm, to ease the future discussion on performance.
Definition 2 (Spectral norm guarantee for matrix approximation (MA)). Given a matrix M ∈ Rn1×n2 , two constants ε > 0, δ < 12 , we say that its approximation matrix M̃ ∈ R
n1×n2 satisfies (ε, δ)-MA property for M , if
P { ‖M − M̃‖ > ε‖M‖ } < δ. (2)
In previous works, the direct analysis of the approximation error to the entire output D−1AV in the (ε, δ)-MA manner is usually spared due to the difficulty caused by the complex softmax structure. In this paper, with the new kernelized attention, we are allowed to perform the analysis through the existing theoretical results in kernel methods. Consequently, in Section 4.5 we are able to give a relatively precise error analysis on the approximation of Skyformer to the entire kernelized attention, which eases the future comparison with other methods approximating kernelized attention.

heading: 4 Method

heading: 4.1 Kernelized Attention

text:
Kernelized Attention replaces the softmax structure in vanilla self-attention with a Gaussian kernel, and the new attention model is stated as:
Kernelized-Attention(Q,K,V ) = CV := κ ( Q
p1/4 , K p1/4
) V , (3)
where we define the n-by-n matrix C as the kernelized attention score matrix κ(Q/p1/4,K/p1/4).
The justification for using the kernelized attention model is as follows. A significant advantage of softmax attention is that tokens are allowed to attend to a limited number of other important tokens in the sequence. We observe that Gaussian kernel function can play a similar role. The expression of a Gaussian kernel is κ(qi,kj) := exp ( −‖qi − kj‖2/2 ) . Via this expression, for token i in the query, Gaussian kernel assigns a large attention score to the token j when kj is close to qi. The distance-based weight assignment is indeed considered as a major reason why kernel methods are powerful. The form of kernelized attention also leads to an automatic normalization. Based on Equation (1), the new attention model can be rewritten in terms of the un-normalized attention score matrix A as
Kernelized-Attention(Q,K,V ) = ( D −1/2 Q ·A ·D −1/2 K ) V ,
where DQ (resp. DK) is a diagonal matrix with elements (DQ)ii = exp ( ‖qi‖2√
p
) (resp. (DK)ii =
exp ( ‖ki‖2√
p
) ), ∀i ∈ [n]. We remark the kernelized attention model can thus be formally taken as
a variant of the original self-attention, which instead normalizes the matrix A in a form of D−1A. The intrinsic normalization allows kernelized attention to have a more reasonable condition number than self-attention, which benefits the stability of model training. To demonstrate the improvement in stability, we additionally provide a toy experiment in Appendix F, which shows the “condition number" of kernelized attention is smaller than self-attention. Moreover, empirical evaluation in Section 5 supports our claim that the new attention model can attain a comparable performance to the original attention model.

heading: 4.2 Skyformer: a modified Nyström method

text:
Before jumping into details of Skyformer, we first propose a method to apply Nyström method to approximate an asymmetric (and thus non-PSD) empirical kernel matrix B constructed with any PSD kernel φ(·, ·). Specifically, with two different n-by-p design matrices Q and K, its element bij from the i-th row and j-th column in B is equal to φ(qi,kj), where qi (resp. kj) is the i-th (resp. j-th) row in Q (resp. K). We remark this type of empirical kernel matrices involves the un-normalized attention score matrix A := SM(Q,K), and the empirical Gaussian kernel matrix C := κ(Q/p1/4,K/p1/4). Therefore this method leads to a low-rank approximation to the output of either self-attention D−1AV or Kernelized Attention CV . (D in self-attention can be obtained by computing A · 1, and thus a low-rank approximation to A also implies an approximation to D.) Computational details are stated as follows. To tackle the challenge of approximating a non-PSD matrix B, our first step is to complete the matrix into a PSD matrix B̄:
B̄ := φ (( Q K ) , ( Q K )) . (4)
Then we approximate B̄ with ˜̄B through
˜̄B = B̄S(ST B̄S)†ST B̄, (5)
where S is a 2n-by-d uniform sub-sampling matrix as defined in Definition 1. The final approximation will be given as
B̃ := (I,0) ˜̄B(0, I)T . (6)
The original matrix B can be well-approximated by B̃ due to the following inequality
‖B − B̃‖ = ‖(I,0)(B̄ − ˜̄B)(0, I)T ‖ ≤ ‖B̄ − ˜̄B‖, and thus we show our task of approximating the non-PSD matrix B boils down to well approximating the PSD matrix B̄.
Remark. The reason why B can be well approximated by a low-rank B̃ is that as an empirical kernel matrix the eigenvalues in B̄ usually decay fast, and thus there are many small eigenvalues in the long tail. In this case, theoretically a low-rank matrix (e.g. truncated singular value decomposition (SVD) of B̄) has enough potential to well approximate the original matrix B̄ (and B accordingly) in terms of spectral norm.
With the derivation above, we officially introduce our proposed Skyformer as an approximation to Kernelized Attention, which applies the modified Nyström method to the kernelized attention score matrix C. The next two subsections will continue our discussion on it, and respectively state the theoretical analysis of its approximation error and some details of its implementation in practice.

heading: 4.3 Error analysis of Skyformer

text:
As mentioned, an implicit advantage of using Kernelized Attention is that we can leverage the existing conclusions for kernel methods to analyze the theoretical properties of the model. In this subsection, we aim to provide some theoretical analysis of its approximation error.
We state a high probability bound on the size d of the sub-sampling matrix used in Skyformer to attain (ε, δ)-MA property for the kernelized attention score matrix C by the following theorem. We refer the readers to the proof in Appendix D to take a closer look at our claim that the matrix to be approximated should be PSD is a key to the theoretical guarantee of Nyström method. Theorem 2 (Adapted from Lemma 9 and Theorem 3 [Musco and Musco, 2017]). Consider the query, key, and value matrix Q,K,V ∈ Rn×p and two positive constants ε < 1, δ < 12 . For the empirical Gaussian kernel matrix C := κ(Q/p1/4,K/p1/4) defined above, we let λ := ε‖C‖ < ‖C̄‖, where C̄ is the completion of C (similar to B̄, constructed as substituting the Gaussian kernel with bandwidth p1/4 for the arbitrary kernel function φ in Equation (4)). We comment λ serves as the regularization coefficient as well as the approximation error bound. To ease the analysis, we specifically define the i-th diagonal element of C̄(C̄ + λI2n)−1 as leverage score `i,∀i = 1, . . . , 2n, and define their sum Tr ( C̄(C̄ + λI)−1 ) as the statistical dimension dstat, which increases with 1/ε as λ ∝ ε. Suppose S is a uniform sub-sampling matrix, and assume there exists a constant β ∈ (0, 1] such that β ≤ dstat2n`i ,∀i = 1, . . . , 2n. For the approximation matrix
˜̄C constructed with C̄,S as in Equation (5), there exists a constant C such that if
d ≥ C dstat β log n δ
then ˜̄C 4 C̄ 4 ˜̄C + λI with probability 1 − δ. Here 4 denotes the Loewner ordering: B 4 A means A−B is positive semidefinite. Furthermore, for our approximation C̃ in Equation (6) to the kernelized attention score C, we have
‖C̃ −C‖ ≤ λ = ε‖C‖.
This theorem implies the time and space complexity of our proposed approximation depends on the statistical dimension dstat. If we directly use the conclusion from Gaussian kernels, dstat should be Õ(1) (complexity modulo poly-log term) [Yang et al., 2017] due to the exponential eigenvalue decay rate of Gaussian kernels, which is comparable to the complexity of most other efficient transformers. However, different than the case in the classical kernel methods, the distribution of the query and key matrix Q and K changes during the training procedure, which may invalidate the conclusion about dstat. We leave the exact non-asymptotic analysis of the computational complexity for future work.

heading: 4.4 Workaround in implementation

text:
A potential limitation with the implementation of the proposed method lies in the tricky fact that the matrix inversion on GPU is much slower and numerically less stable than the same operation on CPU
due to the different back-end libraries in the two platforms. We attempt to circumvent the problem by adapting the strategy in Nyströmformer [Xiong et al., 2021] to our setting. Specifically, we use the matrix-product-based iterative method [Razavi et al.] for finding approximate inverses, instead of some division-based methods (such as the conjugate gradient method) which induces some instability in model training.
To apply the iterative method and inverse matrix M = ST C̄S, we need to satisfy its assumption [Razavi et al., Theorem 2] that ‖I −M‖ < 1. In practice, we instead pass the matrix D−1/2M (M + γI)D
−1/2 M as an input to the iterative method, where γ > 0 is a small constant and the diagonal
matrix DM is defined as diag ((M + γI)1). We give the following lemma to justify our practical usage of the method. The proof is deferred to Appendix E.
Lemma 3. Given a constant γ > 0, if matrices M is constructed as ST C̄S, and DM are defined as above, then all the singular values of D−1/2M (M + γI)D −1/2 M are within (0, 1), which implies that ‖I −D−1/2M (M + γI)D −1/2 M ‖ < 1.
We further comment that numerically an implicit risk of the Schulz-type iterative method we use is the unintended consequence of “zero fill-in". If we use some sparse kernels (e.g. test functions with bounded support) other than Gaussian kernels, the empirical kernel matrices are sparse while the approximate inverse will converge to a dense matrix, which increases the computational cost.

heading: 4.5 Empirical approximation evaluation

text:
Spectral norm, the maximum singular value of a matrix, is a computation-light indicator of matrix approximation performance. In this work, we compare the spectral norm of the difference between the outputs from attention functions and the output from vanilla self-attention with the same input.
We use the initialized and pretrained bert-base-cased models from Huggingface’s implementation [Wolf et al., 2019] . The input vector X is embedded from the tokenized raw text in Wikitext-2 dataset [Merity et al., 2017]. The query, key and value weight matrices in initialized or pretrained models transform input X into Q,K, V of different distributions. We compare the results with different sequence lengths and different numbers of features used in attention approximation methods. We set the number of features in the range of 24 to 28. More features usually require more computation resources.
Figure 1 shows the performance of the modified Nyström method on approximation error with regards to the number of features. We conclude that for Skyformer the approximation is significantly better with the increased number of features, while for other methods the gain is not obvious. The good performance of the modified Nyström method also validates our previous claim that the Nyström method is currently one of the most powerful methods in large-scale kernel machines acceleration.
Remark. Although in a single step the modified Nyström method in Section 4.2 can give low approximation error, we do not recommend directly applying it to the original self-attention. With some exploratory experiments on classification tasks, we find the variant suffers a more severe gradient explosion issue than usual transformers. We speculate that it is because the matrix ST ĀS (in the middle of Equation (6)) inherits the high condition number of the original attention score matrices A, while the derivative of matrix inverse ( ( A−1 )′ = −A−1A′A−1) further amplifies the condition number during backpropagation.

heading: 5 Experimental Results

text:
Tasks and Datasets. We evaluate the proposed methods on five classification tasks on LRA benchmark [Tay et al., 2020b], which focuses on model quality under long-context scenarios: ListOps [Nangia and Bowman, 2018], Text Classification on IMDb review dataset [Maas et al., 2011], Document Retrieval on AAN dataset [Radev et al., 2013], Pathfinder [Linsley et al., 2018], and Image Classification on CIFAR-10 [Krizhevsky et al., 2009]. The LRA benchmark covers diverse longsequence tasks in sequence length, task difficulty, and inspected model abilities. For example, ListOps and Pathfinder evaluate the abilities to capture the long-range hierarchical dependency and spatial dependency, respectively, which poses challenges for sparse attention pattern based methods. We
report the classification accuracy on the test set, training time, and peak memory usage during training for each task.
Baselines. Aside from the vanilla quadratic self-attention, we compare with Big Bird [Zaheer et al., 2020], Performer [Choromanski et al., 2020], Linformer [Wang et al., 2020], Nyströmformer [Xiong et al., 2021], Informer [Zhou et al., 2020], and Reformer [Kitaev et al., 2020]. Most methods are approximating the vanilla full attention for efficiency and thus are not expected to have better performance. As it is not realistic to exhaustively fine-tune all models and search for the best performance under limited computation resources, we instead only replace the self-attention module with the various attention methods and keep other experimental settings the same for fair comparisons.
Implementation Details. We conduct each experiment on one Tesla V100 SXM2 16GB. We use the LRA evaluation benchmark reimplemented in PyTorch by Xiong et al. [2021]. We use a 2-layer transformer model with 64 embedding dimension, 128 hidden dimension, 2 attention heads, and mean pooling for classification. Batch size is selected conditioned on the memory requirements of the standard self-attention method, which leads to 16 for Text Classification, 32 for ListOps, 16 for Document Retrieval, 128 for Pathfinder, and 256 for Image Classification. Learning rate is set to 1e − 4 for Text Classification, ListOps, and Image Classification, and 2e − 4 for Retrieval and Pathfinder. Each model on each task is trained for 50k steps, during which the best checkpoint with the highest accuracy on the development set will be saved for evaluation. For comparable computation complexity, we control the number of features to be 128 used in all methods (except Big Bird), under which setting the models will visit 128 ·n elements in the attention matrix. For numerical consistency, all experiment results are averaged across three runs with different random seeds.
We do not follow all settings in [Xiong et al., 2021] due to the hardware limitation. The compromises, such as approximation dimension and gradient accumulations steps, might bring performance differences comparing to results reported in [Xiong et al., 2021]. The training instability problem also helps explain the performance gap.
Results. The training process of the standard softmax-based method is unstable as observed in Figure 2: it takes more steps to reach the stationary distribution of its long-time limit, and it is more easily getting stuck in a local minimum. Runs with different random seeds may bring divergent performances, and probably leads to lower averaged scores. We have also tried directly approximating the self-attention method with the Nyström method and observed numerical instability during training.
Replacing the softmax structure with Gaussian kernel somehow alleviates this instability problem with boosted performance as shown in Table 1. However, the time and space requirement of Kernelized Attention is not significantly improved compared to the original version, which serves as the motivation to approximate Kernelized Self-Attention with Nyström method.
Though not necessarily the fastest, our proposed Skyformer can efficiently converge to the longtime limit with comparable general performance in classification accuracy (Table 1) and resource consumption (Table 2). The advantages over the standard self-attention are significant with consistently less training time and generally better performance. For example, Skyformer brings nearly 4
times speed-up on text classification and document retrieval while with 2.75% and 1.37% accuracy improvement over the standard self-attention.
Limitations. The applications of Skyformer might be limited to long sequence tasks because for small sequence length n the statistical dimension dstat might be close to n.
To make the claim above clear, we first reiterate that the efficiency of Skyformer is related to dstat. As implied by Theorem 2, the intrinsic difficulty of approximating a raw attention score matrix is concluded as dstat, which corresponds to the effective rank of matrix C̄. The complexity of Skyformer depends on the sub-sample size d (the size of the sub-sampling matrix S). A large dstat leads to a large d , and an inefficient application of the Nyström method.
The classical theory for statistical dimension only guarantees that dstat is small (compared to n) when n is large enough, and it is possible the statistical dimension associated with a short sequence might be even close to the sequence n. Therefore a large n serves as a condition to make the method work. Figure 1 empirically shows that our method performs better with larger n’s.

heading: 6 Conclusions and future work

text:
Motivated by the connection between kernel methods and self-attention, we introduce Kernelized Attention, which replaces the softmax structure in self-attention with a Gaussian kernel. We also propose Skyformer, which adapts the Nyström method to Kernelized Attention to improve its efficiency. We expect the new model can enjoy more stable training while inheriting the strong performance from self-attention. Extensive experiments verify our intuitions and show that both Kernelized Attention and its Nyström approximation variant have comparable accuracy to the original Transformer on the LRA benchmark.
Direct development of this work is the incorporation of further computation tricks in kernel methods, such as the local and global approximation for gram matrix [Snelson and Ghahramani, 2007] and the importance sampling in Nyström methods [Musco and Musco, 2017, Chen and Yang, 2021b,a]. Other related questions include the choice of the kernel other than the Gaussian kernel in our kernelized attention model. It is expected that for different tasks there will be specific kernels more proper than the original self-attention. The results in this work also shed new light on the design of the attention mechanism, which may benefit board downstream NLP tasks.
Acknowledgments and Disclosure of Funding
This research is based upon work in part supported by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via contract No. FA8650-17-C-9116, and U.S. DARPA KAIROS Program No. FA8750-19-2-1004. This work is also in part supported by NSF grant DMS-1810831. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of DARPA, ODNI, IARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.

title: Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method


summaries: This paper proposes Augmented Policy Cloning (APC) approach to improve the data efficiency of expert behavior by augmenting expert trajectory data with virtual, perturbed states and the expert actions in these virtual states. This method is simple by applying the existing image-based data augmentation method, but it can increase the data efficiency of policy cloning and transfer high-DoF behaviors. The proposed method outperforms BC and Naive ABC.

summaries: The paper explores a novel approach to using expert policies in RL context. The main idea is to enrich expert demonstrations by perturbing a state by Gaussian noise and querying expert policy to produce an action from the perturbed state. These additional state-action pairs are used in the policy learning algorithm. The authors show experimentally that the proposed Augmented Policy Cloning (APC) can significantly improve the quality of cloned policy, compared to Behaviour Cloning.

summaries: This paper proposed an data-augmentation technique -- "augmented policy cloning (APC)" to enable efficient policy learning from parametric experts. It achieve a high level efficiency in transferring knowledge from an expert to a student policy for high Degrees of Freedom environment. The augmentation is introduce not only for the states but also to the actions of underlying policy.

summaries: The paper proposes a policy cloning method that transfers an expert to a student. In the considered setting, the expert policy is accessible at all times, while there is no access to the environment to perform RL. The proposed method leverages data augmentation to successfully train the student policy.

summaries: This paper proposes to study the non-interactive imitation learning setting where access to the parametric expert policy is available. The paper proposes to augment the logged expert data’s states with noise, and then determine then sample corresponding actions from the expert’s policy. These perturbed states and corresponding action are then use to train the agent. They show this method outperforms Behavior cloning and an approach where the logged states are perturbed but the corresponding actions are kept fixed on imitation learning tasks and RL tasks where IL is used for warmstarting.

abstractText: We present a simple, yet powerful data-augmentation technique to enable data1 efficient learning from parametric experts. Whereas behavioral cloning refers to 2 learning from samples of an expert, we focus here on what we refer to as the policy 3 cloning setting which allows for offline queries of an expert or expert policy. This 4 setting arises naturally in a number of problems, especially as a component of 5 other algorithms. We achieve a very high level of data efficiency in transferring 6 behavior from an expert to a student policy for high Degrees of Freedom (DoF) 7 control problems using our augmented policy cloning (APC) approach, which 8 combines conventional image-based data augmentation to build invariance to 9 image perturbations with an expert-aware offline data augmentation approach that 10 induces appropriate feedback-sensitivity in a region around expert trajectories. We 11 show that our method increases data-efficiency of policy cloning, enabling transfer 12 of complex high-DoF behaviors from just a few trajectories, and we also show 13 benefits of our approach in the context of algorithms in which policy cloning is a 14 constituent part. 15

venue: arXiv preprint arXiv:2010.14274,

text:
We present a simple, yet powerful data-augmentation technique to enable data-1 efficient learning from parametric experts. Whereas behavioral cloning refers to2 learning from samples of an expert, we focus here on what we refer to as the policy3 cloning setting which allows for offline queries of an expert or expert policy. This4 setting arises naturally in a number of problems, especially as a component of5 other algorithms. We achieve a very high level of data efficiency in transferring6 behavior from an expert to a student policy for high Degrees of Freedom (DoF)7 control problems using our augmented policy cloning (APC) approach, which8 combines conventional image-based data augmentation to build invariance to9 image perturbations with an expert-aware offline data augmentation approach that10 induces appropriate feedback-sensitivity in a region around expert trajectories. We11 show that our method increases data-efficiency of policy cloning, enabling transfer12 of complex high-DoF behaviors from just a few trajectories, and we also show13 benefits of our approach in the context of algorithms in which policy cloning is a14 constituent part.15
1 Introduction16
In various control and reinforcement learning settings, there is a need to transfer behavior from an17 expert policy to a student policy. Broadly, when only samples from the expert policy are available, the18 standard approach is to employ a version of regression from states to actions. This class of approaches19 for producing a policy is known as behavioral cloning [Pomerleau, 1989, Michie and Sammut, 1996].20 Behavioral cloning is quite flexible and supports the setting where the expert trajectories come from a21 human teleoperating the relevant system directly, as well as various settings where the trajectories are22 sampled from other controllers, which themselves may have been trained or scripted. However, for23 any of the settings where the expert policy is actually available, rather than just samples from the24 expert, it is reasonable to suspect that sampling random rollouts from the expert policy followed by25 performing behavioral cloning is not the optimally efficient approach for transferring behavior from26 the expert to the student. Once a trajectory has been sampled via an expert rollout, there is actually27 additional information available that can be ascertained in the neighborhood of the trajectory, without28 having to perform an additional rollout, via the local feedback properties of the expert.29
In this work, we refer to this setting, where we want to transfer from an expert policy to a student30 policy, while assuming the expert policy can be queried, as policy cloning. Naturally, there is still31 often an incentive to reduce the total number of rollouts, which may require actually collecting data32 in an unsafe or costly fashion, especially for real-world control problems. As such, there is an aim to33 characterize any efficiency that can be gained in learning from small numbers of rollouts without as34 much concern for how many offline queries are required of the expert policy. If one has primarily35 encountered behavioral cloning in the context of learning from human demonstrations, policy cloning,36
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.
with an available expert policy may seem contrived. However, policy cloning naturally arises in many37 settings. For example, we may have multiple experts that we wish to consolidate into a single neural38 network policy or there may be memory considerations that motivate compressing a large expert39 network into a smaller model with similar behavior. Perhaps most natural are settings in which an40 expert policy is costly or slow to execute, for example due to running a compute intensive procedure41 such as model predictive control (MPC) on specialized hardware (e.g. GPU); in such settings, the42 aim is to transfer expert behavior to a parametric student policy that amortizes the cost. DAGGER is43 one well known approach for efficiently transferring behavior from an expert to a student under these44 kinds of constraints [Ross et al., 2011]. In a separate setting, the expert may be suboptimal and the45 student needs to learn from expert while also being able to exceed the expert performance, perhaps46 by continuing to learn from a task via RL. This problem has been described as kickstarting in one47 incarnation [Schmitt et al., 2018], but also can arise when learning from behavioral priors [Tirumala48 et al., 2020], [Galashov et al., 2019], as also happens, for example, in Distral [Teh et al., 2017].49
To improve data-efficiency in supervised settings generally, including in behavioral cloning settings,50 it is reasonable to consider data augmentation. Data augmentation refers to applying perturbations to51 a finite training dataset to effectively amplify its diversity, usually in the hopes of producing a model52 that is invariant to the class of perturbations performed. For example, in the well studied problem53 of object classification from single images, it is known that applying many kinds of perturbation54 should not affect the object label, so a model can be trained with many input perturbations all yielding55 the same output [Shorten and Khoshgoftaar, 2019]. This setting is fairly representative, with data56 augmentation usually intended to make the model “robust" to nuisance perturbations of the input.57 This class of image-perturbation has also been recently demonstrated to be effective in the context of58 control problems in the offline RL setting [Yarats et al., 2021, Laskin et al., 2020].59
Critically, for control problems it is not the case that the action should be invariant to the input state.60 Or rather, while it does make sense for a control policy to be invariant to certain classes of sensor61 noise, an important class of robustness is that the policy is appropriately feedback-responsive. This is62 to say that for small perturbations of the state of the control system, the optimal action is different63 in precisely the way that the expert implicitly knows. This has been recognized and exploited in64 previous research that has distilled feedback-control plans into controllers [Mordatch and Todorov,65 2014, Mordatch et al., 2015, Merel et al., 2019]. A similar intuition also underlies schemes which66 inject noise into the expert during rollouts to sample more comprehensively the space of how the67 expert recovers from perturbations [Laskey et al., 2017, Merel et al., 2019].68
In this work, we leverage this insight to develop a highly efficient policy cloning approach that69 makes use of both classes of data augmentation. For a high-DoF control problem that operates only70 from state (humanoid run task from DeepMind control suite [Tunyasuvunakool et al., 2020]), we71 demonstrate the feasibility of policy cloning that employs state-based data augmentation with expert72 querying to transfer the feedback-sensitive behavior of the expert in a region around a small number73 of rollouts. Then on a more difficult high-DoF control problem that involves both state-derived and74 egocentric image observations (humanoid running through corrdiors task from DeepMind control75 suite [Tunyasuvunakool et al., 2020]), we combine the state-based expert-aware data augmentation76 with a separate image augmentation intended to induce invariance to image perturbations. Essentially77 our expert-aware data augmentation involves applying random perturbations to the state-derived78 observations, and training the student to match the expert-queried optimal action at each perturbed79 state, thereby gaining considerable knowledge from the expert without performing excessive rollouts80 simply to cover the state space around existing trajectories. Our approach compares favorably to81 sensible baselines, including the naive approach of attempting to perform behavioral cloning with82 state perturbations, which seeks to induce invariance (as proposed in [Laskin et al., 2020]) rather than83 feedback-sensitivity to state-derived observations.84
In the presentation that follows, we will describe the problem setting (Section 2) as well as our85 approach (Section 3), describe the domains we employ and present our initial experiments (Section 4),86 show that our augmented policy cloning approach works well when used as a component of other87 algorithms like DAGGER and kickstarting (Section 5), and finally close with a discussion (Section 6).88
2 Problem description89
2.1 Expert-driven learning90
We start by introducing a notion of expert-driven learning that will be used throughout the paper.91 At first, we present a general form of the expert-driven objective and then introduce a few concrete92 examples. We consider a standard Reinforcement Learning (RL) problem. We present the domain as93 an MDP with continuous states for simplicity, however the problem definition is similar for a POMDP94 with observations derived from the state. Formally, we describe the MDP in terms of a continuous state95 space S ∈ Rn (for some n > 0), an action space A, transition dynamics p(s′|s, a) : S ×A → p(S),96 and a reward function r : S × A → R. Let Π be a set of parametric policies, i.e. of mappings97 πθ : S → p(A) from the state space S to the probability distributions over actions A, where θ ∈ Rm98 for some m > 0. For simplicity of the notation, we omit the parameter in front of the policy, i.e.99 π = πθ and optimizing over the set of policies would be equivalent to the optimizing over a set of100 parameters. A reinforcement learning problem consists in finding such a policy π that it maximizes101 the expected discounted future reward:102
J(π) = Ep(τ) [∑ t γtr(at|st) ] , (1)
where p(τ) = p(s0) ∏ t p(at|st)p(st+1|st, at) is a trajectory distribution. We assume the existence103 of an expert policy πE(a|s). This policy could be used to simplify the learning of a new policy on the104 same problem. Formally, we construct a new learning objective which aims to maximize the expected105 reward of the problem in hand as well as to clone the expert policy:106
J(π, πE) = αJ(π)− λD(π, πE), (2) where D is some measure of distance of π from πE and α ≥ 0, λ ≥ 0 are parameters measuring107 importance of both objectives. In most of the applications, α ∈ {0, 1} and λ ≥ 0 represents a relative108 importance of cloning an expert policy with respect to the RL objective.109
2.2 Behavioral cloning (BC)110
One important instance of the objective (2) with α = 0, λ = 1 is behavioral cloning. In this case, the111 measure of distance is defined as:112
DBC(π, πE) = −E(a,s)∈BE [log π(a|s)] (3) Here, BE = {(si, ai), i = 1, . . . , N}, N > 0 is a fixed dataset containing expert data. Minimizing113 the objective (3) is be equivalent to maximizing the likelihood of the expert data under the policy π.114 The action in eqn. (3) can be replaced by πE(s) for deterministic policies or by the mean or the mode115 for stochastic policies (e.g., by the mean µE(s) for Gaussian policies πE(·|s) = N (µE(s), σE(s))).116
2.3 DAGGER117
Performance of Behavioral Cloning (BC) can be limited due to the fixed dataset, since the resulting118 policy may fail to generalize to states outside the training distribution. A different approach, known119 in the literature as DAGGER [Ross et al., 2011] was proposed to overcome this limitation. In this120 setting, the expert is queried in states visited by the student, thus reducing distribution shift. In our121 notation, this corresponds to α = 0, λ = 1 in eqn. (2) and the measure of distance is defined as:122
DDAGGER(π, πE) = −Epβ(τ)[log π(a ′ t|st)], (4)
where pβ(τ), β ∈ [0, 1] is a trajectory distribution where actions are sampled according to the mixture123 policy between a student and an expert:124
p(a|s) = βπ̃(a|s) + (1− β)πE(a|s), (5) The action a′t in eqn. (4) is resampled from the expert policy for the state st, i.e., a ′ t ∼ πE(·|st).125 As in Section 2.2, for stochastic experts this action can be replaced by the mean or mode of the126 distribution in eqn. (4). The policy π̃(a|s) corresponds to a frozen version of student policy π so127 that the gradient ∇πDDAGGER(π, πE) is not taken with respect to p(a|s). Note that even though, in128 eqn. (4) we collect data from the environment, the setting nevertheless corresponds to pure imitation129 learning since expected reward is not directly maximized.130
2.4 Kickstarting131
In eqn. (2), we combine both maximization of expected task reward and minimization of distance to132 the expert. In literature, it is known as Kickstarting [Schmitt et al., 2018]. In this case, in the objective133 from eqn. (2), α = 1, and λ ≥ 0. As the measure of distance, we use the cross-entropy from expert134 to a student, similarly to [Schmitt et al., 2018]:135
J(π, πE) = J(π)− λEp(τ) [ −EπE(a|s) log π(a|s) ] (6)
where p(τ) is a trajectory distribution, where actions are sampled according to the student policy136 π(·|s). Usually, in the Kickstarting setting, the expert is sub-optimal and the goal is to train a policy137 that eventually outperforms the expert. Thus, it is customary to reduce λ over the course of training.138 Yet, for simplicity, in our experiments we keep this coefficient fixed.139
3 Augmented policy cloning140
The previous section has demonstrated that the objective corresponding to the cloning behavior from141 the parametric expert policy could arise in multiple scenarios. In this section we propose a new and142 simple method which can significantly improve the data efficiency of the approaches described in143 Section 2. We explain the basic idea for BC, but its generalization to other expert-driven learning144 approaches described in Section 2 is straightforward. In Section 5 we show results for these problems.145
When optimizing the objective (3), for every state s ∈ DE from the expert trajectories dataset, we146 consider a small Gaussian state perturbation:147
δs ∼ N (0, σ2s) (7)
which produces a new virtual state:148 s′ = s+ δs (8)
Then, for this state we query the expert and obtain a new action149
a′ ∼ πE(·|s+ δs) (9)
We then augment the dataset DE with these new pairs of virtual states and actions. More explicitly150 the idea can be expressed in terms of the following objective:151
D(π, πE)APC = E(a,s)∈BE [log π(a|s) + Eδs∼N (0,σ2s),a′∼πE(·|s+δs) log π(a ′|s+ δs)] (10)
We call this approach Augmented Policy Cloning (APC) as it queries the expert policy to augment152 the training data. This approach is different from a naive data-augmentation technique, where a new153 state would be generated, but associated with the original action (and not a new one). It therefore154 allows to build policies which are feedback-responsive with respect to the expert. We formulate APC155 algorithm for BC in Algorithm 1.156
4 Core Results: Evaluation of Augmented Policy Cloning157
4.1 Domains158
To study how our method performs on complex control domains, we consider two complex, high-DoF159 continuous control tasks involving control of a physically simulated humanoid body. Both domains160 are implemented using the MuJoCo physics engine [Todorov et al., 2012] and are available in the161 dm_control repository [Tunyasuvunakool et al., 2020]. The first task is the standard control suite162 Run task, where the Humanoid body needs to run at a target speed and observations are based on163 proprioception. The second task is the Walls task which requires the same Humanoid body to run164 along a corridor and avoid walls, using both proprioception and egocentric vision as observations.165 Both of these problems are rather challenging insofar as they require stabilization and locomotion166 control of a relatively complex humanoid body with 21 actuated DoFs, in one case using vision to167 guide the movement. Note these environments are related to the domains that have been proposed168 for use in offline RL benchmarks [Gulcehre et al., 2020]; however, the experiments we perform in169 this work require availability of the expert policy, so we do not use offline data, but instead train170 new experts and perform experiments in the very low data regime. For more details, please refer to171 Section 1.1 in Supplementary Material.172
Algorithm 1 Augmented Policy Cloning (APC) Parametric student policy: πθ Initial parameters: θ0 expert policy: πE Dataset BE = {(si, ai), i = 1, . . . , N}, N > 0 of expert state-action pairs State perturbation noise σs Learning rate α Number of augmented samples: M Number of gradient updates: K Size of a batch: L for k=1,. . . ,K do
Sample a batch of pairs {(ai, si)}Li=1 ∼ BE For each state si, sample M perturbations δsj ∼ N (0, σs), j = 1, . . . ,M Construct M virtual states s′i,j = si + δsj , i = 1, . . . , L, j = 1, . . . ,M Resample new actions from expert a′i,j ∼ πE(·|s′i,j) For Gaussian experts, the action ai = µE(si) and the new actions are a′i,j = µE(s ′ i,j) Compute the empirical negative log-likelihood: L = − [ log πθk(ai|si) + 1M ∑M j=1 log πθk(a ′ i,j |s′i,j)
] Update the parameters θk+1 = θk − α∇θL
end for
For each task, we train expert policies to convergence using the MPO algorithm Abdolmaleki et al.173 [2018]. Since the expert policy essentially saturates task performance, for each task, we keep three174 partially trained experts such that we can assess the ability of the kickstarting approach to outperform175 sub-optimal experts. We refer to the different experts as Low, achieving approximatively 25 % of the176 optimal policy reward, Medium, achieving around 50 % of the performance and High, corresponding177 to the converged policy. Each expert is represented by a Gaussian policy. For more details, please178 refer to Section 1.2 in Supplementary Material.179
4.2 Applying Augmented Policy Cloning180
First, we evaluate the performance of APC in fitting a fixed dataset of expert trajectories. In order181 to study the data efficiency of the method, we construct datasets containing different numbers of182 expert trajectories. The expert policies are represented by conditionally Gaussian distributions,183 i.e. πE(·|s) = N (µE(s), σ(s)). Thus, to assess the robustness of our method to expert noise we184 produce trajectories using the experts’ mean but adding different levels of (homoscedastic) zero-mean185 Gaussian noise σE .186 a ∼ N (µE(s), σE) Note that in addition to policy noise σE which is introduced when sampling trajectories, initial pose187 and environment layout (for the Walls task) are also sampled randomly for each episode. We consider188 4 levels of expert policy noise: Deterministic, which uses the Gaussian mean for the action, Low,189 with σE = 0.2, Medium σE = 0.5 and High σE = 1.0.190
For the APC method, we rely on Algorithm 1. For baselines, we consider BC algorithm from eqn. (3)191 as well as a simple modification of BC, where we apply, similar to APC, state perturbations as in192 eqn. (7) and eqn. (8), but we do not produce a new action from the expert. We call this approach193 Naive Augmented Behavior Cloning (Naive ABC) which essentially corresponds to robustification194 of the student policies with respect to state perturbation and is similar in spirit to standard data-195 augmentation approaches. For vision-based tasks, we consider random crop augmentations of size196 48x48 (downsampled from the input image of 64x64), similar to Laskin et al. [2020]. When the197 image augmentations are used we add "with image" to the method name. On top of that, we consider198 a variant, where only image augmentation is used, which we call Naive ABC (image only). For199 all methods, as an action in the objective from eqn. (3), we use an expert mean µE(s). We train200 all approaches to convergence (300K learning iterations on Walls and 13M learning iterations on201 Run). Each learning iteration corresponds applying gradients to 64 trajectories, each containing202 10 time steps. After each learning iteration, we evaluate the policy on both a validation set (50203 random instances of the environment) and a test set (150 random instances of the environments).204
We apply early-stopping based on the validation set performance to select the best model and205 report corresponding performance on the test set. For more details, please refer to Section 1.3 in206 Supplementary Material. As an additional evaluation, we test robustness of the obtained policies to a207 fixed amount of noise during execution. For a learned student policy π(·|s) = N (µ(s), σ(s)), we208 evaluate it by executing an action:209 a ∼ N (µ(s), σ), where σ is the fixed amount of student noise. We consider similar noise magnitudes as for the210 expert. For APC and Naive ABC, we sweep over state perturbation noise levels and choose the ones211 performing the best on the validation set. For APC, we use σs = 0.1 for Run and σs = 1.0 for Walls.212 For Naive ABC, we use σs = 0.001 for Run and σs = 0.01 for Walls. The ablation experiments213 over noise levels for APC and Naive ABC are presented in Section 2 of the Supplementary Material214 (Figure 1 and Figure 2).215
The first of results, in Figure 1 (left) demonstrates the increased data efficiency of APC over BC and216 Naive ABC in terms of number of trajectories. The noise level of the expert and student are fixed to217 Low for the ease of comparison. We also see that Naive ABC performs similarly to BC. To further218 push the limits of data efficiency, we conducted a variant where a dataset contains only 1 full trajectory219 (1000 timesteps for Run and around 2k timesteps for Walls) along with multiple short trajectories220 (200 time-step only). This dramatically reduces the amount of expert data available to learn from.221 However, we hypothesise that in the environments considered, much of the diversity of the trajectories222 arises due to initial state variation. This setting might arise in domains where execution is costly,223 such as robotics applications. In such setting we might have a few longer trajectories along with a224 patchwork of shorter trajectories covering more diverse parts of the state space. The results for this225 experiment are given in figure 1 (right). Again, we see that APC is significantly more efficient than226 BC and Naive ABC. Interesting to note that APC picks up quite a high performance after observing227
10 (1 full and 9 short) trajectories for Run task and 100 (1 full and 99 short) trajectories for the Walls228 task. In Section 3 in Supplementary material, we provide additional results for Walls task when we229 use image-based perturbations.230
In the next experiment, in order to understand how robust our method to noise, we study the impact231 of different levels of student and expert noises on performance. For each run, we use a dataset of232 100 trajectories. The results are given in figure 2, where each column corresponds to a different233 level of expert noise, and the X-axis represents different levels of student noise. At first, we observe234 that APC is consistently more robust than BC and Naive ABC for any level of expert and student235 noise. On top of that, we can notice that for any fixed level of expert noise, the performance degrades236 when a student noise increases. Finally, we see that for higher noise levels of expert, the learned237 student performs better in the high noise regime. It is consistent with the intuition - training on noisy238 trajectories leads to a more robust policy. Overall, APC leads the most robust policy.239
5 Additional Results: Augmented Policy Cloning as a subroutine240
5.1 DAGGER with data augmentation241
As described in Section 2.3, DAGGER [Ross et al., 2011] is a more sophisticated approach where242 data is collected from the real environment by executing a policy from eqn. (5), which is a mixture243 between a student and an expert. In this section we study how data augmentation approaches affect244 the data efficiency of the DAGGER algorithm.245
We consider similar baselines for both tasks as in the previous section. For an expert policy that has246 been pre-trained via MPO [Abdolmaleki et al., 2018], we perform online rollouts for two values of247 the expert-student mixing coefficient, β = 0 and β = 0.3 (see eqn. 5). Since both student and expert248 are Gaussian distributions, instead of using a log π in eqn. (4), we could use a state-conditional cross249 entropy from an expert to a student,H[πE(·|s)||π(·|s)]. Empirically, we found that it worked better250 than using log π. We demonstrate a comparison in Section 4 in Supplementary Material. We run251 the experiments in a data-restricted setup such that for every collected trajectory (10 time-steps), we252 apply 10 gradient steps, using a replay-buffer to store the past experience. Additional experimental253 details are given in Section 1.4 in Supplementary Material. Results are shown in Figure 3. We see254 that APC and its vision variant outperform BC and Naive ABC similarly to the behavior cloning255 experiments. While we observe that image augmentation can help, we see that the primary advantage256 comes from the state-based augmentation for APC. For the Run task, we observe that all DAGGER257 methods achieve slightly lower performance than an expert policy. We speculate that this is due to258 insufficient coverage of the state space during training.259
5.2 Kickstarting with data augmentation260
1.0 e9 2.0 e9 3.0 e9 4.0 e9
0
250
500
750
Ep iso
de re
wa rd
Run, High
Method APC BC Naive ABC APC with image Naive ABC (image only) Naive ABC with image MPO
1.0 e9 2.0 e9 3.0 e9 4.0 e9
Run, Medium
1.0 e9 2.0 e9 3.0 e9 4.0 e9
Run, Low
0.0 e9
0.5 e9
1.0 e9
0
250
500
750
Walls, High
0.0 e9
0.5 e9
1.0 e9
Walls, Medium
0.0 e9
0.5 e9
1.0 e9
Environment steps
Walls, Low
Figure 4: Kickstarting results. On the X-axis we show the number of environment steps. On the Y-axis we report averaged across 3 seeds episodic reward achieved by the student. We report confidence intervals in the shaded areas. For Run task, the confidence intervals are very small and are not visible. Each row indicates a task, whereas a column corresponds to the expert type. Dashed black line shows the expert performance.
A similar in spirit approach is kickstarting Schmitt et al. [2018], where we solve an RL task as well261 as cloning the expert policy. Similarly to previous section, we apply APC in kickstarting on the262 cross entropy term in eqn. (6). We use 3 types of expert policy as described in Section 4.1. We run263 experiments using a distributed setup with 64 acting policies and 1 learner, querying the batches264 of trajectories (of size 10) from a replay buffer. On top of running BC methods, we also report the265 performance of MPO Abdolmaleki et al. [2018] learning from scratch on the task of interest. All266 details are given in Section 1.5 in Supplementary Material. The results are given in figure 4.267
We observe that APC performs better than Naive ABC on Run task and similarly on Walls task. Both268 approaches perform better than BC and learning from scratch. We hypothesise that the reason of not269 seeing a consistent advantage could be due to two factors. As we are in a high-data regime, since270 there is no limit on relative acting / learning ratio, and acting policies are not restricted to collect271 trajectories, it is unclear whether data-augmentation should help. In addition, we use reward signal272 which makes the impact of expert cloning less important. Note that the resulting agent is less data273 efficient in these experiments; this is because we do not control the relative ratio between acting274 and learning (i.e., no rate-limiting on the learner, due to instability of kickstarting experiments when275 rate-limiting was explored). Furthermore, unlike in kickstarting Schmitt et al. [2018], we do not276 use an annealing schedule of λ to make the experiments simpler, but we still observe that a fixed277 coefficient helps to kickstart an experiment and outperform an expert policy. On top of that, we see278 that image-based augmentation have less of impact in this setting.279
6 Discussion280
Many expert-driven learning approaches actually have access to an expert that can be queried;281 however, this opportunity is rarely exploited fully. In this work we demonstrated a general scheme for282 more efficient transfer of expert behavior by augmenting expert trajectory data with virtual, perturbed283 states as well as the expert actions in these virtual states. This data augmentation technique is widely284 applicable and we demonstrated that it improves data efficiency when used in place of behavioral285 cloning both in the offline setting or when behavioral cloning is used as a step within DAGGER or286 kickstarting.287
Critically, data efficiency is generally very important in realistic applications, where new data288 acquisition cost could be high. In particular, settings involving deployment of policies in the real289 world, such as robotics applications, may benefit from an ability to efficiently transfer expert policy290 behavior from one neural network to another (for compression or execution speed reasons), or to291 combine behavior from multiple experts into a single neural network. While overall, we consider the292 present work to be fairly basic research with limited ethical impact, insofar as our approach decreases293 the amount of data which needs to be collected through processes which could potentially be unsafe294 or costly, there is a potential positive social value.295
The limitations of our approach consist in the reliance on the ability to query expert policy for the296 perturbed states which reduces the amount of applications where the method could be used. Another297 limitation is the reliance on the continuous state spaces. In discrete state spaces, it is unclear whether298 a small perturbation in state would result in a valid action from an expert.299
In future work, we plan to explore how our proposed augmentation technique can be leveraged in the300 context of KL-regularized RL with behavior priors.301
References302 Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin303 Riedmiller. Maximum a posteriori policy optimisation, 2018.304
Alexandre Galashov, Siddhant M. Jayakumar, Leonard Hasenclever, Dhruva Tirumala, Jonathan305 Schwarz, Guillaume Desjardins, Wojciech M. Czarnecki, Yee Whye Teh, Razvan Pascanu, and306 Nicolas Heess. Information asymmetry in kl-regularized rl, 2019.307
Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna,308 Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. Rl unplugged:309 A collection of benchmarks for offline reinforcement learning. Advances in Neural Information310 Processing Systems, 33, 2020.311
Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for312 robust imitation learning. In Conference on robot learning, pages 143–156. PMLR, 2017.313
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.314 Reinforcement learning with augmented data, 2020.315
Josh Merel, Leonard Hasenclever, Alexandre Galashov, Arun Ahuja, Vu Pham, Greg Wayne,316 Yee Whye Teh, and Nicolas Heess. Neural probabilistic motor primitives for humanoid con-317 trol, 2019.318
Donald Michie and Claude Sammut. Behavioural Clones and Cognitive Skill Models, page 387–395.319 Oxford University Press, Inc., USA, 1996. ISBN 019853860X.320
Igor Mordatch and Emo Todorov. Combining the benefits of function approximation and trajectory321 optimization. In Robotics: Science and Systems, volume 4, 2014.322
Igor Mordatch, Kendall Lowrey, Galen Andrew, Zoran Popovic, and Emanuel V Todorov. Interactive323 control of diverse complex characters with neural networks. Advances in Neural Information324 Processing Systems, 28:3132–3140, 2015.325
Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Technical report,326 Carnegie-Mellon, 1989.327
Stephane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell. A reduction of imitation learning and328 structured prediction to no-regret online learning, 2011.329
Simon Schmitt, Jonathan J. Hudson, Augustin Zidek, Simon Osindero, Carl Doersch, Wojciech M.330 Czarnecki, Joel Z. Leibo, Heinrich Kuttler, Andrew Zisserman, Karen Simonyan, and S. M. Ali331 Eslami. Kickstarting deep reinforcement learning, 2018.332
Connor Shorten and Taghi Khoshgoftaar. A survey on image data augmentation for deep learning.333 Journal of Big Data, 6, 07 2019. doi: 10.1186/s40537-019-0197-0.334
Yee Whye Teh, Victor Bapst, Wojciech Marian Czarnecki, John Quan, James Kirkpatrick, Raia335 Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning.336 arXiv preprint arXiv:1707.04175, 2017.337
Dhruva Tirumala, Alexandre Galashov, Hyeonwoo Noh, Leonard Hasenclever, Razvan Pascanu,338 Jonathan Schwarz, Guillaume Desjardins, Wojciech Marian Czarnecki, Arun Ahuja, Yee Whye339 Teh, et al. Behavior priors for efficient reinforcement learning. arXiv preprint arXiv:2010.14274,340 2020.341
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control.342 In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026–5033.343 IEEE, 2012.344
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom345 Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for346 continuous control. Software Impacts, 6:100022, 2020.347
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep348 reinforcement learning from pixels. In 9th International Conference on Learning Representations,349 ICLR, volume 2021, 2021.350
7 NeuRIPS 2021 checklist351
• Do the main claims made in the abstract and introduction accurately reflect the pa-352 per’s contributions and scope?353 Yes, we believe that the claims in the abstract and introduction are accurately reflected by354 the paper’s contributions and scope.355
• Have you read the ethics review guidelines and ensured that your paper conforms to356 them?357 Yes.358
• Did you discuss any potential negative societal impacts of your work?359 Yes. We included this in the Discussion section.360
• Did you describe the limitations of your work?361 Yes, we discuss these in the "Discussion" section.362 • Did you state the full set of assumptions of all theoretical results363
Not applicable.364 • Did you include complete proofs of all theoretical results365
Not applicable.366 • Did you include the code, data, and instructions needed to reproduce the main exper-367 imental results ?368 We included complete details in the supplementary material which could be used to reproduce369 the experimental results.370
• Did you specify all the training details?371 Yes, in the main paper and supplementary material.372 • Did you report error bars?373
We did. For every plot we reported the mean and standard deviations averaged across a374 number of independent random evaluations.375
• Did you include the amount of compute and the type of resources used?376 Yes, we did, in the supplementary material.377 • If your work uses existing assests, did you cite the creators?378
Yes. We use DeepMind control suite tasks and we cite the appropriate publications.379 • Did you mention the license of the assets?380
We did not explicitly mention the licence as it is clear from the reference.381 • Did you include any new assets either in the supplementary material os as a URL?382
No.383 • Did you discuss whether and how consent was obtained from people whose data you’re384 using / curating?385 Since the assets and references are in the public access, no consent was required.386
• Did you discuss whether the data you are using / curating contains personally identifi-387 able information or offensive content?388 The assets which we used did not include any personally identifiable information nor389 offensive content.390
• Did you include the full text of instructions given to participants and screenshots, if391 applicable?392 Not applicable.393
• Did you describe any potential participant risks, with links to Institutional Review394 Board (IRB) approvals, if applicable395 Not applicable.396
• Did you include the estimated hourly wage paid to participants and the total amount397 spent on participant compensation?398 Not applicable.399

title: Data augmentation for efficient learning from parametric experts