Question: Represent the scientific article to find a short TLDR

Question: It has long been assumed that high dimensional continuous control problems cannot be solved effectively by discretizing individual dimensions of the action space due to the exponentially large number of bins over which policies would have to be learned. In this paper, we draw inspiration from the recent success of sequence-to-sequence models for structured prediction problems to develop policies over discretized spaces. Central to this method is the realization that complex functions over high dimensional spaces can be modeled by neural networks that predict one dimension at a time. Specifically, we show how Q-values and policies over continuous spaces can be modeled using a next step prediction model over discretized dimensions. With this parameterization, it is possible to both leverage the compositional structure of action spaces during learning, as well as compute maxima over action spaces (approximately). On a simple example task we demonstrate empirically that our method can perform global search, which effectively gets around the local optimization issues that plague DDPG. We apply the technique to off-policy (Q-learning) methods and show that our method can achieve the state-of-the-art for off-policy methods on several continuous control tasks. Reinforcement learning has long been considered as a general framework applicable to a broad range of problems. However, the approaches used to tackle discrete and continuous action spaces have been fundamentally different. In discrete domains, algorithms such as Q-learning leverage backups through Bellman equations and dynamic programming to solve problems effectively. These strategies have led to the use of deep neural networks to learn policies and value functions that can achieve superhuman accuracy in several games (; where actions lie in discrete domains. This success spurred the development of RL techniques that use deep neural networks for continuous control problems BID12 BID20 . The gains in these domains, however, have not been as outsized as they have been for discrete action domains. This disparity is, in part, a of the inherent difficulty in maximizing an arbitrary function on a continuous domain, even in low-dimensional settings. Furthermore, it becomes harder to apply dynamic programming methods to back up value function estimates from successor states to parent states in continuous control problems. Several of the recent continuous control reinforcement learning approaches attempt to borrow characteristics from discrete problems by proposing models that allow maximization and backups more easily BID12 .One way in which continuous control can avail itself of the above advantages is to discretize each of the dimensions of continuous control action spaces. As noted in, doing this naively, however, would create an exponentially large discrete space of actions. For example with M dimensions being discretized into N bins, the problem would balloon to a discrete space with M N possible actions. We leverage the recent success of sequence-to-sequence type models BID32 to train such discretized models, without falling into the trap of requiring an exponentially large number of actions. Our method relies on a technique that was first introduced in BID3, which allows us to escape the curse of dimensionality in high dimensional spaces by modeling complicated probability distributions using the chain rule decomposition. In this paper, we similarly parameterize functions of interest -Q-values -using a decomposition of the joint function into a sequence of conditional values tied together with the bellman operator. With this formulation, we are able to achieve fine-grained discretization of individual domains, without an explosion in the number of parameters; at the same time we can model arbitrarily complex distributions while maintaining the ability to perform (approximate) global maximization. These benefits come at the cost of shifting the exponentially complex action space into an exponentially complex MDP BID5 BID10. In many settings, however, there are relationships between transitions that can be leveraged and large regions of good solutions, which means that this exponential space need not be fully explored. Existing work using neural networks to perform approximate exponential search is evidence of this BID37; BID2.While this strategy can be applied to most function approximation settings in RL, we focus on off-policy settings with an algorithm akin to DQN. Empirical on an illustrative multimodal problem demonstrates how our model is able to perform global maximization, avoiding the exploration problems faced by algorithms like NAF BID13 and DDPG. We also show the effectiveness of our method on a range of benchmark continuous control problems from hopper to humanoid. In this paper, we introduce the idea of building continuous control algorithms utilizing sequential, or autoregressive, models that predict over action spaces one dimension at a time. Here, we use discrete distributions over each dimension (achieved by discretizing each continuous dimension into bins) and apply it using off-policy learning. We briefly describe the notation we use in this paper. Let s t ∈ R L be the observed state of the agent, a ∈ R N be the N dimensional action space, and E be the stochastic environment in which the agent operates. Finally, let a i:j = a i · · · a j T be the vector obtained by taking the sub-range/slice of a DISPLAYFORM0 At each step t, the agent takes an action a t, receives a reward r t from the environment and transitions stochastically to a new state s t+1 according to (possibly unknown) dynamics p E (s t+1 |s t, a t). An episode consists of a sequence of such steps (s t, a t, r t, s t+1), with t = 1 · · · H where H is the last time step. An episode terminates when a stopping criterion F (s t+1) is true (for example when a game is lost, or when the number of steps is greater than some threshold length H max).Let R t = H i=t γ i−1 r i be the discounted reward received by the agent starting at step t of an episode. As with standard reinforcement learning, the goal of our agent is to learn a policy π (s t) that maximizes the expected future reward E [R H] it would receive from the environment by following this policy. Because this paper is focused on off-policy learning with Q-Learning BID39, we will provide a brief description of the algorithm. Q-learning is an off-policy algorithm that learns an action-value function Q (s, a) and a corresponding greedy-policy, π Q (s) = argmax a Q (s, a). The model is trained by finding the fixed point of the Bellman operator, i.e. DISPLAYFORM0 This is done by minimizing the Bellman Error, over the exploration distribution, ρ β (s) DISPLAYFORM1 Traditionally, Q is represented as a table of state action pairs or with linear function approximators or shallow neural networks BID39 BID35. Recently, there has been an effort to apply these techniques to more complex domains using non-linear function approximators that are deep neural networks (; . In these models, a Deep Q-Network (DQN) parameterized by parameters, θ, is used to predict Q-values, i.e. Q(s, a) = f (s, a; θ). The DQN parameters, θ, are trained by performing gradient descent on the error in equation 2, without taking a gradient through the Q-values of the successor states (although, see BID1 for an approach that takes this into account).Since the greedy policy, π Q (s), uses the action value with the maximum Q-value, it is essential that any parametric form of Q be able to find a maxima easily with respect to actions. For a DQN where the output layer predicts the Q-values for each of the discrete outputs, it is easy to find this max -it is simply the action corresponding to the index of the output with the highest estimated Q-value. In continuous action problems, it can be tricky to formulate a parametric form of the Q-value where it is easy to find such a maxima. Existing techniques either use a restrictive functional form, such as NAF BID13. DDPG employs a second neural network to approximate this max, in addition to the Q function approximator. This second network is trained to maximize / ascend the Q function as follows: DISPLAYFORM2 where ρ β is the state distribution explored by some behavioral policy, β and µ(·; θ µ) is the deterministic policy. In this work we modify the form of our Q-value function while still retaining the ability to find local maxima over actions for use in a greedy policy. In this section, we outline our proposed model, Sequential DQN (SDQN). This model decomposes the original MDP model with N -D actions to a similar MDP which contains sequences of 1-D actions. By doing this, we have 2 layer hierarchy of MDP -the "upper" containing the original environment, and the "lower" containing the transformed, stretched out, MDP. Both MDP model the same environment. We then combine these MDP by noting equality of Q values at certain states and doing bellman backups against this equality. See figure 1 for a pictorial view of this hierarchy. Consider an environment with states s t and actions a ∈ R N. We can perform a transformation to this environment into a similar environment replacing each N -D action into a sequence of N 1-D actions. This introduces a new MDP consisting of states u st k where superscript denotes alignment to the state s t, above, and subscript k to denote time offset on the lower MDP from s t. As a , u st k = (s t, a 1:k) is a tuple containing the state s t from original MDP and a history of additional states in the new MDP -in our case, a concatenation of actions a 1:k previously selected. The transitions of this new MDP can be defined by two rules: when all 1-D actions are taken we compute 1 step in the N -D environment receiving a new state, s t+1, a reward r t, and resetting a. In all other transitions, we append the previously selected action in a and receive 0 reward. This transformation reduces the N -D actions to a series of 1-D actions. We can now discretize the 1-D output space and directly apply Q-learning. Note that we could apply this strategy to continuous values, without discretization, by choosing a conditional distribution, such as a mixture of 1-D Gaussians, over which a maxima can easily be found. As such, this approach is equally applicable to pure continuous domains as compared to discrete approximations. The downside to this transformation is that it increases the number of steps needed to solve the transformed MDP. In practice, this transformation makes learning a Q-function considerably harder. The extra steps of dynamic programming coupled with learned function approximators causes large overestimation and stability issues. This can be avoided by learning Q-values for both MDPs at the same time and performing the bellman backup from the lower to the upper MDP for the transitions, s t, where Q-values should be equal. We define Q U (s, a) ∈ R, a ∈ R N, as a function that is the Q-value for the top MDP. Next, we define Q L (u, a i) ∈ R where a i ∈ R as the Q value for the lower MDP. We would like to have consistent Q-values across both MDPs when they perform one step in the environment. To make this possible, we must define how the time discounting works. We define the lower MDP to have zero discount for all steps except for when the real environment changes state. Thus, the discount is 0 for all all u st k where k < N, and the same as the top MDP when k = N. By doing this, the following is then true: DISPLAYFORM0 where u st N −1 contains the upper state and the previous N − 1 actions: (s t, a t 1:N −1). This equality allows us to "short circuit" the backups. Backups are only needed up until the point where the upper MDP can be used improving training and stability. During training, we parameterize Q U and Q L as neural networks. We learn Q U via TD-0 learning by minimizing: DISPLAYFORM1 Next, we learn Q L by also doing Q-learning, but we make use of of the equality noted in equation 3 and zero discounting. There is no new information nor environment dynamics for this MDP. As such we can draw samples from the same replay buffer used for learning Q U. For states u st k where k < N we minimize the bellman error as follows: DISPLAYFORM2 When Q U and Q L should be equal, as defined in equation 3, we do not backup we instead enforce soft equality by MSE. DISPLAYFORM3 In practice, as in DQN, we can also make use of target networks and/or double DQN BID15 when training Q U and Q L for increased stability. When using this model as a policy we compute the argmax over each action dimension of the lower MDP. As with DQN, we employ exploration when training with either epsilon greedy exploration or Boltzmann exploration. Q U is a MLP whose inputs are state and actions and outputs are Q values. Unlike in DDPG, the loss function does not need to be smooth with respect to actions. As such, we also feed in discretized representation of each action dimension to make training simpler. We worked with two parameterizations for Q L. First, we looked at a recurrent LSTM model BID17. This model has shared weights and passes information via hidden activations from one action dimension to another. The input at each time step is a function of the current state from the upper MDP, s t, and a single action dimension, a i. As it's an LSTM, the hidden state is capable of accumulating the previous actions. Second, we looked at a version with separate weights for each step of the lower MDP. The lower MDP does not have a fixed size input as the amount of action dimensions it has as inputs varies. To combat this, we use N separate models that are switched between depending on the state index of the lower MDP. We call these distinct models Q i where i ∈ [1, N]. These models are feed forward neural networks that take as input a concatenation of all previous action selections, a 1 t: i, as well as the upper state, s t. Doing this in switching Q L with the respective Q i in every use. Empirically we found that this weight separation led to more stable training. In more complex domains, such as vision based control tasks for example, one should untie only a subset of the weights and keep common components -a vision system -the same. In practice, we found that for the simplistic domains we worked in fully untied weights was sufficient. Architecture exploration for these kinds of models is still ongoing work. For full detail of model architectures and training procedures selection see Appendix C.3 RELATED WORK Our work was motivated by two distinct desires -to learn policies over exponentially large discrete action spaces, and to approximate value functions over high dimensional continuous action spaces effectively. In our paper we used a sequential parameterization of policies that help us to achieve this without making an assumption about the actual functional form of the model. Other prior work attempts to handle high dimensional action spaces by assuming specific decompositions. For example, BID26 were able to scale up learning to extremely large action sizes by factoring the action value function and use product of experts to learn policies. An alternative strategy was proposed in using action embeddings and applying k-nearest neighbors to reduce scaling of action sizes. By laying out actions on a hypercube, BID23 are able to perform a binary search over actions ing in a logarithmic search for the optimal action. Their method is similar to SDQN, as both construct a Q-value from sub Q-values. Their approach presupposes these constraints, however, and optimizes the Bellman equation by optimizing hyperplanes independently thus enabling optimizing via linear programming. Our approach is iterative and refines the action selection, which contrasts to their independent sub-plane maximization. and BID22 proposes a transformation similar to ours where a continuous action MDP is converted to a sequence of transitions representing a binary search over the continuous actions. In our setting, we used a 2-layer hierarchy of variable width as opposed to a binary tree. Additionally, we used the original MDP as part of our training procedure to reduce estimation error. We found this to be critical to reduce overestimation error when working with function approximators. Along with the development of discrete space algorithms, researchers have innovated specialized solutions to learn over continuous state and action environments including BID30 BID13. More recently, novel deep RL approaches have been developed for continuous state and action problems. TRPO BID27 and A3C uses a stocastic policy parameterized by diagonal covariance Gaussian distributions. NAF BID13 relies on quadratic advantage function enabling closed form optimization of the optimal action. Other methods structure the network in a way such that they are convex in the actions while being non-convex with respect to states BID0 or use a linear policy BID24.In the context of reinforcement learning, sequential or autoregressive policies have previously been used to describe exponentially large action spaces such as the space of neural architectures, BID40 and over sequences of words BID2 BID29. These approaches rely on policy gradient methods whereas we explore off-policy methods. Hierarchical/options based methods, including BID9 ) which perform spatial abstraction or BID34 that perform temporal abstraction pose another way to factor action spaces. These methods refine their action selection from time where our approaches operates on the same timescale and factors the action space. A vast literature on constructing sequential models to solve tasks exists outside of RL. These models are a natural fit when the data is generated in a sequential process such as in language modeling This in more sample efficiency but yields poor approximations of the Q surface outside of where the policy is. Right: Reward achieved over time. DDPG quickly converges to a local maximum. SDQN has high variance performance initially as it searches the space, but then quickly converges to the global maximum as the Q surface estimate becomes more accurate. BID4. One of the first and most effective deep learned sequence-to-sequence models for language modeling was proposed in BID33, which used an encoder-decoder architecture. In other domains, techniques such as NADE BID19 To consider the effectiveness of our algorithm, we consider a deterministic environment with a single time step, and a 2D action space. This can be thought of as being a two-armed bandit problem with deterministic rewards, or as a search problem in 2D action space. We chose our reward function to be a multimodal distribution as shown in the first column in FIG2. A large suboptimal mode and a smaller optimal mode exist. As with bandit problems, this formulation helps us isolate the ability of our method to find an optimal policy, without the confounding effect that arises from backing up rewards via the Bellman operator for sequential problems. As in traditional RL, we do exploration while learning. We consider uniformly sampling (-greedy with = 1) as well as sampling data from a normal distribution centered at the current policy -we refer to this as "local." A visualization of the final Q surfaces as well as training curves can be found in FIG2. DDPG uses local optimization to learn a policy on a constantly changing estimate of Q values predicted by a critic. The form of the Q distribution is flexible and as such there is no closed form properties we can make use of for learning a policy. As such, gradient descent, a local optimization algorithm, is used. This algorithm can get stuck in a sub-optimal policy. We hypothesize that these local maximum in policy space exist in more realistic simulated environments as well. Traditionally, deep learning methods use local optimizers and avoid local minima or maxima by working in a high dimensional parameter space BID8. In RL, however, the action space of a policy is relatively small dimensional thus it is much more likely that they exist. For example, in the hopper environment, a common failure mode we experienced when training algorithms like DDPG is to learn to balance instead of moving forward and hopping. We contrast this to SDQN. As expected, this model is capable of completely representing the Q surface (under the limits of discretization). The optimization of the policy is not done locally however enabling convergence to the optimal policy. Much like DDPG, the Q surface learned can be done on uniform, off policy, data. Unlike DDPG, however, the policy will not get stuck in a local maximum. In the uniform behavior policy setting, the model slowly reaches the right solution.1 With a behavior policy that is closer to being on-policy (such as the stochastic Gaussian greedy policy referred to above), the rate of convergence increases. Much of the error occurs from selecting over estimated actions. When sampling more on policy, the over estimated data points get sampled more frequently ing in faster training. To evaluate the relative performance of these models we perform a series of experiments on common continuous control tasks. We test the hopper (3-D action space), swimmer (2-D action space), half cheetah (6-D action space), walker2d (6-D action space) and the humanoid environment (17-D action space) from the OpenAI gym suite BID6. 2 We performed a wide hyper parameter search over various parameters in our models (described in Appendix C), and selected the best performing runs. We then ran 10 random seeds of the same hyper parameters to evaluate consistency and to get a more realistic estimate of performance. We believe this replication is necessary as many of these algorithms are not only sensitive to both hyper parameters but random seeds. First, we look at learning curves of some of the environments tested in FIG4. Our method quickly achieves good policies much faster than DDPG. For a more qualitative analysis, we use the best reward achieved while training averaged across over 25,000 steps and with evaluations sampled every 5,000 steps. Again we perform an average over 10 different random seeds. This metric gives a much better sense of stability than the traditionally reported instantaneous max reward achieved during training. We compare our algorithm to the current state-of-the-art in off-policy continuous control: DDPG. Through careful model selection and extensive hyper parameter tuning, we train DDPG agents with performance better than previously published for some of these tasks. Despite this search, however, we believe that there is still space for significant performance gain for all the models given different neural network architectures and hyper parameters. See for discussion on implementation variability and performance. Results can be seen in Figure 4. Our algorithm achieves better performance on four of the five environments we tested. Unlike existing continuous control algorithms, we have a choice over the number of discritization bins we choose, B. To test the effect of this we first take the best performing half cheetah hyper parameter configuration found above, and rerun it varying the number of bins. For statistical significance we run 10 trials per tested value of B. Results can be found in Figure 5. These suggest that SDQN is robust to this hyper parameter, working well in all bin amounts greater than 4. Lower than 4 bins does not yield enough fine grain enough control to solve this task effectively. Next we look to the effect of action order. In most existing environments there is no implicit "ordering" to environments. Given the small action space dimensionality, we hypothesized that this ordering would not matter. We test this hypothesis by taking our hyper parameters for half cheetah found in section 4.2 and reran them with random action ordering, 10 times for statistical significance. Half cheetah has 6 action dimensions thus we are only exploring a subset of orderings. Seed 0 represents the ordering used when finding the original hyper parameters. Results can be found in Figure 5. While there is some variability, the overall changes in performance are small validating our original assumption. Conceptually, our approach centers on the idea that action selection at each stage can be factored and sequentially selected. In this work we use 1-D action spaces that are discretized as our base component. Existing work in the image modeling domain suggests that using a mixture of logistic units BID25 greatly speeds up training and would also satisfy our need for a closed form max. Additionally, this work imposes a prespecified ordering of actions which may negatively impact training for certain classes of problems (with much larger number of action dimensions). To address this, we could learn to factor the action space into the sequential order for continuous action spaces or learn to group action sets for discrete action spaces. Another promising direction is to combine this approximate max action with gradient based optimization procedure. This would relieve some of the complexity of the modeling task of the maxing network, at the cost of increased compute when sampling from the policy. Finally, the work presented here is exclusively on off-policy methods. We chose to focus on these methods due to their sample efficiency. Use of an sequential policies with discretized actions could also be used as the policy for any stochastic policy optimization algorithm such as TRPO BID27 or A3C . In this work we present a continuous control algorithm that utilize discretized action spaces and sequential models. The technique we propose is an off-policy RL algorithm that utilizes sequential prediction and discretization. We decompose our model into a hierarchy of Q function. The effectiveness of our method is demonstrated on illustrative and benchmark tasks, as well as on more complex continuous control tasks. Sampling an Action To gain insight into the characteristics of Q that our SDQN algorithm learns, we visualized from the hopper environment as it is complex but has a small dimensional action space. Figure B.3: Exploration of the sub-DQN during after training. The top row shows the Q i predictions for a given frame (action dimensions correspond to the joint starting at the top and moving toward the bottom -action 3 is the ankle joint). The bottom row shows the corresponding rendering of the current state. For insensitive parts of the gait, such as when the hopper is in the air (e.g. frame 430, 442, 490, 502), the network learns to be agnostic to the choice of actions; this is reflected in the flat Q-value distribution, viewed as a function of action index. On the other hand, for critical parts of the gait, such as when the hopper is in contact with the ground (e.g. frames 446, 478), the network learns that certain actions are much better than others, and the Q-distribution is no longer a flat line. This reflects the fact that taking wrong actions in these regimes could lead to bad such as tripping, yielding a lower reward. First we compute each action dimension's Q distribution, Q L / Q i, and compare those distributions to that of the top MDP for the full action dimentionality, Q U. A figure containing these distributions and corresponding state visualization can be found in Figure B.3.For most states in the hopper walk cycle, the Q distribution is very flat. This implies that small changes in the action taken in a state will have little impact on future reward. This makes sense as the system can recover from any action taken in this frame. However, this is not true for all statescertain critical states exist, such as when the hopper is pushing off, where not selecting the correct action value greatly degrades performance. This can be seen in frame 466.Our algorithm is trained with a number of soft constraints. First, if fully converged, we would expect Q i−1 >= Q i as every new sub-action taken should maintain or improve the expected future discounted reward. Additionally, we would expect Q N (s, a) = Q U (s, a) (from equation 6). In the majority of frames these properties seem correct, but there is certainly room for improvement. Next, we attempt to look at Q surfaces in a more global manner. We plot 2D cross sections for each pair of actions and assume the third dimension is zero. Figure B.4 shows the . As seen in the previous visualization, the surface of both the sequential Q surface and the Q U is not smooth, which is expected as the environment action space for Hopper is expected to be highly non-linear. Some regions of the surface seem quite noisy which is not expected. Interestingly though, these regions of noise do not seem to lower the performance of the final policy. In Q-learning, only the maximum Q value regions have any impact on the taken policy. Future work is needed to better characterize this effect. We would like to explore techniques that use "soft"; BID28; BID14. These techniques will use more of the Q surface thus smooth the representations. Additionally, we notice that the dimensions of the autoregressive model are modeled differently. The last action, a 3 has considerably more noise than the previous two action dimensions. This large difference in the smoothness and shape of the surfaces demonstrates that the order of the actions dimensions matters. This figure suggests that the model has a harder time learning sharp features in the a 1 dimension. In future work, we would like to explore learned orderings, or bidirectional models, to combat this. Finally, the form of Q U is extremely noisy and has many cube artifacts. The input of this function is both a one hot quantized action, as well as the floating point representation. It appears the model uses the quantization as its main feature and learns a sharp Q surface. The most sensitive hyper parameters were the learning rate of the two networks, reward scale, and finally, discount factor. Parameters such as batch size, quantization amounts, and network sizes mattered to a lesser extent. We found it best to have exploration be done on less than 10% of the transitions. We didn't see any particular value below this that gave better performance. In our experiments, the structure of model used also played a large impact in performance, such as, using tied versus untied weights for each sequential prediction model. In future work, we would like to lower the amount of hyper parameters needed in these algorithms and study the effects of various hyper parameters more thoroughly. In this model, we looked at a number of configurations. Of the greatest importance is the form of the model itself. We looked at an LSTM parameterization as well as an untied weight variant. The untied weight variant's parameter search is listed below. To compute Q i we first take the state and previous actions and do one fully connected layer of size "embedding size". We then concatenate these representations and feed them into a 2 hidden layer MLP with "hidden size" units. The output of this network is "quantization bins" wide with no activation. return a

A method to do Q-learning on continuous action spaces by predicting a sequence of discretized 1-D actions.

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

We propose Dual Actor-Critic algorithm, which is derived in a principled way from the Lagrangian dual form of the Bellman optimality equation. The algorithm achieves the state-of-the-art performances across several benchmarks.


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

Question: Model-based reinforcement learning (MBRL) aims to learn a dynamic model to reduce the number of interactions with real-world environments. However, due to estimation error, rollouts in the learned model, especially those of long horizon, fail to match the ones in real-world environments. This mismatching has seriously impacted the sample complexity of MBRL. The phenomenon can be attributed to the fact that previous works employ supervised learning to learn the one-step transition models, which has inherent difficulty ensuring the matching of distributions from multi-step rollouts. Based on the claim, we propose to learn the synthesized model by matching the distributions of multi-step rollouts sampled from the synthesized model and the real ones via WGAN. We theoretically show that matching the two can minimize the difference of cumulative rewards between the real transition and the learned one. Our experiments also show that the proposed model imitation method outperforms the state-of-the-art in terms of sample complexity and average return. Reinforcement learning (RL) has become of great interest because plenty of real-world problems can be modeled as a sequential decision-making problem. Model-free reinforcement learning (MFRL) is favored by its capability of learning complex tasks when interactions with environments are cheap. However, in the majority of real-world problems, such as autonomous driving, interactions are extremely costly, thus MFRL becomes infeasible. One critique about MFRL is that it does not fully exploit past queries over the environment, and this motivates us to consider the model-based reinforcement learning (MBRL). In addition to learning an agent policy, MBRL also uses the queries to learn the dynamics of the environment that our agent is interacting with. If the learned dynamic is accurate enough, the agent can acquire the desired skill by simply interacting with the simulated environment, so that the number of samples to collect in the real world can be greatly reduced. As a , MBRL has become one of the possible solutions to reduce the number of samples required to learn an optimal policy. Most previous works of MBRL adopt supervised learning with 2 -based errors (; ; or maximum likelihood , to obtain an environment model that synthesizes real transitions. These non-trivial developments imply that optimizing a policy on a synthesized environment is a challenging task. Because the estimation error of the model accumulates as the trajectory grows, it is hard to train a policy on a long synthesized trajectory. On the other hand, training on short trajectories makes the policy short-sighted. This issue is known as the planning horizon dilemma . As a , despite having a strong intuition at first sight, MBRL has to be designed meticulously. Intuitively, we would like to learn a transition model in a way that it can reproduce the trajectories that have been generated in the real world. Since the attained trajectories are sampled according to a certain policy, directly employing supervised learning may not necessarily lead to the mentioned especially when the policy is stochastic. The resemblance in trajectories matters because we estimate policy gradient by generating rollouts; however, the one-step model learning adopted by many MBRL methods do not guarantee this. Some previous works propose multi-step training (; ;); however, experiments show that model learning fails to benefit much from the multi-step loss. We attribute this outcome to the essence of super- Supervised learning Rollout from real world Figure 1: Distribution matching enables the learned transition to generate similar rollouts to the real ones even when the policy is stochastic or the initial states are close. On the other hand, training with supervised learning does not ensure rollout similarity and the ing policy gradient may be inaccurate. This figure considers a fixed policy sampling in the real world and a transition model. vised learning, which elementally preserves only one-step transition and the similarity between real trajectories and the synthesized ones cannot be guaranteed. In this work, we propose to learn the transition model via distribution matching. Specifically, we use WGAN to match the distributions of state-action-next-state triple (s, a, s) in real/learned models so that the agent policy can generate similar trajectories when interacting with either the true transition or the learned transition. Figure 1 illustrates the difference between methods based on supervised learning and distribution matching. Different from the ensemble methods proposed in previous works, our method is capable of generalizing to unseen transitions with only one dynamic model because merely incorporating multiple models does not alter the essence that one-step (or few-step) supervised learning fails to imitate the distribution of multi-step rollouts. Concretely, we gather some transitions in the real world according to a policy. To learn the real transition T, we then sample fake transitions from our synthesized model T with the same policy. The synthesized model serves as the generator in the WGAN framework and there is a critic that discriminates the two transition data. We update the generator and the critic alternatively until the synthesized data cannot be distinguished from the real one, which we will show later that it gives T → T theoretically. Our contributions are summarized below: • We propose an MBRL method called model imitation (MI), which enforces the learned transition model to generate similar rollouts to the real one so that policy gradient is accurate; • We theoretically show that the transition can be learned by MI in the sense that T → T by consistency and the difference in cumulative rewards can be bounded by the training error of WGAN; • To stabilize model learning, we deduce guarantee for our sampling technique and investigate training across WGANs; • We experimentally show that MI is more sample efficient than state-of-the-art MBRL and MFRL methods and outperforms them on four standard tasks. In this section, we introduce our motivation inspired by learning from demonstration (LfD) and give a brief survey of MBRL methods. A straightforward approach to LfD is to leverage behavior cloning (BC), which reduces LfD to a supervised learning problem. Even though learning a policy via BC is time-efficient, it cannot imitate a policy without sufficient demonstration because the error may accumulate without the guidance of expert . Generative Adversarial Imitation Learning (GAIL) is another state-of-the-art LfD method that learns an optimal policy by utilizing generative adversarial training to match occupancy measure . GAIL learns an optimal policy by matching the distribution of the trajectories generated from an agent policy with the distribution of the given demonstration. shows that the two distributions match if and only if the agent has learned the optimal policy. One of the advantages of GAIL is that it only requires a small amount of demonstration data to obtain an optimal policy but it requires a considerable number of interactions with environments for the generative adversarial training to converge. Our intuition is that transition learning (TL) is similar to learning from demonstration (LfD) by exchanging the roles of transition and policy. In LfD, trajectories sampled from a fixed transition are given, and the goal is to learn a policy. On the other hand, in TL, trajectories sampled from a fixed policy are given, and we would like to imitate the underlying transition. That being said, from LfD to TL, we interchange the roles of the policy and the transition. It is therefore tempting to study the counterpart of GAIL in TL; i.e., learning the transition by distribution matching. Fortunately, by doing so, the pros of GAIL remain while the cons are insubstantial in MBRL because sampling with the learned model is considered to be much cheaper than sampling in the real one. That GAIL learns a better policy than what BC does suggests that distribution matching possesses the potential to learn a better transition than supervised learning. For deterministic transition, @ -based error is usually utilized to learn the transition model. , an approach that uses supervised learning with mean-squared error as its objective, is shown to perform well under fine-tuning. To alleviate model bias, some previous works adopt ensembles , where multiple transition models with different initialization are trained at the same time. In a slightly more complicated manner, utilizes meta-learning to gather information from multiple models. Lastly, on the theoretical side, SLBO is the first algorithm that develops from solid theoretical properties for model-based deep RL via a joint model-policy optimization framework. For the stochastic transition, maximum likelihood estimator or moment matching are natural ways to learn a synthesized transition, which is usually modeled by the Gaussian distribution. Following this idea, Gaussian process and Gaussian process with model predictive control are introduced as an uncertainty-aware version of MBRL. Similar to the deterministic case, to mitigate model bias and foster stability, an ensemble method for probabilistic networks is also studied. An important distinction between training a deterministic or stochastic transition is that although the stochastic transition can model the noise hidden within the real world, the stochastic model may also induce instability if the true transition is deterministic. This is a potential reason why an ensemble of models is adopted to reduce variance. We consider the standard Markov Decision Process (MDP) . MDP is represented by a tuple S, A, T, r, γ, where S is the state space, A is the action space, T (s t+1 |s t, a t) is the transition density of state s t+1 at time step t + 1 given action a t made under state s t, r(s, a) is the reward function, and γ ∈ is the discount factor. A stochastic policy π(a|s) is a density of action a given state s. Let the initial state distribution be α. The performance of the triple (α, π, T) is evaluated in the expectation of the cumulative reward in the γ-discounted infinite horizon setting: Equivalently, R(α, π, T) is the expected cumulative rewards in a length-H trajectory {s t, a t} H−1 t=0 generated by (α, π, T) with H ∼ Geometric(1 − γ). When α and T are fixed, R(·) becomes a function that only depends on π, and reinforcement learning algorithms aim to find a policy π to maximize R(π). Given initial state distribution α(s), policy π(a|s) and transition T (s |s, a), the normalized occupancy measure ρ where P(·) is the probability measure and will be replaced by a density function if S or A is continuous. Intuitively, ρ α,π T (s, a) is a distribution of (s, a) in a length-H trajectory {s t, a t} H−1 t=0 with H ∼ Geometric(1 − γ) following the laws of (α, π, T). , the relation between ρ α,π T and (α, π, T) is characterized by the Bellman flow constraint. Specifically, x = ρ α,π T as defined in Eq. 2 is the unique solution to: In addition, Theorem 2 of gives that π(a|s) and ρ α,π T (s, a) have an one-to-one correspondence with α(s) and T (s |s, a) fixed; i.e., π(a|s) is the only policy whose occupancy measure is ρ. With the occupancy measure, the cumulative reward Eq. 1 can be represented as The goal of maximizing the cumulative reward can then be achieved by adjusting ρ α,π T, and this motivates us to adopt distribution matching approaches like WGAN to learn a transition model. In this section, we present a consistency and error bounds for WGAN. All proofs of the following theorems and lemmas can be found in Appendix A. In the setting of MBRL, the training objective for WGAN is By Kantorovich-Rubinstein duality , the optimal value of the inner maximization is exactly is the discounted distribution of (s, a, s). Thus, by minimizing over the choice of T, we are essentially finding p that minimizes W 1 (p(s, a, s)||p (s, a, s)), which gives the consistency . Proposition 1 (Consistency for WGAN). Let T and T be the true and synthesized transitions respectively. If WGAN is trained to its optimal point, we have The support constraint is inevitable because the training data is sampled from ρ T and guaranteeing anything beyond it can be difficult. Still, we will empirically show that the support constraint is not an issue in our experiments because the performance boosts up in the beginning, indicating that Supp(ρ α,π T) may be large enough initially. Now that training with WGAN gives a consistent estimate of the true transition, it is sensible to train a synthesized transition upon it. However, the consistency is too restrictive as it only discusses the optimal case. The next step is to analyze the non-optimal situation and observe how the cumulative reward deviates w.r.t. the training error. Theorem 1 indicates that if WGAN is trained properly, i.e., having small, the cumulative reward on the synthesized trajectory will be close to that on the true trajectory. As MBRL aims to train a policy on the synthesized trajectory, the accuracy of the cumulative reward over the synthesized trajectory is thus the bottleneck. Theorem 1 also implies that WGAN's error is linear to the (expected) length of the trajectory (1 − γ) −1. This is a sharp contrast to the error bounds in most RL literature, as the dependency on the trajectory length is usually quadratic , or of an even higher order. Since WGAN gives us a better estimation of the cumulative reward in the learned model, the policy update becomes more accurate. In this section, we present a practical MBRL method called model imitation (MI) that incorporates the transition learning mentioned in Section 4. Due to the long-term digression, it is hard to train the WGAN directly from a long synthesized trajectory. To tackle this issue, we use the synthesized transition T to sample N short trajectories with initial states sampled from the true trajectory. To analyze this sampling technique, let β < γ be the discount factor of the short trajectories so that the expected length is E is upper bounded by the training error of WGAN on short trajectories, which can be small empirically because the short ones are easier to imitate. and Lemma 1, where d is the dimension of (s, a). , where diam(·) is the diameter. The second term encourages β to be large while the third term does the opposite. Besides, β need not be large if N is large enough; in practice we may sample N short trajectories to reduce the error from W 1 (ρ T ||ρ T) to W 1 (ρ β T ||ρ T). Finally, since ρ β T is the occupancy measure we train on, from the proof of Theorem 1 we deduce that. Thus, WGAN may perform better under this sampling technique. To learn the real transition based on the occupancy measure matching mentioned in Section 4, we employ a transition learning scheme by aligning the distribution of (s, a, s) between the real and the learned environments. Inspired by how GAIL learns to align (s, a) via solving an MDP with rewards extracted from a discriminator, we formulate an MDP with rewards from a discriminator over (s, a, s). Specifically, the WGAN critic f (s, a, s) in Eq. 5 is used as the (psuedo) rewards r(s, a, s) of our MDP. Interestingly, there is a duality between GAIL and our transition learning: for GAIL, the transition is fixed and the objective is to train a policy to maximize the cumulative pseudo rewards, while for our transition learning, the policy is fixed and the objective is to train a synthesized transition to maximize the cumulative pseudo rewards. In practice, since the policy is updated alternatively with the synthesized model, we are required to train a number of WGANs along with the change of the policy. Although the generators across WGANs correspond to the same transition and can be similar, we observe that WGAN may get stuck at a local optimum when we switch from one WGAN training to another. The reason is that unlike GAN that mimics the Jensen-Shannon divergence and hence its inner maximization is upper bounded by log, WGAN mimics the Wasserstein distance and the inner maximization is unbounded from above. Intuitively, such unboundedness makes the WGAN critic so strong that the WGAN generator (the synthesized transition) cannot find a way out and get stuck at a local optimum. Thereby, we have to modify the WGAN objective to alleviate such a situation. To ensure the boundedness, for a fixed δ > 0, we introduce cut-offs at the WGAN objective so that the inner maximization is upper bounded by 2δ: As δ → ∞, Eq. 6 recovers the WGAN objective, Eq. 5. Therefore, this is a truncated version of WGAN. To comprehend Eq. 6 further, notice that it is equivalent to which is a hinge loss version of the generative adversarial objective. Such WGAN is introduced in , where the consistency is provided and further experiments are evaluated in. According to , the inner minimization can be interpreted as the soft-margin SVM. Consequently, it provides a geometric intuition of maximizing margin, which potentially enhances robustness. Finally, because the objective of transition learning is to maximize the cumulative pseudo rewards on the MDP, T does not directly optimize Eq. 7. Note that the truncation only takes part in the inner minimization: which gives us a WGAN critic f (s, a, s). As mentioned, f will be the pseudo reward function. Later, we will introduce a transition learning version of PPO to optimize the cumulative pseudo reward. for N epochs do for n policy epochs do 10: update π θ by TRPO on the data generated by T φ 11: end for 12: end for 13: end for After modifying the WGAN objective, to include both the stochastic and (approximately) deterministic scenarios, the synthesized transition is modeled by a Gaussian distribution T (s |s, a) = T φ (s |s, a) ∼ N (µ φ (s, a), Σ φ (s, a) ). Although the underlying transitions of tasks like MuJoCo are deterministic, modeling by a Gaussian does not harm the transition learning empirically. Recall that the synthesized transition is trained on an MDP whose reward function is the critic of the truncated WGAN. To achieve this goal with proper stability, we employ PPO , which is an efficient approximation of TRPO . Note that although the PPO is originally designed for policy optimization, it can be adapted to transition learning with a fixed sampling policy and the PPO objective (Eq. 7 of) where r t (φ) = T φ (s t+1 |s t, a t) T φold (s t+1 |s t, a t),Â t: advantage func. derived from the pseudo reward f (s t, a t, s t+1). To enhance stability of the transition learning, in addition to PPO, we also optimize maximum likelihood, which can be regarded as a regularization. We empirically observe that jointly optimizing both maximum likelihood and the PPO objective attains better transition model for policy gradient. The overall loss of the transition learning becomes where L mle is the loss of MLE, which is policy-agnostic and can be estimated with all collected real transitions. For more implementation details, please see Appendix B.1. We consider a training procedure similar to SLBO , where they consider the fact that the value function is dependent on the varying transition model. As a , unlike most of the MBRL methods that have only one pair of model-policy update for each real environment sampling, SLBO proposes to take multiple update pairs for each real environment sampling. Our proposed model imitation (MI) method is summarized in Algorithm 1. In the experiment section, we would like to answer the following questions. Does the proposed model imitation outperforms the state-of-the-art in terms of sample complexity and average return? Does the proposed model imitation benefit from distribution matching and is superior to its model-free and model-based counterparts, TRPO and SLBO? To fairly compare algorithms and enhance reproducibility, we adopt open-sourced environments released along with a model-based benchmark paper , which is based on a physical simulation engine, MuJoCo . Specifically, we evaluate the proposed algorithm MI on four continuous control tasks including Hopper, HalfCheetah, Ant, and Reacher. For hyperparameters mentioned in Algorithm 1 and coefficients such as entropy regularization λ, please refer to Appendix B.2. Table 1: Proportion of bench-marked RL methods that are inferior to MI in terms of 5% t-test. x/y indicates that among y approaches, MI is significantly better than x approaches. The detailed performance can be found in Table 1 of. It should be noted that the reported in We compare to two model-free algorithms, TRPO and PPO , to assess the benefit of utilizing the proposed model imitation since our MI (Algorithm 1) uses TRPO for policy gradient to update the agent policy. We also compare MI to four model-based methods. SLBO gives theoretical guarantees of monotonic improvement for model-based deep RL and proposes to update a joint model-policy objective. PETS propose to employ uncertainty-aware dynamic models with sampling-based uncertainty to capture both aleatoric and epistemic uncertainty. METRPO shows that insufficient data may cause instability and propose to use an ensemble of models to regularize the learning process. STEVE dynamically interpolates among model rollouts of various horizon lengths and favors those whose estimates have lower error. Figure 2 shows the learning curves for all methods. In Hopper, HalfCheetah, and Ant, MI converges fairly fast and learns a policy significantly better than competitors'. In Ant, even though MI does not improve the performance too much from the initial one, the fact that it maintains the average return at around 1,000 indicates that MI can capture a better transition than other methods do with only 5,000 transition data. Even though we do not employ an ensemble of models, the curves show that our learning does not suffer from high variance. In fact, the performance shown in Figure 2 indicates that the variance of MI is lower than that of methods incorporating ensembles such as METRPO and PETS. The questions raised at the beginning of this section can now be answered. The learned model enables TRPO to explore the world without directly access real transitions and therefore TRPO equipped with MI needs much fewer interactions with the real world to learn a good policy. Even though MI is based on the training framework proposed in SLBO, the additional distribution matching component allows the synthesized model to generate similar rollouts to that of the real environments, which empirically gives superior performance because we rely on long rollouts to estimate policy gradient. To better understand the performance presented in Figure 2, we further compare MI with benchmarked RL algorithms recorded in including state-of-the-art MFRL methods such as TD3 and SAC . It should be noted that the reported of are the final performance after 200k time-steps but we only use up to 100k time-steps to train MI. Table 1 indicates that MI significantly outperforms most of the MBRL and MFRL methods with 50% fewer samples, which verifies that MI is more sample-efficient by incorporating distribution matching. We have pointed out that the state-of-the-art methods concentrate on learning synthesized models in a supervised fashion, which does not guarantee that the policy is able to reproduce a similar trajectory in the learned model and therefore the model may not be accurate enough to estimate long rollouts. We have proposed to incorporate WGAN to achieve occupancy measure matching between the real transition and the synthesized model and theoretically shown that matching indicates the closeness in cumulative rewards between the synthesized model and the real environment. To enable stable training across WGANs, we have suggested using a truncated version of WGAN to prevent training from getting stuck at local optimums. The empirical property of WGAN application such as imitation learning indicates its potential to learn the transition with fewer samples than supervised learning. We have confirmed it experimentally by further showing that MI converges much faster and obtains better policy than state-of-the-art model-based and model-free algorithms. A.1 PROOF FOR WGAN Proposition 1 (Consistency for WGAN). Let α(s), π(a|s), T (s |s, a) be initial state distribution, policy and synthesized transition. Let T be the true transition, p(s, a, s) = ρ α,π T (s, a)T (s |s, a) be the discounted distribution of the triple (s, a, s) under the true transition. If the WGAN is trained to its optimal point, we have Proof. Because the loss function of WGAN is the 1-Wasserstein distance, we know p(s, a, s) = p (s, a, s) at its optimal points. Plug in to the Bellman flow constraint Eq., where the first inequality holds because r(s, a)/L r is 1-Lipschitz and the last equality follows from Kantorovich-Rubinstein duality. Since W 1 distance is symmetric, the same holds if we interchange T and T, so we arrive at |R(π, T) − R(π, T)| ≤ L r /(1 − γ). To enhance state exploration, we sample real transitions according to policy β ∼ N (µ θ (s), σ), where µ(s) is the mean of our Gaussian parameterized policy π θ and σ is a fixed standard deviation. In addition, since model the transition as a Gaussian distribution, we found that matching ρ α,π θ T with ρ α,β T is empirically more stable and more sample-efficient than matching ρ α,β T with ρ α,β T. For policy update, it is shown that using the mean µ φ of the Gaussian-parameterized transition can accelerate policy optimization and better balance exploration and exploitation. In order to enforce the Lipschitz constraint to the WGAN critic f, we employ gradient penalty Table 2: List of hyper-parameters adopted in our experiments.

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

Our method incorporates WGAN to achieve occupancy measure matching for transition learning.

Represent the TLDR for finding the corresponding scientific article.

We identify and formalize the memorization problem in meta-learning and solve this problem with novel meta-regularization method, which greatly expand the domain that meta-learning can be  applicable to and effective on.


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

Question: Batch Normalization (BN) and its variants have seen widespread adoption in the deep learning community because they improve the training of deep neural networks. Discussions of why this normalization works so well remain unsettled. We make explicit the relationship between ordinary least squares and partial derivatives computed when back-propagating through BN. We recast the back-propagation of BN as a least squares fit, which zero-centers and decorrelates partial derivatives from normalized activations. This view, which we term {\em gradient-least-squares}, is an extensible and arithmetically accurate description of BN. To further explore this perspective, we motivate, interpret, and evaluate two adjustments to BN. Training deep neural networks has become central to many machine learning tasks in computer vision, speech, and many other application areas. BID10 showed empirically that Batch Normalization (BN) enables deep networks to attain faster convergence and lower loss. Reasons for the effectiveness of BN remain an open question BID12. Existing work towards explaining this have focused on covariate shift; described how BN makes the loss function smoother. This work examines the details of the back-propagation of BN, and recasts it as a least squares fit. This gradient regression zero-centers and decorrelates partial derivatives from the normalized activations; it passes on a scaled residual during back-propagation. Our view provides novel insight into the effectiveness of BN and several existing alternative normalization approaches in the literature. Foremost, we draw an unexpected connection between least squares and the gradient computation of BN. This motivates a novel view that complements earlier investigations into why BN is so effective. Our view is consistent with recent empirical surprises regarding ordering of layers within ResNet residual maps BID5 and within shake-shake regularization branches BID6. Finally, to demonstrate the extensibility of our view, we motivate and evaluate two variants of BN from the perspective of gradient-least-squares. In the first variant, a least squares explanation motivates the serial chaining of BN and Layer Normalization (LN) BID0. In the second variant, regularization of the least-squares leads to a version of BN that performs better on batch size two. In both variants, we provide empirical support on CIFAR-10.In summary, our work presents a view, which we term gradient-least-squares, through which the back-propagation of BN and related work in a neural network can be recast as least squares regression. This regression decomposes gradients into an explained portion and a residual portion; BN back-propagation will be shown to remove the explained portion. Hopefully, gradient-least-squares will be broadly useful in the future design and understanding of neural network components. Figure 1 reviews normalization with batch statistics, and illustrates our main theorem. Gradient-least-squares relates quantities shown in hexagons Figure 1: The left figure reviews, for a single channel at a particular layer within a single batch, notable quantities computed during the forward pass and during back-propagation of BN. Let DISPLAYFORM0 Let L be a function dependent on the normalized activations z i defined for each j by z j = (x j − µ) σ This, along with partial derivatives, are shown in the left figure. Our work establishes a novel identity on the quantities shown in hexagons. The right figure illustrates our main in a scatter plot, in which each pair z i, ∂L ∂z i is shown as a data point in the regression. Consider any particular channel within which {x i} are activations to be normalized in BN moment calculations. BID10 defined BN as DISPLAYFORM0 where σ, µ are batch moments, but b and c are learned per-channel parameters persistent across batches. In BN, the batch dimension and spatial dimensions are marginalized out in the computation of batch moments. For clarity, we consider a simplified version of BN. We ignore the variables b and c in equation 1 responsible for a downstream channel-wise affine transformation. Ignoring b and c is done without loss of generality, since the main observation in this work will focus on the Gaussian normalization and remains agnostic to downstream computations. We also ignore a numerical stability hyperparameter.We examine back-propagation of partial derivatives through this normalization, where µ and σ are viewed as functions of x. Notably, µ and σ are functions of each x i, and thus the division by σ is not affine. We write the normalized output as DISPLAYFORM1 We review ordinary least squares of a single variable with intercept BID2.Let g j = α + βz j + j where α and β are parameters, z and g are observations. z j and g j are entries in z and g respectively. j are i.i.d. Gaussian residuals. We wish to fit α and β α,β = arg min DISPLAYFORM2 The least-squares problem in equation 3 is satisfied byβ = Cov(z, g) DISPLAYFORM3 When z are normalized activations and g are partial derivatives, then Ez = 0 and Var(z) = 1. In this special case, the solution simplifies intô DISPLAYFORM4 Theorem 1 (Normalization gradients are least-squares residuals). Let i ∈ {1 . . . N} be indices over some set of activations {x i}. Then the moment statistics are defined by µ = N i=1x i N and DISPLAYFORM5 Let L be a function dependent on the normalized activations z i defined for each j by z j = (x j − µ) σ. Then, the gradients of L satisfy, for all j ∈ {1, . . ., N}, the following: DISPLAYFORM6 where DISPLAYFORM7 Proof: Normalization gradients are least-squares residuals. The proof involves a derivation of partial derivatives by repeated applications of the chain rule and rules of total derivative. Because {z i} normalized over i has mean 0 and variance 1, the partial derivatives can be rearranged to satisfy the single variable ordinary least squares framework. Fix j. We expand ∂L ∂x j as a linear combination of DISPLAYFORM8 We state ∂z i ∂x j directly. Steps are in Appendix A under Lemma 1. DISPLAYFORM9 Through substitution of equations 10 into 9, we get DISPLAYFORM10 Noting that {z i} normalized over i has mean 0 and variance 1, we recoverβ andα, in the sense of equations 4 and 5, from equation 13. DISPLAYFORM11 Finally, we rearrange equations 15 and 14 into 13 to conclude, as desired, DISPLAYFORM12 During back-propagation of a single batch, the normalization function takes in partial derivatives ∂L ∂z (·), and removes that which can be explained by least squares of ∂L ∂z (·) against z (·). As illustrated in Figure 1, during back-propagation, the residual then divides away σ to become ∂L ∂x (·), the gradient for the unnormalized activations. CALCULATIONS BN aims to control its output to have mean near 0 and variance near 1, normalized over the dataset; this is related to the original explanation termed internal covariate shift BID10. Most existing work that improve or re-purpose BN have focused on describing the distribution of activations. Definition 1. In the context of normalization layers inside a neural network, activations are split into partitions, within which means and variances are computed. We refer to these partitions as normalization partitions. Definition 2. Within the context of a normalization partition, we refer to the moments calculated on the partitions as partition statistics. Theorem 1 shows that BN has least squares fitting built into the gradient computation. Gradients of the activations being normalized in each batch moment calculation are fit with a single-variable with-intercept least squares model, and only a rescaled residual is kept during back-propagation. We emphasize that the data on which the regression is trained and applied is a subset of empirical activations within a batch, corresponding to the normalization partitions of BN.To show extensibility, we recast several popular normalization techniques into the gradient-leastsquares view. We refer to activations arising from a single member of a particular batch as an item. BHW C refers to dimensions corresponding to items, height, width, and channels respectively. In non-image applications or fully connected layers, H and W are 1. BN marginalizes out the items and spatial dimensions, but statistics for each channel are kept separate. In the subsequent sections, we revisit several normalization methods from the perspective of the gradient. FIG1 reviews the normalization partitions of these methods, and places our main theorem about gradient-least-squares into context. BID0 introduced Layer Normalization (LN) in the context of large LSTM models and recurrent networks. Only the (H, W, C) dimensions are marginalized in LN, whereas BN marginalizes out the (B, H, W) dimensions. In our regression framework, the distinction can be understood as changing the data point partitions in which least squares are fit during back-propagation. LN marginalizes out the channels, but computes separate statistics for each batch item. To summarize, the regression setup in the back-propagation of LN is performed against other channels, rather than against other batch items. introduced Instance Normalization (IN) in the context of transferring styles across images. IN is is closely related to contrast normalization, an older technique used in image processing. IN emphasizes end-to-end training with derivatives passing through the moments. Only the (H, W) dimensions are marginalized in IN, whereas BN marginalizes (B, H, W) dimensions. In our framework, this can be understood as using fewer data points and a finer binning to fit the least squares during back-propagation, as each batch item now falls into its own normalization partition. introduced Group Normalization (GN) to improve performance on image-related tasks when memory constrains the batch size. Similar to LN, GN also marginalizes out the (H, W, C) dimensions in the moment computations. The partitions of GN are finer: the channels are grouped into disjoint sub-partitions, and the moments are computed for each sub-partition. When the number of groups is one, GN reduces to LN.In future normalization methods that involve normalizing with respect to different normalization partitions; such methods can pattern match with BN, LN, IN, or GN; the back-propagation can be formulated as a least-squares fit, in which the partial derivatives at normalized activations ∂L ∂z (·) are fitted against the normalized z (·), and then the residual of the fit is rescaled to become ∂L ∂x (·).Figure 2 summarize the normalization partitions for BN, LN, IN, and GN; the figure visualizes, as an example, a one-to-one correspondence between an activation in BN, and a data point in the gradient regression. Theorem 1 is agnostic to the precise nature of how activations are partitioned before being normalized; thus, equation 9 applies directly to any method that partitions activations and performs Gaussian normalization on each partition. DISPLAYFORM0 DISPLAYFORM1 The L2 normalization of weights in WN appears distinct from the Gaussian normalization of activations in BN; nevertheless, WN can also be recast as a least squares regression. BID4 and improved BID5 residual mappings in ResNets. Arrows point in the direction of the forward pass. Dotted lines indicate that gradients are zero-centered and decorrelated with respect to downstream activations in the residual mapping. The improved ordering has BN coming first, and thus constrains that gradients of the residual map must be decorrelated with respect to some normalized activations inside the residual mapping. An update to the popular ResNet architecture showed that the network's residual mappings can be dramatically improved with a new ordering BID5. The improvement moved BN operations into early positions and surprised the authors; we support the change from the perspective of gradient-least-squares. FIG2 reviews the precise ordering in the two versions. BID6 provides independent empirical support for the BN-early order, in shake-shake regularization BID3 architectures. We believe that the surprise arises from a perspective that views BN only as a way to control the distribution of activations; one would place BN after a sequence of convolution layers. In the gradient-least-squares perspective, the first layer of each residual mapping is also the final calculation for these gradients before they are added back into BID0 0.9102 0.3548 the main trunk. The improved residual branch constrains the gradients returning from the residual mappings to be zero-centered and decorrelated with respect to some activations inside the branch. We illustrate this idea in FIG2. Gradient-least-squares views back-propagation in deep neural networks as a solution to a regression problem. Thus, formulations and ideas from a regression perspective would motivate improvements and alternatives to BN. We pursue and evaluate two of these ideas. BN and LN are similar to each other, but they normalize over different partitioning of the activations; in back-propagation, the regressions occur respectively with respect to different partitions of the activations. Suppose that a BN and a LN layer are chained serially in either order. This in a two-step regression during back-propagation; in reversed order, the residual from the first regression is further explained by a second regression on a different partitioning. In principle, whether this helps would depend on the empirical characteristics of the gradients encountered during training. The second regression could further decorrelate partial gradients from activations. Empirically, we show improvement in a reference ResNet-34-v2 implementation on CIFAR-10 relative to BN with batch size 128. In all cases, only a single per-channel downstream affine transformation is applied, after both normalization layers, for consistency in the number of parameters. See table 1 for CIFAR-10 validation performances. We kept all default hyperparameters from the reference implementation: learning schedules, batch sizes, and optimizer settings. BN performs less well on small batches BID9. Gradient-least-squares interprets this as gradient regressions failing on correlated data, an issue typically addressed by regularization. We pursue this idea to recover some performance on small batches by use of regularization. Our regularization uses streaming estimates of past gradients to create virtual data in the regression. This performed better than standard BN on the same batch size, but we did not recover the performance of large batches; this is consistent with the idea that regularization could not in general compensate for having much less data. See Appendix C for CIFAR-10 validation performances. DISPLAYFORM0 to rescale the contributions to the batch mean for each normalization scheme. It uses an analogous set of parameters λ k and activations w k for variances. We sketch the back-propagation of a simplified version of SN in the perspective of gradient-least-squares. We ignore both the division and downstream affine z → c · z + b. The normalization calculation inside SwN can be written as: DISPLAYFORM1 where Ω = {BN, LN, IN}. There is potentially a unique mean and variance used for each activation. Equation 19 bears similarities to the setup in Theorem 1, but we leave unresolved whether there is a gradient-least-squares regression interpretation for SN. Decorrelated Batch Normalization (DBN) ) is a generalization of BN that performs Mahalanobis ZCA whitening to decorrelate the channels, using differentiable operations. On some level, the matrix gradient equation resemble the least squares formulation in Theorem 1.Spectral Normalization (SpN) BID15 is an approximate spectral generalization of WN. For DBN and SpN, the regression interpretations remain unresolved. BN has been instrumental in the training of deeper networks BID10. Subsequent work ed in Batch Renormalization BID9, and further emphasized the importance of passing gradients through the minibatch moments, instead of a gradient-free exponential running average. In gradient-least-squares, use of running accumulators in the training forward pass would stop the gradients from flowing through them during training, and there would be no least-squares. BID5 demonstrate empirically the unexpected advantages of placing BN early in residual mappings of ResNet. Santurkar et al. FORMULA1 showed that BN makes the loss landscape smoother, and gradients more predictable across stochastic gradient descent steps. BID1 found evidence that spatial correlation of gradients explains why ResNet outperforms earlier designs of deep neural networks. BID11 proved that BN accelerates convergence on least squares loss, but did not consider back-propagation of BN as a least squares residual. BID14 has recast BN as a stochastic process, ing in a novel treatment of regularization. This work makes explicit how BN back-propagation regresses partial derivatives against the normalized activations and keeps the residual. This view, in conjunction with the empirical success of BN, suggests an interpretation of BN as a gradient regression calculation. BN and its variants decorrelate and zero-center the gradients with respect to the normalized activations. Subjectively, this can be viewed as removing systematic errors from the gradients. Our view also support empirical in literature preferring early BN placement within neural network branches. Leveraging gradient-least-squares considerations, we ran two sets of normalization experiments, applicable to large batch and small batch settings. Placing a LN layer either before or after BN can be viewed as two-step regression that better explains the residual. We show empirically on CIFAR-10 that BN and LN together are better than either individually. In a second set of experiments, we address BN's performance degradation with small batch size. We regularize the gradient regression with streaming gradient statistics, which empirically recovers some performance on CIFAR-10 relative to basic BN, on batch size two. Why do empirical improvements in neural networks with BN keep the gradient-least-squares residuals and drop the explained portion? We propose two open approaches for investigating this in future work. A first approach focuses on how changes to the gradient regression in different formulations; the two empirical experiments in our work contribute to this. A second approach examines the empirical relationships between gradients of activations evaluated on the same parameter values; we can search for a shared noisy component arising from gradients in the same normalization partition. Suppose that the gradient noise correlates with the activations -this is plausible because the population of internal activations arise from using shared weights -then normalizations could be viewed as a layer that removes systematic noise during back-propagation. In DISPLAYFORM0 Then, the partial derivatives satisfy DISPLAYFORM1 Proof. In deriving ∂z j ∂x i, we will treat the cases of when j = i and when j = i separately. We start by examining intermediate quantities of interest as a matter of convenience for later use. We define helper quantities u i = x i − µ. Note that each u j depends on all of x i via µ. Next, we write out useful identities DISPLAYFORM2 We prepare to differentiate with rule of total derivative: DISPLAYFORM3 Making use of equations 21, 22, 23 and 25, We simplify ∂σ ∂x i for any i as follows. DISPLAYFORM4 We apply the quotient rule on ∂z j ∂x i when j = i, then substitute equation 33 DISPLAYFORM5 Similarly, when i = j, inputs in batch b. In our work, we keep track of am exponential running estimates across batches, DISPLAYFORM6 DISPLAYFORM7 DISPLAYFORM8 that marginalize the (B, H, W) dimensions into accumulators of shape C. The b subscript of the outer expectation is slightly abusive notation indicating thatα * andβ * are running averages across recent batches with momentum as a hyperparameter that determines the weighting. We regularize the gradient regression with virtual activations and virtual gradients, defined as follows. We append two virtual batch items, broadcast to an appropriate shape, x + = µ b + σ b and x − = µ b − σ b. Here, µ b and σ b are batch statistics of the real activations. The concatenated tensor undergoes standard BN, which outputs the usual {z i} for the real activations, but z + = 1 and z − = −1 for the virtual items. The z + and z − do not affect the feed forward calculations, but they receive virtual gradients during back-propagation: DISPLAYFORM9 Virtual data z +, ∂L ∂z + and z −, ∂L ∂z − regularizes the gradient-least-squares regression. ∂L ∂z + and ∂L ∂z − eventually modify the gradients received by the real x i activations. The virtual data can be weighted with hyperparameters. In our experiments, we see improvements, robust to a hyperparameter cross-product search over the weightings and the momentum forα * andβ *. The momentum forα * andβ * were in {.997, .5} and the virtual item weights were in {2 i−1} i∈{0,1,2,3}. The performance of larger batches are not recovered; regularized regression could not be reasonably expected to recover the performance of regressing with more data. See table 2 for final validation performances with a reference Tensorflow ResNet-34-v2 implementation on batch size of two. The baseline evaluation with identity (no normalization) experienced noticeable overfitting in terms of cross entropy but not accuracy. The base learning rate was multiplied by 1 64 relative to the baseline rate used in runs with batch size 128.

Gaussian normalization performs a least-squares fit during back-propagation, which zero-centers and decorrelates partial derivatives from normalized activations.

We show shortcut connections should be placed in patterns that minimize between-layer distances during backpropagation, and design networks that achieve log L distances using L log(L) connections.