Question: Represent the scientific article to find a short TLDR!

Question: State-of-the-art performances on language comprehension tasks are achieved by huge language models pre-trained on massive unlabeled text corpora, with very light subsequent fine-tuning in a task-specific supervised manner. It seems the pre-training procedure learns a very good common initialization for further training on various natural language understanding tasks, such that only few steps need to be taken in the parameter space to learn each task. In this work, using Bidirectional Encoder Representations from Transformers (BERT) as an example, we verify this hypothesis by showing that task-specific fine-tuned language models are highly close in parameter space to the pre-trained one. Taking advantage of such observations, we further show that the fine-tuned versions of these huge models, having on the order of $10^8$ floating-point parameters, can be made very computationally efficient. First, fine-tuning only a fraction of critical layers suffices. Second, fine-tuning can be adequately performed by learning a binary multiplicative mask on pre-trained weights, \textit{i.e.} by parameter-sparsification. As a , with a single effort, we achieve three desired outcomes: learning to perform specific tasks, saving memory by storing only binary masks of certain layers for each task, and saving compute on appropriate hardware by performing sparse operations with model parameters. One very puzzling fact about overparameterized deep neural networks is that sheer increases in dimensionality of the parameter space seldom make stochastic gradient-based optimization more difficult. Given an effective network architecture reflecting proper inductive biases, deeper and/or wider networks take just about the same, if not a lower, number of training iterations to converge, a number often by orders of magnitude smaller than the dimensionality of the parameter space. For example, ResNet-18 (parameter count 11.7M) and ResNet-152 (parameter count 60.2M) both train to converge, at similar convergence rates, in no more than 600K iterations on Imagenet . Meaningful optimization seems to happen in only a very low-dimensional parameter subspace, viz. the span of those relatively few weight updates, with its dimensionality not ostensibly scaling with the model size. In other words, the network seems already perfectly converged along most of the parameter dimensions at initialization, suggesting that training only marginally alters a high-dimensional parameter configuration. This phenomenon is epitomized in fine-tuning of pre-trained models. Pre-training is a, often unsupervised, learning procedure that yields a good common initialization for further supervised learning of various downstream tasks. The better a pre-trained model is, the fewer iterations are required on average to fine-tune it to perform specific tasks, ing in fine-tuned models hypothetically closer 1 to the pre-trained one in parameter space. However, better pre-trained models are, almost always, larger models , and nowhere is this trend more prominent than recent pretrained language models that achieved state-of-the-art natural language understanding performance, e.g. GPT-2 has 1.5B parameters. Thus, a problem naturally arises hand-in-hand with an obvious hint to its solution: as pre-trained models get larger, on the one hand, computation of each fine-tuned model becomes more expensive in terms of both memory and compute for inference, while on the other hand, greater closeness between the pre-trained and fine-tuned models in the parameter space prescribes a higher degree of computational redundancy that could be potentially avoided. Additionally, there might exist more computationally efficient fine-tuned networks that are not necessarily close to, but cheaply attainable from, the pre-trained parameters, which are shared across all tasks. In this study, we seek to address these questions, using Bidirectional Encoder Representations from Transformers (BERT) and the General Language Understanding Evaluation (GLUE) benchmark tasks as a working example. We first found that the fine-tuned and pre-trained parameters are both L 1 -close and angular-close in parameter space, consistent with the small number of fine-tuning iterations separating them. Next, we demonstrated that there also exist good fine-tuned models that are L 0 -close (i.e. having a small number of different components) to the pre-trained one. Further, we showed that there exist good fine-tuned parameters that are L 0 -small (i.e. sparse, or having a large fraction of zero components). Finally, we successfully found fine-tuned language models that are both L 0 -small and L 0 -close to the pre-trained models. We remark the practical implications of these constraints. By forcing fine-tuned parameters to be L 0 -close to the pre-trained ones, one only needs to store a small number of different weights per task, in addition to the common pre-trained weights, substantially saving parameter memory. By forcing fine-tuned parameters to be sparse, one potentially saves memory and compute, provided proper hardware acceleration of sparse linear algebraic operations. Surprisingly, our findings also reveal an abundance of good task-specific parameter configurations within a sparse L 0 -vicinity of large pre-trained language models like BERT: a specific task can be learned by simply masking anywhere between 1% to 40% of the pre-trained weights to zero. See Figure 1 for an explanation of the L 0 -and sparse L 0 -vicinities. Figure 1: An illustration of the L 0 -vicinity and the sparse L 0 -vicinity of a pre-trained parameter in a three-dimensional parameter space. The L 0 -vicinity is continuous and contains parameters that are L 0 -close, whereas the sparse L 0 -vicinity is a discrete subset of L 0 -close parameters that are also L 0 -small. Our search for L 0 -close fine-tuning solutions is motivated by the observation that sensitivities of the optimization objective to different layers in a network are highly variable . trained fine-grain sparse connectivity patterns over randomly initialized network parameters, termed supermasks, suggesting a similar and complementary role model sparsification plays to gradient-based learning of the objective. This is also related to network architecture search (NAS). The most similar study to ours is piggyback and its variants , where in a multi-task visual object classification scenario, the authors trained taskspecific binary masks on top of a shared set of pre-trained parameters. In this work, we not only applied similar techniques further to larger pre-trained language models, but also studied the trade- off between L 0 -closeness and sparseness in a systematic way. Also, note that randomly generated high-dimensional masks can also support multi-task learning, e.g.. A large body of literature is concerned with sparsification of large networks for efficient inference. Here, we employed iterative pruning during fine-tuning to produce highperformance fine-grain sparse models. To enforce parameter sparsity differentiably in combination with the L 0 -closeness constraint, instead of principled approaches to imposing L 0 -regularization , we used simpler straight-through estimator, much like binary quantization techniques ; note that this is also used by and. Consider a pre-trained network F θ: x → F (x; θ), parameterized by θ, noted as subscript for convenience. The fine-tuning procedure to perform a task t ∈ T can be described as a supervised training procedure of model G task-specific last layer unique to task t, and • denotes function composition. In the case of BERT, we have a stack of modules among which E is the embedding layer, P a final pooling layer and each B is a transformer block where collects all the learnable parameter matrices in the block. A(·, ·, ·) represents scaled dot-product attention , DO(·) dropout, LN (·) layer normalization and GeLU(·) the Gaussian error linear unit activation function . We experimented with the BERT BASE model , for which L = 12, with total parameter count of 109M 2. See Table 1 for additional taskspecific parameter counts, all 5 orders of magnitude smaller than the total parameter count. Optimization of them alone fails to fine-tune (see Appendix A). The GLUE benchmark is a collection of diverse natural language understanding tasks . We fine-tune on these tasks and report the evaluation performances. We exclude the problematic WNLI set 3. F1 scores are reported for QQP and MRPC, Spearman correlations are reported for STS-B, and accuracy scores are reported for all other tasks. For all fine-tuning procedures, we use the exact hyperparameters as described in the original paper unless specified otherwise, with additional constraints described as follows. No constraints are imposed on task-specific last layers G (t). L 0 -close fine-tuning To search for fine-tuned solutions that are L 0 -close to the pre-trained parameters, we selectively fix certain parameter matrices at pre-trained values and perform fine-tuning optimization on a lower-dimensional parameter space. Sparse (L 0 -small) fine-tuning We use iterative pruning during fine-tuning to produce sparse models. Pruning is based on weight magnitudes in each layer and is performed periodically during fine-tuning with sparsity gradually increasing from 0% to a final level according to a cubic schedule. Iterative pruning successfully ensures that parameters are L 0 -small (see Appendix B). Supermask training as fine-tuning (sparse and L 0 -close) In order to search for fine-tuned networks that are both sparse and L 0 -close to the pre-trained one, we reparameterize the model by a multiplicative binary mask θ =θ µ, whereθ is the pre-trained parameters, and µ ∈ {0, 1} N the mask, N being the dimensionality of the parameter space and the Hadamard product. If learning is purely through optimizing the mask µ while holdingθ constant, the mask is called a supermask . Since µ is discrete-valued and thus not differentiable, we reparame- where Bern(p) denotes an element-wise independent Bernoulli sampler with probability p, and σ(·) the sigmoid function, applied element-wise on ν ∈ R N, the continuous mask parameter. We treat gradient backpropagation through µ as a straight-through estimator, similar to the techniques used in;. Same fine-tuning hyperparameters were used except for the learning rate (see Appendix C). Control over the final sparsity is exerted by initialization of µ for fine-tuning. We initialize ν according to a soft magnitude-based pruning mask: a fraction of small-magnitude values are initialized to ν = −5 and the rest to ν = 5. We found that the initial sparsity directly controls the final sparsity (see Appendix D), allowing us to produce masks with sparsity levels ranging from 1% to 89%. We observe that the original fine-tuning procedures for GLUE tasks all take 10 2 to 10 4 parameter update steps (Table 1), negligible compared to the dimensionality of the parameter space, viz. 10 8. Thus, we first questioned whether fine-tuned parameters are indeed close to the pre-trained ones in parameter space. We measured the L 1 -distances, i.e. L 1 -norm of parameter difference, and angular distances (Table 2). Specifically, we inspect the weight matrices in self-attention layers, of size 768× 768 where 768 is the hidden state dimension. We report the minimum and maximum values across GLUE tasks: MNLI showed the largest values, and RTE showed the smallest values. Evidently, we see a significantly higher L 1 -and angular-closeness between fine-tuned and pre-trained parameters as compared to the expected distance between two independent random initializations. This suggests that, during the course of fine-tuning, the very few model parameter updates traversed a very short distance in the parameter space. Comparing the parameter distance across GLUE tasks, we find that it scales with the number of fine-tuning iterations (see Appendix E). Further, we inspect the closeness in parameter subspaces for each layer. We found that, though all layers change very little during fine-tuning, there is nevertheless a high degree of variability across different parameter matrices (Figure 2). Blocks deeper in the encoder stack are less L 1 -close but more angular-close than shallower ones. In all self-attention modules, value and output projection Next, inspired by the high degree of variability in each layer's parameter change during fine-tuning, we ask whether effective fine-tuning can be achieved by optimizing only a fraction of layers while having others fixed at pre-trained values, ing in fine-tuned models L 0 -close in parameter space. Our suggest this is indeed feasible (Table 3). Informed by different layers' sensitivity to finetuning, we performed fine-tuning experiments by progressively excluding: key projection layers in self-attention across all encoder stack layers, the penultimate and ultimate encoder stacks, and the word embedding layer. Each of these exclusions independently or all three combined do not significantly degrade performance, while reducing the number of parameters to fine-tune by up to 40% (from 109M to 66M). Encouraged by these , we ask whether more aggressive constraints can be imposed to the finetuning process to further cut computational cost. Though L 0 -close fine-tuning obviates optimization of a large fraction of parameters, avoiding full storage of all parameters for each fine-tuned task, all operations still need to be performed at inference time. In order to reduce operations, we seek to sparsify parameters. Combined with strict L 0 -closeness, this amounts to a search over a binary mask in a high-dimensional parameter space. We adopt supermask training (see Section 3) to this end. Figure 3 shows fine-tuned model performance across GLUE tasks obtained by supermask training. Final sparsity level of the supermask is controlled by its initialization (see Section 3 and Appendix D). We note that there is little task performance degradation between 1% and 40% of parameter sparsity. Layer-wise sparsity levels of supermasks also demonstrate systematic trends (Figure 4): across GLUE tasks, W Q, W K and W I tend to be sparser than W V, W D and W O, and shallower encoder stack layers are sparser than deeper ones. Moreover, we show that supermask fine-tuning of only a fraction of sensitive layers could also achieve high performance (Table 3). Figure 4: Supermask sparsity levels across layers. Shown is the low-sparsity MNLI supermask with a global sparsity level of 12.9%, but similar patterns are observed across all GLUE tasks. How does the learning of a supermask actually work? Does a supermask trivially learn to prune away the weights with smallest magnitudes? To address these questions, we inspect the magnitudes of the pre-trained weights zeroed by the supermasks (Figure 5, Table 4). These weights turn out to have remarkably higher magnitudes than the next smallest entries, suggesting the learning of supermaks is mechanistically distinct from trivial magnitude-based pruning. Table 4: Comparison between supermask pruned weights and magnitude-based pruned weights. Specifically, we compare the weights pruned with low-sparsity supermasks (initialized at 0% sparsity) and weights pruned with one-shot magnitude-based pruning at the same final sparsity. We report the maximum and mean magnitude of the pruned weights. The last row shows percentages of the overlap between the supermask and the magnitude-based pruning mask, i.e. the percentages of weights zeroed by the supermask that are also the smallest weights. One surprising finding of this study is the many occurrences of good fine-tuned parameters among the 2 N configurations in the set θ: θ =θ µ, µ ∈ {0, 1} N, viz. vertices of an N -dimensional hypercube. First, there exist supermasks up to 40% sparse without significant performance degrada- Table 5: Low-sparsity supermask performance. We report the sparsity levels achieved when the supermasks were initialized at 0% sparsity. For several tasks, fine-tuning is achieved with less than 3% of pre-trained weights pruned. For the supermask evaluation , we include the mean and standard deviation of 10 Bernoulli samplings of a single run. tion for all GLUE tasks, for some tasks even sparser (Figure 3). It is natural that performance drops as the mask becomes sparser; however, it is rather counterintuitive that there exist good supermasks at the dense end (Figure 3), because we know that the pre-trained model with only the task-specific last layer fine-tuned utterly fails to learn any task (Appendix A). Figure 6: Low-sparsity supermask performance, i.e. task performance of super-masks initialized at 0% sparsity, compared against baseline. To shed more light on this phenomenon, we study the supermasks trained with all-dense initialization (Figure 6). Surprisingly, these low-sparsity supermasks successfully learned to perform all the tasks without significant degradation from baseline. Essentially, complicated tasks like MNLI and QQP can be learned by clamping 12 − 13% of the pre-trained weights to zero, whereas for simple tasks like MRPC and RTE, setting only 1 − 2% of the pre-trained weight entries to zero suffices to learn the task (Table 5). Fine-tuning can indeed be very fine, suggesting relative frequent occurrences of good solutions within a sparse L 0 -neighborhood of the pre-trained parameters. Finally, we question whether the supermasks learned to perform different tasks share commonalities. Specifically, we quantify the amount of overlapping zeros in learned supermasks across different tasks (Figure 7). It seems the overlaps are not substantially larger than what would have been caused by chance, suggesting that, even though there seem to be many good supermasks for each task, these masks are largely distinct from each other, each unique to the task it learns. Figure 7: Fractions of overlap of zero elements in supermasks across GLUE tasks, compared to randomly generated masks. Each value in the grid shows the fraction of pruned elements in one task (horizontal axis) that are also pruned in the other (vertical axis). Here, we show low-sparsity supermasks (initialized at 0% sparsity) and compare the masks in the value layer of the first encoder, which is one of the most sparse layers in the entire model. We show that, due to surprisingly frequent occurrences of good parameter configurations in the sparse L 0 -vicinity of large pre-trained language models, two techniques are highly effective in producing efficient fine-tuned networks to perform specific language understanding tasks: optimizing only the most sensitive layers and learning to sparsify parameters. In contrast to commonly employed post-training compression methods that have to trade off with performance degradation, our procedure of generating sparse networks is by itself an optimization process that learns specific tasks. Optimization of only task-specific layers does not lead to successful fine-tuning. For instance, for the MRPC task, freezing parameter weights in the pre-trained model and optimizing the task-specific last layer alone yields a low-performing model. Across 10 independent runs, the model consistently predicts all 1's for the paraphrase classification task, yielding an F1 score of 81.2. This is a significant degradation compared to the baseline performance of 89.4 ± 0.7 across multiple runs (Table 3). Thus, it is critical to fine-tune layers in the pre-trained model and not just the task-specific layers alone. Iterative pruning during fine-tuning (Figure 8) outperforms supermask training (Figure 3) at higher sparsity levels. While supermask training remains successful up to 40% sparsity, iterative pruning produces binary masks up to 50% sparse and for some tasks even sparser without significant performance degradation. However, while iterative pruning saves compute on appropriate hardware by performing sparse operations with model parameters, supermask training further saves memory by storing only binary masks of certain layers for each task. Figure 8: Iterative pruning during fine-tuning. We plot the evaluation performance at sparsity levels from 10% to 90% across GLUE tasks. Note the baseline performance for each task marked by the leftmost end of each curve (0% sparsity). Supermask training requires a much larger learning rate compared to typical training . While a learning rate of 2 × 10 −5 is used for optimizing weights, a learning rate of 2 × 10 −1 is used for optimizing masks. We notice a degradation in performance at smaller learning rates for supermask training (Table 6). Table 6: MRPC low-sparsity supermask performance at learning rates from 2 × 10 −5 and 2 × 10 −1. We note that supermask training requires a much higher learning rate than typical parameter training. At low learning rates, the model significantly degrades in performance and predicts all 0's for the paraphrase classification task, yielding an F1 score of 0.0. This pattern holds true across GLUE tasks. There is no straightforward to control the amount of weights pruned when training supermasks . We find that setting the initial sparsity through a soft magnitude-based pruning mask controls the final sparsity level. Figure 9 shows this correlation between initial and final sparsities of supermasks across GLUE tasks. The supermasks are initialized to a soft magnitude-based pruning mask and the sparsity level shifts during supermask training. This figure shows the initial sparsity level plotted against the final sparsity level. We note that, at lower initial sparsity levels, the supermask is pushed to a greater sparsity level, whereas at higher sparsity levels, the supermask is pushed to a lower sparsity level. This pattern is similar across GLUE tasks but is most prominent in the MNLI task, scaling with the number of fine-tuning steps (Table 1). We find that parameter distance scales with the number of fine-tuning steps (Figure 10). Figure 10: Scaling of parameter distance with the number of fine-tuning iterations. We find that angular distance correlates with the amount of fine-tuning (Table 1). Each data point corresponds to a different GLUE task.

Sparsification as fine-tuning of language models

Represent the short summary for finding the corresponding scientific article (long):

We propose a extension of the batch normalization, show a first-of-its-kind convergence analysis for this extension and show in numerical experiments that it has better performance than the original batch normalizatin.


Question: Represent the text to find a very short summary (1-2 sentences max)

Question: We present a simple and effective algorithm designed to address the covariate shift problem in imitation learning. It operates by training an ensemble of policies on the expert demonstration data, and using the variance of their predictions as a cost which is minimized with RL together with a supervised behavioral cloning cost. Unlike adversarial imitation methods, it uses a fixed reward function which is easy to optimize. We prove a regret bound for the algorithm in the tabular setting which is linear in the time horizon multiplied by a coefficient which we show to be low for certain problems in which behavioral cloning fails. We evaluate our algorithm empirically across multiple pixel-based Atari environments and continuous control tasks, and show that it matches or significantly outperforms behavioral cloning and generative adversarial imitation learning. Training artificial agents to perform complex tasks is essential for many applications in robotics, video games and dialogue. If success on the task can be accurately described using a reward or cost function, reinforcement learning (RL) methods offer an approach to learning policies which has been shown to be successful in a wide variety of applications However, in other cases the desired behavior may only be roughly specified and it is unclear how to design a reward function to characterize it. For example, training a video game agent to adopt more human-like behavior using RL would require designing a reward function which characterizes behaviors as more or less human-like, which is difficult. Imitation learning (IL) offers an elegant approach whereby agents are trained to mimic the demonstrations of an expert rather than optimizing a reward function. Its simplest form consists of training a policy to predict the expert's actions from states in the demonstration data using supervised learning. While appealingly simple, this approach suffers from the fact that the distribution over states observed at execution time can differ from the distribution observed during training. Minor errors which initially produce small deviations from the expert trajectories become magnified as the policy encounters states further and further from its training distribution. This phenomenon, initially noted in the early work of , was formalized in the work of who proved a quadratic O(T 2) bound on the regret and showed that this bound is tight. The subsequent work of showed that if the policy is allowed to further interact with the environment and make queries to the expert policy, it is possible to obtain a linear bound on the regret. However, the ability to query an expert can often be a strong assumption. In this work, we propose a new and simple algorithm called DRIL (Disagreement-Regularized Imitation Learning) to address the covariate shift problem in imitation learning, in the setting where the agent is allowed to interact with its environment. Importantly, the algorithm does not require any additional interaction with the expert. It operates by training an ensemble of policies on the demonstration data, and using the disagreement in their predictions as a cost which is optimized through RL together with a supervised behavioral cloning cost. The motivation is that the policies in the ensemble will tend to agree on the set of states covered by the expert, leading to low cost, but are more likely to disagree on states not covered by the expert, leading to high cost. The RL cost thus pushes the agent back towards the distribution of the expert, while the supervised cost ensures that it mimics the expert within the expert's distribution. Our theoretical show that, subject to realizability and optimization oracle assumptions, our algorithm obtains a O(κ T) regret bound for tabular MDPs, where κ is a measure which quantifies a tradeoff between the concentration of the demonstration data and the diversity of the ensemble outside the demonstration data. We evaluate DRIL empirically across multiple pixel-based Atari environments and continuous control tasks, and show that it matches or significantly outperforms behavioral cloning and generative adversarial imitation learning, often recovering expert performance with only a few trajectories. Denote by S the state space, A the action space, and Π the class of policies the learner is considering. Let T denote the task horizon and π the expert policy whose behavior the learner is trying to mimic. For any policy π, let d π denote the distribution over states induced by following π. Denote C(s, a) the expected immediate cost of performing action a in state s, which we assume is bounded in. In the imitation learning setting, we do not necessarily know the true costs C(s, a), instead we observe expert demonstrations. Our goal is to find a policy π which minimizes an observed surrogate loss between its actions and the actions of the expert under the induced distribution of states, i.e.π For the following, we will assume is the total variation distance (denoted by ·), which is an upper bound on the 0 − 1 loss. Our goal is thus to minimize the following quantity, which represents the distance between the actions taken by our policy π and the expert policy π: The following shows that if represents an upper bound on the 0 − 1 loss and C satisfies certain smoothness conditions, then minimizing this loss within translates into an O(T) regret bound on the task cost Theorem 1. Let π be such that J exp (π) =, and Unfortunately, it is often not possible to optimize J exp directly, since it requires evaluating the expert policy on the states induced by following the current policy. The supervised behavioral cloning cost J BC, which is computed on states induced by the expert, is often used instead: Minimizing this loss within yields a quadratic regret bound on regret: Furthermore, this bound is tight: as we will discuss later, there exist simple problems which match the worst-case lower bound. Our algorithm is motivated by two criteria: i) the policy should perform similarly to the expert on the expert's data distribution, and ii) the policy should move towards the expert's data distribution if it is away from it. These two criteria are addressed by combining two losses: a standard behavior cloning loss, and an additional loss which represents the variance over actions induced by sampling different policies from the posterior given the demonstration data D. We call this the uncertainty cost, which is defined as: 2: Initialize policy π and policy ensemble E = {π e} 3: for e = 1, E do 4: Sample D e ∼ D with replacement, with |D e | = |D|. Train π e to minimize J BC (π e) on D e to convergence. 6: end for 7: for i = 1,... do Perform one gradient update to minimize J BC (π) using a minibatch from D. Perform one step of policy gradient to minimize E s∼dπ,a∼π(·|s) [C clip U (s, a)]. 10: end for The motivation is that the variance over plausible policies is high outside the expert's distribution, since the data is sparse, but it is low inside the expert's distribution, since the data there is dense. Minimizing this cost encourages the policy to return to regions of dense coverage by the expert. Intuitively, this is what we would expect the expert policy π to do as well. The total cost which the algorithm optimizes is given by: The first term is a behavior cloning loss and is computed over states generated by the expert policy, of which the demonstration data D is a representative sample. The second term is computed over the distribution of states generated by the current policy and can be optimized using policy gradient. A number of methods can be used to approximate the posterior p(π|D). In our experiments, we use an ensemble E = {π e} |E| e=1 of models with different initializations which are trained on different bootstrap samples of the demonstration data. Note that the demonstration data is fixed, and this ensemble can be trained once offline. We then interleave the supervised behavioral cloning updates and the policy gradient updates which minimize the variance of the posterior. The full algorithm is shown in Algorithm 1. Other methods for approximating the posterior include Bayesian neural networks using diagonal Gaussian approximations or MCdropout , which we also found to work well (see Appendix D.2). In practice, for the supervised loss we optimize the KL divergence between the actions predicted by the policy and the expert actions, which is an upper bound on the total variation distance due to Pinsker's inequality. We also found it helpful to use a clipped uncertainty cost: where the threshold q is a top quantile of the raw uncertainty costs computed over the demonstration data. The threshold q defines a normal range of uncertainty based on the demonstration data, and values above this range incur a positive cost (or negative reward). The RL cost can be optimized using any policy gradient method. In our experiments we used advantage actor-critic (A2C) , which estimates the expected cost using rollouts from multiple parallel actors all sharing the same policy (see Appendix C for details). We note that model-based RL methods could in principle be used as well if sample efficiency is a constraint. We now analyze DRIL for tabular MDPs. We will show that, subject to assumptions that the policy class contains an optimal policy and that we are able to optimize costs within of their global minimum, our algorithm obtains a regret bound which is linear in κ T, where κ is a quantity specific to the environment and d π. Intuitively, κ represents a tradeoff between how concentrated the demonstration data is and how high the variance of the posterior is outside the expert distribution. Assumption 1. (Realizability) π ∈ Π. Assumption 2. (Optimization Oracle) For any given cost function J, our minimization procedure returns a policyπ ∈ Π such that J(π) ≤ arg min π∈Π J(π) +. The motivation behind our algorithm is that the policies in the ensemble agree inside the expert's distribution and disagree outside of it. This defines a reward function which pushes the learner back towards the expert's distribution if it strays away. However, what constitutes inside and outside the distribution, or sufficient agreement or disagreement, is ambiguous. Below we define quantities which makes these ideas precise. Definition 1. For any set U ⊆ S, define the maximum probability ratio between the state distributions induced by the expert policy and by policies in the policy class inside of U as α(U) = max π∈Π s∈U. For a set U, α(U) will be low if the expert distribution has high mass inside of U, and the states in U is reachable by policies in the policy class. Definition 2. Define the minimum variance of the posterior outside of U as We now define the κ coefficient as the minimum ratio of these two quantities over all possible subsets of S. Definition 3. We define κ(U) = α(U) β(U), and κ = min U ⊆S κ(U). We can view minimizing κ(U) over different U ⊆ S as minimizing a tradeoff between coverage by the expert policy inside of U, and variance of the posterior outside of U. Note that by making U very small, it may be easy to make α(U) small, but doing so may also make β(U) small and κ(U) large. Conversely, making U large may make β(U) large but may also make α(U) large as a . We now establish a relationship between the κ coefficient just defined, the cost our algorithm optimizes, and J exp defined in Equation which we would ideally like to minimize and which translates into a regret bound. All proofs can be found in Appendix A. This shows that if κ is not too large, and we are able to make our cost function J alg (π) small, then we can ensure J exp (π) is also be small. This is only useful if our cost function can indeed achieve a small minimum. The next lemma shows that this is the case. Here is the threshold specified in Assumption 2. Combining these two lemmas with the previous of , we get a regret bound which is linear in κ T. Theorem 3. Letπ be the of minimizing J alg using our optimization oracle, and assume that Our bound is an improvement over that of behavior cloning if κ is less than O(T). Note that DRIL does not require knowledge of κ. The quantity κ is problem-dependent and depends on the environment dynamics, the expert policy and the class of policies available to the learner. We next compute κ exactly for a problem for which behavior cloning is known to perform poorly, and show that it is independent of T. Example 1. Consider the tabular MDP given in as an example of a problem where behavioral cloning incurs quadratic regret, shown in Figure 1. There are 3 states S = (s 0, s 1, s 2) and two actions (a 1, a 2). Each policy π can be represented as a set of probabilities π(a 1 |s) for each state s ∈ S 1. The posterior p(π(a 1 |s)|D) is given by a Beta distribution with parameters Beta(n 1 + 1, n 2 + 1) where n 1, n 2 are the number of times the pairs (s, a 1) and (s, a 2) occur, respectively, in the demonstration data D. The agent always starts in s 0 and the expert's policy is given by π (a 1 |s 0) = 1, π (a 1 |s 1) = 0, π (a 1 |s 2) = 1. due to the dynamics of the MDP, so dπ(s) d π (s) ≤ 1 for s ∈ {s 0, s 1}. Furthermore, since s 2 is never visited in the demonstration data, p(π(a 1 |s 2)|D) = Beta = U nif orm, which also implies that p(π(a 1 |s 2)|D) = U nif orm. It follows that Var π∼p(π|D) (π(a|s 2)) is equal to the variance of a uniform distribution over, i.e. 1 12. Therefore: Applying our from Theorem 3, we see that our algorithm obtains an O(T) regret bound on this problem, in contrast to the O(T 2) regret of behavioral cloning 2. The idea of learning through imitation dates back at least to the work of , who trained a neural network to imitate the steering actions of a human driver using images as input. The problem of covariate shift was already observed, as the author notes: "when driving for itself, the network may occasionally stray from the center of the road and so must be prepared to recover by steering the vehicle back to the center of the road". This issue was formalized in the work of , who on one hand proved an O(T 2) regret bound, and on the other hand provided an example showing this bound is tight. The subsequent work proposed the DAGGER algorithm which obtains linear regret, provided the agent can both interact with the environment, and query the expert policy. Our approach also requires environment interaction, but importantly does not require the ability to query the expert. Also of note is the work of , which extended DAGGER to time series prediction problems by using the true targets as expert corrections. Imitation learning has been used within the context of modern RL to help improve sample efficiency or overcome exploration . These settings assume the reward is known and that the policies can then be fine-tuned with reinforcement learning. In this case, covariate shift is less of an issue since it can be corrected using the reinforcement signal. The work of also proposed a method to address the covariate shift problem when learning from demonstrations when the reward is known, by conservatively extrapolating the value function outside the training distribution using negative sampling. This addresses a different setting from ours, and requires generating plausible states which are off the manifold of training data, which may be challenging when the states are high dimensional such as images. The work of proposed to treat imitation learning within the Q-learning framework, setting a positive reward for all transitions inside the demonstration data and zero reward for all other transitions in the replay buffer. This rewards the agent for repeating (or returning to) the expert's transitions. The work of also incorporates a mechanism for reducing covariate shift by fitting a Q-function that classifies whether the demonstration states are reachable from the current state. Random Expert Distillation uses Random Network Distillation (RND) to estimate the support of the expert's distribution in state-action space, and minimizes an RL cost designed to guide the agent towards the expert's support. This is related to our method, but differs in that it minimizes the RND prediction error rather than the posterior variance and does not include a behavior cloning cost. The behavior cloning cost is essential to our theoretical and avoids certain failure modes, see Appendix B for more discusion. Generative Adversarial Imitation Learning (GAIL) is a state-of-the-art algorithm which addresses the same setting as ours. It operates by training a discriminator network to distinguish expert states from states generated by the current policy, and the negative output of the discriminator is used as a reward signal to train the policy. The motivation is that states which are outside the training distribution will be assigned a low reward while states which are close to it will be assigned a high reward. This encourages the policy to return to the expert distribution if is strays away from it. However, the adversarial training procedure means that the reward function is changing over time, which can make the algorithm unstable or difficult to tune. In contrast, our approach uses a simple fixed reward function. We include comparisons to GAIL in our experiments. Using disagreement between models in an ensemble to represent uncertainty has recently been explored in several contexts. The works of used disagreement between different dynamics models to drive exploration in the context of model-based RL. Conversely, used variance across different dropout masks to prevent policies from exploiting error in dynamics models. Ensembles have also been used to represent uncertainty over Q-values in model-free RL in order to encourage exploration . Within the context of imitation learning, the work of used the variance of the ensemble together with the DAGGER algorithm to decide when to query the expert demonstrator to minimize unsafe situations. Here, we use disagreement between different policies sampled from the posterior to address covariate shift in the context of imitation learning. As a first experiment, we applied DRIL to the tabular MDP of shown in Figure 1. We computed the posterior over the policy parameters given the demonstration data using . Shaded region represents range between 5 th and 95 th quantiles, computed across 500 trials. Behavior cloning exhibits poor worstcase regret, whereas DRIL has low regret across all trials. a separate Beta distribution for each state s with parameters determined by the number of times each action was performed in s. For behavior cloning, we sampled a single policy from this posterior. For our method, we sampled 5 policies and used their negative variance to define an additional reward function. We combined this with a reward which was the probability density function of a given state-action pair under the posterior distribution, which corresponds to the supervised learning loss, and used tabular Q-learning to optimize the sum of these two reward functions. This experiment was repeated 500 times for time horizon lengths up to 500 and N = 1, 5, 10 expert demonstration trajectories. Figure 2 shows plots of the regret over the 500 different trials across different time horizons. Although the average performance of BC improves with more expert demonstrations, it exhibits poor worst-case performance with some trials incurring very high regret, especially when using fewer demonstrations. Our method has low regret across all trials, which stays close to constant independantly of the time horizon, even with a single demonstration. This performance is better than that suggested by our analysis, which showed a worst-case linear bound with respect to time horizon. We next evaluated our approach on six different Atari environments. We used pretrained PPO agents from the stable baselines repository to generate N = {1, 3, 5, 10, 15, 20} expert trajectories. We compared against two other methods: standard behavioral cloning (BC) and Generative Adversarial Imitation Learning (GAIL). Results are shown in Figure 3a. DRIL outperforms behavioral cloning across most environments and numbers of demonstrations, often by a substantial margin. In the worst case its performance matches that of behavior cloning. In many cases, our method is able to match the expert's performance using a small number of trajectories. Figure 3b shows the evolution of the uncertainty cost and the policy reward throughout training. In all cases, the test reward improves while the uncertainty cost decreases. Interestingly, there is correspondence between the change in the uncertainty cost during training and the gap in performance between behavior cloning and DRIL. For example, in MsPacman there is both a small improvement in uncertainty cost over time and a small gap between behavior cloning and our method, whereas in Breakout there is a large improvement in uncertainty cost and a large gap between behavior cloning and our method. This suggests that the gains from our method comes from redirecting the policy back towards the expert manifold, which is manifested as a decrease in the uncertainty cost. We were not able to obtain meaningful performance for GAIL on these domains, despite performing a hyperparameter search across learning rates for the policy and discriminator, and across different numbers of discriminator updates. We additionally experimented with clipping rewards in an effort to stabilize performance. These are consistent with those of , who also reported negative when running GAIL on images. While improved performance might be possible with more sophisticated adversarial training techniques, we note that this contrasts with our method which uses a fixed reward function obtained through simple supervised learning. In Appendix D we provide ablation experiments examining the effects of different cost function choices and the role of the BC loss. We also compare the ensemble approach to a dropout-based approach for posterior approximation and show that DRIL works well in both cases. We next report of running our method on a 6 different continuous control tasks from the PyBullet 3 and OpenAI Gym environments. We again used pretrained agents to generate expert demonstrations. Results are shown in Figure 4. In these environments we found behavior cloning to be a much stronger baseline than for the Atari environments, and in many tasks it was able to match expert performance using as little as 3 trajectories. Our method exhibits a modest improvement on Walker2D and BipedalWalkerHardcore when a single trajectory is used, and otherwise has similar performance to behavior cloning. The fact that our method does not perform worse than behavior cloning on tasks where covariate shift is likely less of an issue provides evidence of its robustness. Addressing covariate shift has been a long-standing challenge in imitation learning. In this work, we have proposed a new method to address this problem by penalizing the disagreement between an ensemble of different policies sampled from the posterior. Importantly, our method requires no additional labeling by an expert. Our experimental demonstrate that DRIL can often match expert performance while using only a small number of trajectories across a wide array of tasks, ranging from tabular MDPs to pixel-based Atari games and continuous control tasks. On the theoretical side, we have shown that our algorithm can provably obtain a low regret bound for tabular problems in which the κ parameter is low. There are multiple directions for future work. On the theoretical side, extending our analysis to continuous state spaces and characterizing the κ parameter on a larger array of problems would help to better understand the settings where our method can expect to do well. Empirically, there are many other settings in structured prediction (Daumé et al., 2009) where covariate shift is an issue and where our method could be applied. For example, in dialogue and language modeling it is common for generated text to become progressively less coherent as errors push the model off the manifold it was trained on. Our method could potentially be used to fine-tune language or translation models after training by applying our uncertainty-based cost function to the generated text. A PROOFS Proof. We will first show that for any π ∈ Π and U ⊆ S, we have. We can rewrite this as: We begin by bounding the first term: We next bound the second term: Now observe we can decompose the RL cost as follows: Putting these together, we get the following: Here we have used the fact that β(U) ≤ 1 since 0 ≤ π(a|s) ≤ 1 and α(U) ≥ s∈U Taking the minimum over subsets U ⊆ S, we get J exp (π) ≤ κJ alg (π). Proof. Plugging the optimal policy into J alg, we get: We will first bound Term 1: We will next bound Term 2: The last step follows from our optimization oracle assumption: Combining the bounds on the two terms, we get J alg (π) ≤ 2. Since π ∈ Π, the follows. Theorem 1. Letπ be the of minimizing J alg using our optimization oracle, and assume that Proof. By our optimization oracle and Lemma 2, we have Combining with Lemma 1, we get: Applying Theorem 1 from , we get J(π) ≤ J(π) + 3uκ T. The following example shows how minimizing the uncertainty cost alone without the BC cost can lead to highly sub-optimal policies if the demonstration data is generated by a stochastic policy which is only slightly suboptimal. Consider the following deterministic chain MDP: The agent always starts in s 1, and gets a reward of 1 in s 3 and 0 elsewhere. The optimal policy is given by: Assume the demonstration data is generated by the following policy, which is only slightly suboptimal: Let us assume realizability and perfect optimization for simplicity. If both transitions (s 2, a 0) and (s 2, a 1) appear in the demonstration data, then Random Expert Distillation (RED) will assign zero cost to both transitions. If we do not use bootstrapped samples to train the ensemble, then DRIL without the BC cost (we will call this UO-DRIL for Uncertainty-Only DRIL) will also assign zero cost to both transitions since all models in the ensemble would recover the Bayes optimal solution given the demonstration data. If we are using bootstrapped samples, then the Bayes optimal solution for each bootstrapped sample may differ and thus the different policies in the ensemble might disagree in their predictions, although given enough demonstration data we would expect these differences (and thus the uncertainty cost) to be small. Note also that since no samples at the state s 0 occur in the demonstration data, both RED and UO-DRIL will likely assign high uncertainty costs to state-action pairs at (s 0, a 0), (s 0, a 1) and thus avoid highly suboptimal policies which get stuck at s 0. Now consider policiesπ 1,π 2 given by:π Both of these policies only visit state-action pairs which are visited by the demonstration policy. In the case described above, both RED and UO-DRIL will assignπ 1 andπ 2 similarly low costs. However,π 1 will cycle forever between s 1 and s 2, never collecting reward, whileπ 2 will with high probability reach s 3 and stay there, thus achieving high reward. This shows that minimizing the uncertainty cost alone does not necessarily distinguish between good and bad policies. However,π 1 will incur a higher BC cost thanπ 2, sinceπ 2 more closely matches the demonstration data at s 2. This shows that including the BC cost can be important for further disambiguating between policies which all stay within the distribution of the demonstration data, but have different behavior within that distribution. C EXPERIMENTAL DETAILS C.1 ATARI ENVIRONMENTS All behavior cloning and ensemble models were trained to minimize the negative log-likelihood classification loss on the demonstration data for 500 epochs using Adam and a learning rate of 2.5 · 10 −4. For our method, we initially performed a hyperparameter search on Space Invaders over the following values: We then chose the best values and kept those hyperparameters fixed for all other environments. All other A2C hyperparameters follow the default values in the repo : policy networks consisted of 3-layer convolutional networks with 8−32−64 feature maps followed by a single-layer MLP with 512 hidden units. For GAIL, we used the implementation in and replaced the MLP discriminator by a CNN discriminator with the same architecture as the policy network. We initially performed a hyperparameter search on Breakout with 10 demonstrations over the values shown in Table 2. However, we did not find any hyperparameter configuration which performed better than behavioral cloning. All behavior cloning and ensemble models were trained to minimize the mean-squared error regression loss on the demonstration data for 500 epochs using Adam and a learning rate of 2.5 · 10 −4. Policy networks were 2-layer fully-connected MLPs with tanh activations and 64 hidden units. • DRIL This is the regular DRIL agent, which optimizes both the BC cost and the clipped cost: • DRIL (clipped cost 0/1) This is the same as the regular DRIL agent, except that we use the following clipped cost function: • DRIL (raw cost) This is the same as the regular DRIL agent, except that we use the raw cost C U (s, a) rather than the clipped cost. • DRIL (no BC cost) This is the same as the regular DRIL agent, except that we remove the BC updates and only optimize C clip U (s, a). Results are shown in Figure 5. First, switching from the clipped cost in {−1, +1} to the clipped cost in {0, 1} or the raw cost causes a large drop in performance across most environments. One explanation may be that since the costs are always positive for both variants (which corresponds to a reward which is always negative), the agent may learn to terminate the episode early in order to minimize the total cost incurred. Using a cost/reward which has both positive and negative values avoids this behavior. Second, optimizing the pure BC cost performs better than the pure uncertainty cost for some environments (MsPacman, SpaceInvaders, BeamRider) while optimizing the pure uncertainty cost performs better than BC for others (Breakout, Qbert). DRIL, which optimizes both, has robust performance and performs the best out of the different variants over most environments and numbers of trajectories. We provide additional comparing the ensembling and MC-dropout approaches to posterior estimation. For MC-Dropout we trained a single policy network with a dropout rate of 0.1 applied to all layers except the last, and estimated the variance for each state-action pair using 10 different dropout masks. Similarly to the ensemble approach, we computed the 98 th quantile of the variance on the demonstration data and used this value in our clipped cost. Results are shown for three environments in Figure 6. MC-dropout performs similarly to the ensembling approach, which shows that our method can be paired with different approaches to posterior estimation.

Represent the short summary for finding the corresponding scientific article (long)

Method for addressing covariate shift in imitation learning using ensemble uncertainty

We isolate one factor of RL generalization by analyzing the case when the agent only overfits to the observations. We show that architectural implicit regularizations occur in this regime.


Question: represent the text to find a very short summary (1-2 sentences max)

Question: We present and discuss a simple image preprocessing method for learning disentangled latent factors. In particular, we utilize the implicit inductive bias contained in features from networks pretrained on the ImageNet database. We enhance this bias by explicitly fine-tuning such pretrained networks on tasks useful for the NeurIPS2019 disentanglement challenge, such as angle and position estimation or color classification. Furthermore, we train a VAE on regionally aggregate feature maps, and discuss its disentanglement performance using metrics proposed in recent literature. Fully unsupervised methods, that is, without any human supervision, are doomed to fail for tasks such as learning disentangled representations . In this contribution, we utilize the implicit inductive bias contained in models pretrained on the ImageNet database , and enhance it by finetuning such models on challenge-relevant tasks such as angle and position estimation or color classification. In particular, our submission for challenge stage 2 builds on our submission from stage 1 1, in which we employed pretrained CNNs to extract convolutional feature maps as a preprocessing step before training a VAE . Although this approach already in partial disentanglement, we identified two issues with the feature vectors extracted this way. Firstly, the feature extraction network is trained on ImageNet, which is rather dissimilar to the MPI3d dataset used in the challenge. Secondly, the feature aggregation mechanism was chosen ad-hoc and likely does not retain all information needed for disentanglement. We attempt to fix these issues by finetuning the feature extraction network as well as learning the aggregation of feature maps from data by using the labels of the simulation datasets MPI3d-toy and MPI3d-realistic. Our method consists of the following three steps: supervised finetuning of the feature extraction CNN (section 2.1), extracting a feature vector from each image in the dataset using the finetuned network (section 2.2), training a VAE to reconstruct the feature vectors and disentangle the latent factors of variation (section 2.3). In this step, we finetune the feature extraction network offline (before submission to the evaluation server). The goal is to adapt the network such that it produces aggregated feature vectors that capture the latent variables well. In particular, the network is finetuned by learning to predict the value of each latent factor from the aggregated feature vector of an image. To this end, we use the simulation datasets MPI3d-toy and MPI3d-realistic 2, namely the images as inputs and the labels as supervised classification targets. For the feature extraction network, we use the VGG19-BN architecture of the torchvision package. The input images are standardized using mean and variance across each channel computed from the ImageNet dataset. We use the output feature maps of the last layer before the final average pooling (dimensionality 512ˆ2ˆ2) as the input to a feature aggregation module which reduces the feature map to a 512-dimensional vector 3. This aggregation module consists of three convolution layers with 1024, 2048, 512 feature maps and kernel sizes 1, 2, 1 respectively. Each layer is followed by batch normalization and ReLU activation. We also employ layerwise dropout with rate 0.1 before each convolution layer. Finally, the aggregated feature vector is 2-normalized, which was empirically found to be important for the ing disentanglement performance. Then, for each latent factor, we add a linear classification layer computing the logits of each class from the aggregated feature vector. These linear layers are discarded after this step. We use both MPI3d-toy and MPI3d-realistic for training to push the network to learn features that identify the latent factors in a robust way, regardless of details such as reflections or specific textures. In particular, we use a random split of 80% of each dataset as the training set, and the remaining samples as a validation set. VGG19-BN is initialized with a set of weights ing from ImageNet training 4, and the aggregation module and linear layers were randomly initialized using uniform He initialization . The network is trained for 5 epochs using the RAdam optimizer with learning rate 0.001, β 0 " 0.999, β 1 " 0.9, a batch size of 512 and a weight decay of 0.01. We use a multi-task classification loss consisting of the sum of cross entropies between the prediction and the ground truth of each latent factor. After training, the classification accuracy on the validation set is around 98% for the two degrees of freedom of the robot arm, and around 99.9% for the remaining latent factors. In this step, we use the finetuned feature extraction network to produce a set of aggregated feature vectors. We simply run the network detailed in the previous step on each image of the dataset and store the aggregated 512-dimensional vectors in memory. Again, inputs to the feature extractor are standardized such that mean and variance across each channel correspond to the respective ones from the ImageNet dataset. 2. Pretraining using any data was explicitly stated to be allowed by the challenge organizers. 3. Using aggregated feature vectors instead of feature maps is necessitated by the memory requirements of the challenge. 4. https://download.pytorch.org/models/vgg19_bn-c79401a0.pth Finally, we train a standard β-VAE on the set of aggregated feature vectors ing from the previous step. The encoder network consists of 4 fully-connected layers with 1024 neurons each, followed by two fully-connected layers parametrizing mean and log variance of the factorized Gaussian approximate posterior q pz | xq " N`µ, σ 2w ith C " 16 latent factors. The number of latent factors was experimentally determined. The decoder network consists of 4 fully-connected layers with 1024 neurons each, followed by a fully-connected layer parametrizing the mean of the factorized Gaussian conditional distribution p px | zq " N pμ, Iq. The mean is constrained to range p0, 1q using the sigmoid activation. All fully-connected layers but the final ones use batch normalization and are followed by ReLU activation functions. We use orthogonal initialization for all layers and assume a factorized Gaussian prior p pzq " N p0, Iq on the latent variables. For optimization, we use the RAdam optimizer with a learning rate of 0.001, β 0 " 0.999, β 1 " 0.9 and a batch size of B " 256. The VAE is trained for N " 100 epochs by minimizing where β is a hyperparameter to balance the losses of the MSE reconstruction and the KLD penalty terms. As the scale of the KLD term depends on the numbers of latent factors C, we normalize it by C such that β can be varied independently of C. It can be harmful to start training with too much weight on the KLD term . Therefore, we use the following cosine schedule to smoothly anneal β from β start " 0.001 to β end " 0.4 over the course of training: where βptq is the value for β in training episode t P t0,..., N´1u, and annealing runs from epoch t start " 10 to epoch t end " 49. This schedule lets the model initially learn to reconstruct the data well and only then puts pressure on the latent variables to be factorized which we found to considerably improve performance. On the public leaderboard (i.e. on MPI3D-real), our best submission achieves the first rank on the FactorVAE , and DCI metrics, with a large gap to the second-placed entry. See appendix A for a discussion of the . Unsurprisingly, introducing prior knowledge simplifies the disentanglement task considerably, reflected in improved scores. To do so, our approach makes use of task-specific supervision obtained from simulation, which restricts its applicability. Nevertheless, it constitutes a demonstration that this type of supervision can transfer to better disentanglement on real world data, which was one of the goals of the challenge.

We use supervised finetuning of feature vectors to improve transfer from simulation to the real world

Represent this text to find a much much longer article about it with say hundreds-thousands of words

Redistributing and growing weights according to the momentum magnitude enables the training of sparse networks from random initializations that can reach dense performance levels with 5% to 50% weights while accelerating training by up to 5.6x.