summaries: The paper presents a new approach for model-based offline RL. Here, the model is not used to estimate the return directly using roll-outs, but to generate synthetic data that is then used to train the Q-function. In a way, the method is a combination of CQL and MOPO. What is new is the idea of using this synthetic data to, penalize extrapolation instead of using uncertainty estimation. This modification performs much better on the selected benchmarks. As benchmarks a selection of D4RL benchmarks is used (3 benchmarks) and two vision based benchmarks.

summaries: This paper proposed a new model-based offline RL algorithm, COMBO. The algorithm applies model-based learning to be able to generate random rollouts, and the generated data is combined with the given fixed dataset for value function training. The paper shows by using the generated rollouts as data for Conservative Q-learning provides better constraint and better empirical performance compared to the orignal CQL algorithm.

summaries: This paper introduces a Conservative Offline Model-Based policy Optimization RL algorithm, called COMBO, that learns a pessimistic model by enabling a lower bound optimization on the policy performance without requiring uncertainty quantification. COMBO employs an actor-critic method to learn the value function on both offline and synthetic datasets. The idea is to penalize Q-values learned from state-action pairs that are out-of-support of the offline dataset.

summaries: This work investigates model-based reinforcement learning with offline dataset. In particular, most of the work on mbrl with offline datasets requires uncertainty estimation to determine the out-of-distribution states and actions. This work eliminates the need to have uncertainty estimation by using the existing Conservative Q-Learning(CQL) algorithm for states and actions from the dataset as well as the rollouts generated from the dynamics using the policy being learned. They show that inclusion of new dataset from the generated rollouts and by being pessimistic about their values, they tend to be better than existing Model-based offline Reinforcement Learning which requires uncertainty estimation

abstractText: Model-based reinforcement learning (RL) algorithms, which learn a dynamics model from logged experience and perform conservative planning under the learned model, have emerged as a promising paradigm for offline reinforcement learning (offline RL). However, practical variants of such model-based algorithms rely on explicit uncertainty quantification for incorporating conservatism. Uncertainty estimation with complex models, such as deep neural networks, can be difficult and unreliable. We empirically find that uncertainty estimation is not accurate and leads to poor performance in certain scenarios in offline model-based RL. We overcome this limitation by developing a new model-based offline RL algorithm, COMBO, that trains a value function using both the offline dataset and data generated using rollouts under the model while also additionally regularizing the value function on out-of-support state-action tuples generated via model rollouts. This results in a conservative estimate of the value function for out-of-support state-action tuples, without requiring explicit uncertainty estimation. Theoretically, we show that COMBO satisfies a policy improvement guarantee in the offline setting. Through extensive experiments, we find that COMBO attains greater performance compared to prior offline RL on problems that demand generalization to related but previously unseen tasks, and also consistently matches or outperforms prior offline RL methods on widely studied offline RL benchmarks, including image-based tasks.

authors:
Alekh Agarwal
Nan Jiang
Sham M Kakade

title: Reinforcement learning: Theory and algorithms

venue: CS Dept., UW Seattle,

authors:
Rishabh Agarwal
Dale Schuurmans
Mohammad Norouzi

title: An optimistic perspective on offline reinforcement learning

venue: In International Conference on Machine Learning,

authors:
Arthur Argenson
Gabriel Dulac-Arnold

title: Model-based offline planning

venue: arXiv preprint arXiv:2008.05556,

authors:
Kamyar Azizzadenesheli
Emma Brunskill
Animashree Anandkumar

title: Efficient exploration through bayesian deep q-networks

venue: In ITA,

authors:
Ignasi Clavera
Violet Fu
Pieter Abbeel

title: Model-augmented actor-critic: Backpropagating through paths

venue: arXiv preprint arXiv:2005.08068,

authors:
Thomas Degris
Martha White
Richard S Sutton

title: Off-policy actor-critic

venue: arXiv preprint arXiv:1205.4839,

authors:
Frederik Ebert
Chelsea Finn
Sudeep Dasari
Annie Xie
Alex Lee
Sergey Levine

title: Visual foresight: Model-based deep reinforcement learning for vision-based robotic control

venue: arXiv preprint arXiv:1812.00568,

authors:
Damien Ernst
Pierre Geurts
Louis Wehenkel

title: Tree-based batch mode reinforcement learning

venue: Journal of Machine Learning Research,

authors:
Chelsea Finn
Sergey Levine

title: Deep visual foresight for planning robot motion

venue: IEEE International Conference on Robotics and Automation (ICRA),

authors:
Justin Fu
Aviral Kumar
Ofir Nachum
George Tucker
Sergey Levine

title: D4rl: Datasets for deep data-driven reinforcement learning, 2020

authors:
Scott Fujimoto
Shixiang Shane Gu

title: A minimalist approach to offline reinforcement learning

venue: arXiv preprint arXiv:2106.06860,

authors:
Scott Fujimoto
David Meger
Doina Precup

title: Off-policy deep reinforcement learning without exploration

venue: arXiv preprint arXiv:1812.02900,

authors:
Scott Fujimoto
Herke Van Hoof
David Meger

title: Addressing function approximation error in actor-critic methods

venue: arXiv preprint arXiv:1802.09477,

authors:
Seyed Kamyar Seyed Ghasemipour
Dale Schuurmans
Shixiang Shane Gu

title: Emaq: Expected-max q-learning operator for simple yet effective offline and online rl

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor

venue: arXiv preprint arXiv:1801.01290,

authors:
Thomas Jaksch
Ronald Ortner
Peter Auer

title: Near-optimal regret bounds for reinforcement learning

authors:
Michael Janner
Justin Fu
Marvin Zhang
Sergey Levine

title: When to trust your model: Model-based policy optimization

venue: In Advances in Neural Information Processing Systems,

authors:
Natasha Jaques
Asma Ghandeharioun
Judy Hanwen Shen
Craig Ferguson
Agata Lapedriza
Noah Jones
Shixiang Gu
Rosalind Picard

title: Way off-policy batch deep reinforcement learning of implicit human preferences in dialog

venue: arXiv preprint arXiv:1907.00456,

authors:
Natasha Jaques
Judy Hanwen Shen
Asma Ghandeharioun
Craig Ferguson
Agata Lapedriza
Noah Jones
Shixiang Shane Gu
Rosalind Picard

title: Human-centric dialog training via offline reinforcement learning

authors:
Ying Jin
Zhuoran Yang
Zhaoran Wang

title: Is pessimism provably efficient for offline rl

authors:
Gregory Kahn
Adam Villaflor
Pieter Abbeel
Sergey Levine

title: Composable actionconditioned predictors: Flexible off-policy learning for robot navigation

venue: In Conference on Robot Learning,

authors:
Dmitry Kalashnikov
Alex Irpan
Peter Pastor
Julian Ibarz
Alexander Herzog
Eric Jang
Deirdre Quillen
Ethan Holly
Mrinal Kalakrishnan
Vincent Vanhoucke

title: Scalable deep reinforcement learning for vision-based robotic manipulation

authors:
Rahul Kidambi
Aravind Rajeswaran
Praneeth Netrapalli
Thorsten Joachims

title: Morel: Model-based offline reinforcement learning

venue: arXiv preprint arXiv:2005.05951,

authors:
Ilya Kostrikov
Rob Fergus
Jonathan Tompson
Ofir Nachum

title: Offline reinforcement learning with fisher divergence critic regularization

authors:
Aviral Kumar
Justin Fu
Matthew Soh
George Tucker
Sergey Levine

title: Stabilizing off-policy q-learning via bootstrapping error reduction

authors:
Aviral Kumar
Aurick Zhou
George Tucker
Sergey Levine

title: Conservative q-learning for offline reinforcement learning

venue: arXiv preprint arXiv:2006.04779,

authors:
Sascha Lange
Thomas Gabel
Martin A. Riedmiller

title: Batch reinforcement learning

venue: In Reinforcement Learning,

authors:
Romain Laroche
Paul Trichelair
Remi Tachet Des Combes

title: Safe policy improvement with baseline bootstrapping

authors:
Alex X. Lee
Anusha Nagabandi
Pieter Abbeel
Sergey Levine

title: Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model

authors:
Byung-Jun Lee
Jongmin Lee
Kee-Eung Kim

title: Representation balancing offline modelbased reinforcement learning

venue: In International Conference on Learning Representations,

authors:
Sergey Levine
Aviral Kumar
George Tucker
Justin Fu

title: Offline reinforcement learning: Tutorial, review, and perspectives on open problems

venue: arXiv preprint arXiv:2005.01643,

authors:
Yao Liu
Adith Swaminathan
Alekh Agarwal
Emma Brunskill

title: Off-policy policy gradient with state distribution correction

title: Provably good batch reinforcement learning without great exploration

venue: arXiv preprint arXiv:2007.08202,

authors:
Kendall Lowrey
Aravind Rajeswaran
Sham Kakade
Emanuel Todorov
Igor Mordatch

title: Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Control

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

authors:
Ajay Mandlekar
Fabio Ramos
Byron Boots
Silvio Savarese
Li Fei-Fei
Animesh Garg
Dieter Fox

title: Iris: Implicit reinforcement without interaction at scale for learning control from offline robot manipulation data

authors:
Tatsuya Matsushima
Hiroki Furuta
Yutaka Matsuo
Ofir Nachum
Shixiang Gu

title: Deployment-efficient reinforcement learning via model-based offline optimization

venue: arXiv preprint arXiv:2006.03647,

authors:
Rémi Munos
Csaba Szepesvari

title: Finite-time bounds for fitted value iteration

venue: J. Mach. Learn. Res.,

authors:
Ofir Nachum
Bo Dai
Ilya Kostrikov
Yinlam Chow
Lihong Li
Dale Schuurmans

title: Algaedice: Policy gradient from arbitrary experience

authors:
Ian Osband
Benjamin Van Roy

title: Why is posterior sampling better than optimism for reinforcement learning

authors:
Ian Osband
John Aslanides
Albin Cassirer

title: Randomized prior functions for deep reinforcement learning

authors:
Yaniv Ovadia
Emily Fertig
Jie Ren
Zachary Nado
David Sculley
Sebastian Nowozin
Joshua V Dillon
Balaji Lakshminarayanan
Jasper Snoek

title: Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift

venue: arXiv preprint arXiv:1906.02530,

authors:
Xue Bin Peng
Aviral Kumar
Grace Zhang
Sergey Levine

title: Advantage-weighted regression: Simple and scalable off-policy reinforcement learning

authors:
Marek Petrik
Yinlam Chow
Mohammad Ghavamzadeh

title: Safe policy improvement by minimizing robust baseline regret

venue: arXiv preprint arXiv:1607.03842,

authors:
Doina Precup
Richard S Sutton
Sanjoy Dasgupta

title: Off-policy temporal-difference learning with function approximation

venue: In ICML, pages 417–424,

authors:
Rafael Rafailov
Tianhe Yu
A. Rajeswaran
Chelsea Finn

title: Offline reinforcement learning from images with latent space models

authors:
Paria Rashidinejad
Banghua Zhu
Cong Ma
Jiantao Jiao
Stuart Russell

title: Bridging offline reinforcement learning and imitation learning: A tale of pessimism

venue: arXiv preprint arXiv:2103.12021,

title: Neural fitted q iteration–first experiences with a data efficient neural reinforcement learning method

venue: In European Conference on Machine Learning,

authors:
Stephane Ross
Drew Bagnell

title: Agnostic system identification for model-based reinforcement learning

venue: In ICML,

authors:
Susan M Shortreed
Eric Laber
Daniel J Lizotte
T Scott Stroup
Joelle Pineau
Susan A Murphy

title: Informing sequential clinical decision-making through reinforcement learning: an empirical study

venue: Machine learning,

authors:
Noah Y Siegel
Jost Tobias Springenberg
Felix Berkenkamp
Abbas Abdolmaleki
Michael Neunert
Thomas Lampe
Roland Hafner
Martin Riedmiller

title: Keep doing what worked: Behavioral modelling priors for offline reinforcement learning

venue: arXiv preprint arXiv:2002.08396,

authors:
Avi Singh
Albert Yu
Jonathan Yang
Jesse Zhang
Aviral Kumar
Sergey Levine

title: Cog: Connecting new skills to past experience with offline reinforcement learning

authors:
Samarth Sinha
Animesh Garg

title: S4rl: Surprisingly simple self-supervision for offline reinforcement learning

venue: arXiv preprint arXiv:2103.06326,

authors:
R.S. Sutton
A.G. Barto

title: Reinforcement Learning: An Introduction

title: Dyna, an integrated architecture for learning, planning, and reacting

venue: ACM Sigart Bulletin,

authors:
Richard S Sutton
A Rupam Mahmood
Martha White

title: An emphatic approach to the problem of off-policy temporal-difference learning

venue: The Journal of Machine Learning Research,

authors:
Adith Swaminathan
Thorsten Joachims

title: Batch learning from logged bandit feedback through counterfactual risk minimization

venue: J. Mach. Learn. Res,

authors:
Phillip Swazinna
Steffen Udluft
Thomas Runkler

title: Overcoming model bias for robust offline deep reinforcement learning

venue: arXiv preprint arXiv:2008.05533,

authors:
L. Wang
Wei Zhang
Xiaofeng He
H. Zha

title: Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation

venue: Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining,

authors:
Yifan Wu
George Tucker
Ofir Nachum

title: Behavior regularized offline reinforcement learning

venue: arXiv preprint arXiv:1911.11361,

authors:
Yue Wu
Shuangfei Zhai
Nitish Srivastava
Joshua Susskind
Jian Zhang
Ruslan Salakhutdinov
Hanlin Goh

title: Uncertainty weighted actor-critic for offline reinforcement learning

venue: arXiv preprint arXiv:2105.08140,

authors:
F. Yu
H. Chen
X. Wang
Wenqi Xian
Yingying Chen
Fangchen Liu
V. Madhavan
Trevor Darrell

title: Bdd100k: A diverse driving dataset for heterogeneous multitask learning

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

authors:
Tianhe Yu
Deirdre Quillen
Zhanpeng He
Ryan Julian
Karol Hausman
Chelsea Finn
Sergey Levine

title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning

authors:
Tianhe Yu
Garrett Thomas
Lantao Yu
Stefano Ermon
James Zou
Sergey Levine
Chelsea Finn
Tengyu Ma

title: Mopo: Model-based offline policy optimization

venue: arXiv preprint arXiv:2005.13239,

authors:
Xianyuan Zhan
Xiangyu Zhu
Haoran Xu

title: Model-based offline planning with trajectory pruning

venue: arXiv preprint arXiv:2105.07351,

authors:
Wenxuan Zhou
Sujay Bajracharya
David Held. Plas

title: Latent action space for offline reinforcement learning

venue: arXiv preprint arXiv:2011.07213,

heading: 1 Introduction

text:
Offline reinforcement learning (offline RL) [30, 34] refers to the setting where policies are trained using static, previously collected datasets. This presents an attractive paradigm for data reuse and safe policy learning in many applications, such as healthcare [62], autonomous driving [65], robotics [25, 48], and personalized recommendation systems [59]. Recent studies have observed that RL algorithms originally developed for the online or interactive paradigm perform poorly in the offline case [14, 28, 26]. This is primarily attributed to the distribution shift that arises over the course of learning between the offline dataset and the learned policy. Thus, development of algorithms specialized for offline RL is of paramount importance to benefit from the offline data available in aformentioned applications. In this work, we develop a principled model-based offline RL algorithm that matches or exceeds the performance of prior offline RL algorithms in benchmark tasks.
A major paradigm for algorithm design in offline RL is to incorporate conservatism or regularization into online RL algorithms. Model-free offline RL algorithms [15, 28, 63, 21, 29, 27] directly incorporate conservatism into the policy or value function training and do not require learning a dynamics model. However, model-free algorithms learn only on the states in the offline dataset, which can lead to overly conservative algorithms. In contrast, model-based algorithms [26, 67] learn a pessimistic dynamics model, which in turn induces a conservative estimate of the value function. By generating and training on additional synthetic data, model-based algorithms have the potential for
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
broader generalization and solving new tasks using the offline dataset [67]. However, these methods rely on some sort of strong assumption about uncertainty estimation, typically assuming access to a model error oracle that can estimate upper bounds on model error for any state-action tuple. In practice, such methods use more heuristic uncertainty estimation methods, which can be difficult or unreliable for complex datasets or deep network models. It then remains an open question as to whether we can formulate principled model-based offline RL algorithms with concrete theoretical guarantees on performance without assuming access to an uncertainty or model error oracle. In this work, we propose precisely such a method, by eschewing direct uncertainty estimation, which we argue is not necessary for offline RL.
Our main contribution is the development of conservative offline model-based policy optimization (COMBO), a new model-based algorithm for offline RL. COMBO learns a dynamics model using the offline dataset. Subsequently, it employs an actor-critic method where the value function is learned using both the offline dataset as well as synthetically generated data from the model, similar to Dyna [57] and a number of recent methods [20, 67, 7, 48]. However, in contrast to Dyna, COMBO learns a conservative critic function by penalizing the value function in state-action tuples that are not in the support of the offline dataset, obtained by simulating the learned model. We theoretically show that for any policy, the Q-function learned by COMBO is a lower-bound on the true Q-function. While
the approach of optimizing a performance lower-bound is similar in spirit to prior model-based algorithms [26, 67], COMBO crucially does not assume access to a model error or uncertainty oracle. In addition, we show theoretically that the Q-function learned by COMBO is less conservative than model-free counterparts such as CQL [29], and quantify conditions under which the this lower bound is tighter than the one derived in CQL. This is illustrated through an example in Figure 1. Following prior works [31], we show that COMBO enjoys a safe policy improvement guarantee. By interpolating model-free and model-based components, this guarantee can utilize the best of both guarantees in certain cases. Finally, in our experiments, we find that COMBO achieves the best performance on tasks that require out-of-distribution generalization and outperforms previous latent-space offline model-based RL methods on image-based robotic manipulation benchmarks. We also test COMBO on commonly studied benchmarks for offline RL and find that COMBO generally performs well on the benchmarks, achieving the highest score in 9 out of 12 MuJoCo domains from the D4RL [12] benchmark suite.

heading: 2 Preliminaries

text:
Markov Decision Processes and Offline RL. We study RL in the framework of Markov decision processes (MDPs) specified by the tupleM = (S,A, T, r, µ0, γ). S,A denote the state and action spaces. T (s′|s,a) and r(s,a) ∈ [−Rmax, Rmax] represent the dynamics and reward function respectively. µ0(s) denotes the initial state distribution, and γ ∈ (0, 1) denotes the discount factor. We denote the discounted state visitation distribution of a policy π using dπM(s) := (1− γ) ∑∞ t=0 γ
tP(st = s|π), where P(st = s|π) is the probability of reaching state s at time t by rolling out π in M. Similarly, we denote the state-action visitation distribution with dπM(s,a) := d π M(s)π(a|s). The goal of RL is to learn a policy that maximizes the return, or long term cumulative rewards: maxπ J(M, π) := 11−γE(s,a)∼dπM(s,a)[r(s,a)].
Offline RL is the setting where we have access only to a fixed dataset D = {(s,a, r, s′)}, which consists of transition tuples from trajectories collected using a behavior policy πβ . In other words, the dataset D is sampled from dπβ (s,a) := dπβ (s)πβ(a|s). We defineM as the empirical MDP induced by the dataset D and d(s,a) as sampled-based version of dπβ (s,a). In the offline setting, the goal is to find the best possible policy using the fixed offline dataset.
Model-Free Offline RL Algorithms. One class of approaches for solving MDPs involves the use of dynamic programming and actor-critic schemes [56, 5], which do not explicitly require the learning
of a dynamics model. To capture the long term behavior of a policy without a model, we define the action value function as Qπ(s,a) := E [ ∑∞ t=0 γ
t r(st,at) | s0 = s,a0 = a] , where future actions are sampled from π(·|s) and state transitions happen according to the MDP dynamics. Consider the following Bellman operator: BπQ(s,a) := r(s,a) + γE s′∼T (·|s,a),a′∼π(·|s′) [Q(s′,a′)], and its sample based counterpart: B̂πQ(s,a) := r(s,a) + γQ(s′,a′), associated with a single transition (s,a, s′) and a′ ∼ π(·|s′). The action-value function satisfies the Bellman consistency criterion given by BπQπ(s,a) = Qπ(s,a) ∀(s,a). When given an offline dataset D, standard approximate dynamic programming (ADP) and actor-critic methods use this criterion to alternate between policy evaluation [40] and policy improvement. A number of prior works have observed that such a direct extension of ADP and actor-critic schemes to offline RL leads to poor results due to distribution shift over the course of learning and over-estimation bias in the Q function [14, 28, 63]. To address these drawbacks, prior works have proposed a number of modifications aimed towards regularizing the policy or value function (see Section 6). In this work, we primarily focus on CQL [29], which alternates between:
Policy Evaluation: The Q function associated with the current policy π is approximated conservatively by repeating the following optimization:
Qk+1←argmin Q
β ( Es∼D,a∼µ(·|s)[Q(s,a)]−Es,a∼D[Q(s,a)] ) + 1
2 Es,a,s′∼D
[( Q(s,a)−B̂πQk(s,a) ) 2 ] , (1)
where µ(·|s) is a wide sampling distribution such as the uniform distribution over action bounds. CQL effectively penalizes the Q function at states in the dataset for actions not observed in the dataset. This enables a conservative estimation of the value function for any policy [29], mitigating the challenges of over-estimation bias and distribution shift.
Policy Improvement: After approximating the Q function as Q̂π, the policy is improved as π ← argmaxπ′Es∼D,a∼π′(·|s) [ Q̂π(s,a) ] . Actor-critic methods with parameterized policies and Q
functions approximate argmax and argmin in above equations with a few gradient descent steps.
Model-Based Offline RL Algorithms. A second class of algorithms for solving MDPs involve the learning of the dynamics function, and using the learned model to aid policy search. Using the given dataset D, a dynamics model T̂ is typically trained using maximum likelihood estimation as: minT̂ E(s,a,s′)∼D [ log T̂ (s′|s,a) ] . A reward model r̂(s,a) can also be learned similarly if it is un-
known. Once a model has been learned, we can construct the learned MDP M̂ = (S,A, T̂ , r̂, µ0, γ), which has the same state and action spaces, but uses the learned dynamics and reward function. Subsequently, any policy learning or planning algorithm can be used to recover the optimal policy in the model as π̂ = argmaxπ J(M̂, π). This straightforward approach is known to fail in the offline RL setting, both in theory and practice, due to distribution shift and model-bias [51, 26]. In order to overcome these challenges, offline model-based algorithms like MOReL [26] and MOPO [67] use uncertainty quantification to construct a lower bound for policy performance and optimize this lower bound by assuming a model error oracle u(s,a). By using an uncertainty estimation algorithm like bootstrap ensembles [43, 4, 37], we can estimate u(s,a). By constructing and optimizing such a lower bound, offline model-based RL algorithms avoid the aforementioned pitfalls like model-bias and distribution shift. While any RL or planning algorithm can be used to learn the optimal policy for M̂, we focus specifically on MBPO [20, 57] which was used in MOPO. MBPO follows the standard structure of actor-critic algorithms, but in each iteration uses an augmented dataset D ∪Dmodel for policy evaluation. Here, D is the offline dataset and Dmodel is a dataset obtained by simulating the current policy using the learned dynamics model. Specifically, at each iteration, MBPO performs k-step rollouts using T̂ starting from state s ∈ D with a particular rollout policy µ(a|s), adds the model-generated data to Dmodel, and optimizes the policy with a batch of data sampled from D ∪Dmodel where each datapoint in the batch is drawn from D with probability f ∈ [0, 1] and Dmodel with probability 1− f .

heading: 3 Conservative Offline Model-Based Policy Optimization

text:
The principal limitation of prior offline model-based algorithms (discussed in Section 2) is the assumption of having access to a model error oracle for uncertainty estimation and strong reliance on heuristics of quantifying the uncertainty. In practice, such heuristics could be challenging for complex datasets or deep neural network models [44]. We argue that uncertainty estimation is not
Algorithm 1 COMBO: Conservative Model Based Offline Policy Optimization Require: Offline dataset D, rollout distribution µ(·|s), learned dynamics model T̂θ , initialized policy and critic
πϕ and Qψ . 1: Train the probabilistic dynamics model T̂θ(s′, r|s,a) = N (µθ(s,a),Σθ(s,a)) on D. 2: Initialize the replay buffer Dmodel ← ∅. 3: for i = 1, 2, 3, · · · , do 4: Collect model rollouts by sampling from µ and T̂θ starting from states inD. Add model rollouts toDmodel. 5: Conservatively evaluate πiϕ by repeatedly solving eq. 2 to obtain Q̂ πiϕ ψ using samples from D ∪Dmodel. 6: Improve policy under state marginal of df by solving eq. 3 to obtain πi+1ϕ . 7: end for
imperative for offline model-based RL and empirically show that uncertainty estimation could be inaccurate in offline RL problems especially when generalization to unknown behaviors is required in Section 5.1.1. Our goal is to develop a model-based offline RL algorithm that enables optimizing a lower bound on the policy performance, but without requiring uncertainty quantification. We achieve this by extending conservative Q-learning [29], which does not require explicit uncertainty quantification, into the model-based setting. Our algorithm COMBO, summarized in Algorithm 1, alternates between a conservative policy evaluation step and a policy improvement step, which we outline below.
Conservative Policy Evaluation: Given a policy π, an offline dataset D, and a learned model of the MDP M̂, the goal in this step is to obtain a conservative estimate of Qπ . To achieve this, we penalize the Q-values evaluated on data drawn from a particular state-action distribution that is more likely to be out-of-support while pushing up the Q-values on state-action pairs that are trustworthy, which is implemented by repeating the following recursion:
Q̂k+1←argmin Q
β ( Es,a∼ρ(s,a)[Q(s,a)]−Es,a∼D[Q(s,a)] ) + 1
2 Es,a,s′∼df
[( Q(s,a)−B̂πQ̂k(s,a) )2] . (2)
Here, ρ(s,a) and df are sampling distributions that we can choose. Model-based algorithms allow ample flexibility for these choices while providing the ability to control the bias introduced by these choices. For ρ(s,a), we make the following choice: ρ(s,a) = dπ
M̂ (s)π(a|s), where dπ M̂ (s) is the
discounted marginal state distribution when executing π in the learned model M̂. Samples from dπ M̂ (s) can be obtained by rolling out π in M̂. Similarly, df is an f−interpolation between the offline dataset and synthetic rollouts from the model: dµf (s,a) := f d(s,a) + (1 − f) d µ
M̂ (s,a),
where f ∈ [0, 1] is the ratio of the datapoints drawn from the offline dataset as defined in Section 2 and µ(·|s) is the rollout distribution used with the model, which can be modeled as π or a uniform distribution. To avoid notation clutter, we also denote df := d µ f .
Under such choices of ρ and df , we push down (or conservatively estimate) Q-values on state-action tuples from model rollouts and push up Q-values on the real state-action pairs from the offline dataset. When updating Q-values with the Bellman backup, we use a mixture of both the model-generated data and the real data, similar to Dyna [57]. Note that in comparison to CQL and other model-free algorithms, COMBO learns the Q-function over a richer set of states beyond the states in the offline dataset. This is made possible by performing rollouts under the learned dynamics model, denoted by dµ M̂ (s,a). We will show in Section 4 that the Q function learned by repeating the recursion in Eq. 2 provides a lower bound on the true Q function, without the need for explicit uncertainty estimation. Furthermore, we will theoretically study the advantages of using synthetic data from the learned model, and characterize the impacts of model bias.
Policy Improvement Using a Conservative Critic: After learning a conservative critic Q̂π, we improve the policy as:
π′ ← argmax π
Es∼ρ,a∼π(·|s) [ Q̂π(s,a) ] (3)
where ρ(s) is the state marginal of ρ(s,a). When policies are parameterized with neural networks, we approximate the argmax with a few steps of gradient descent. In addition, entropy regularization can also be used to prevent the policy from becoming degenerate if required [17]. In Section 4.2, we show that the resulting policy is guaranteed to improve over the behavior policy.
Practical Implementation Details. Our practical implementation largely follows MOPO, with the key exception that we perform conservative policy evaluation as outlined in this section, rather than using uncertainty-based reward penalties. Following MOPO, we represent the probabilistic dynamics model using a neural network, with parameters θ, that produces a Gaussian distribution over the next state and reward: T̂θ(st+1, r|s,a) = N (µθ(st,at),Σθ(st,at)). The model is trained via maximum likelihood. For conservative policy evaluation (eq. 2) and policy improvement (eq. 3), we augment ρ with states sampled from the offline dataset, which shows more stable improvement in practice. It is relatively common in prior work on model-based offline RL to select various hyperparameters using online policy rollouts [67, 26, 3, 33]. However, we would like to avoid this with our method, since requiring online rollouts to tune hyperparameters contradicts the main aim of offline RL, which is to learn entirely from offline data. Therefore, we do not use online rollouts for tuning COMBO, and instead devise an automated rule for tuning important hyperparameters such as β and f in a fully offline manner. We search over a small discrete set of hyperparameters for each task, and use the value of the regularization term Es,a∼ρ(s,a)[Q(s,a)]−Es,a∼D[Q(s, s)] (shown in Eq. 2) to pick hyperparameters in an entirely offline fashion. We select the hyperparameter setting that achieves the lowest regularization objective, which indicates that the Q-values on unseen model-predicted state-action tuples are not overestimated. Additional details about the practical implementation and the hyperparameter selection rule are provided in Appendix B.1 and Appendix B.2 respectively.

heading: 4 Theoretical Analysis of COMBO

text: In this section, we theoretically analyze our method and show that it optimizes a lower-bound on the expected return of the learned policy. This lower bound is close to the actual policy performance (modulo sampling error) when the policy’s state-action marginal distribution is in support of the state-action marginal of the behavior policy and conservatively estimates the performance of a policy otherwise. By optimizing the policy against this lower bound, COMBO guarantees policy improvement beyond the behavior policy. Furthermore, we use these insights to discuss cases when COMBO is less conservative compared to model-free counterparts.

heading: 4.1 COMBO Optimizes a Lower Bound

text:
We first show that training the Q-function using Eq. 2 produces a Q-function such that the expected off-policy policy improvement objective [8] computed using this learned Q-function lower-bounds its actual value. We will reuse notation for df and d from Sections 2 and 3. Assuming that the Q-function is tabular, the Q-function found by approximate dynamic programming in iteration k, can be obtained by differentiating Eq. 2 with respect to Qk (see App. A for details):
Q̂k+1(s,a) = (B̂πQk)(s,a)− β ρ(s,a)− d(s,a) df (s,a) . (4)
Eq. 4 effectively applies a penalty that depends on the three distributions appearing in the COMBO critic training objective (Eq. 2), of which ρ and df are free variables that we choose in practice as discussed in Section 3. For a given iteration k of Eq. 4, we further define the expected penalty under ρ(s,a) as:
ν(ρ, f) := Es,a∼ρ(s,a) [ ρ(s,a)− d(s,a)
df (s,a)
] . (5)
Next, we will show that the Q-function learned by COMBO lower-bounds the actual Q-function under the initial state distribution µ0 and any policy π. We also show that the asymptotic Q-function learned by COMBO lower-bounds the actual Q-function of any policy π with high probability for a large enough β ≥ 0, which we include in Appendix A.2. LetM represent the empirical MDP which uses the empirical transition model based on raw data counts. The Bellman backups over the dataset distribution df in Eq. 2 that we analyze is an f−interpolation of the backup operator in the empirical MDP (denoted by BπM) and the backup operator under the learned model M̂ (denoted by B π M̂
). The empirical backup operator suffers from sampling error, but is unbiased in expectation, whereas the model backup operator induces bias but no sampling error. We assume that all of these backups enjoy concentration properties with concentration coefficient Cr,T,δ , dependent on the desired confidence value δ (details in Appendix A.2). This is a standard assumption in literature [31]. Now, we state our main results below. Proposition 4.1. For large enough β, we have Es∼µ0,a∼π(·|s)[Q̂π(s,a)] ≤ Es∼µ0,a∼π(·|s)[Qπ(s,a)], where µ0(s) is the initial state distribution. Furthermore, when ϵs is small, such as in the large
sample regime, or when the model bias ϵm is small, a small β is sufficient to guarantee this condition along with an appropriate choice of f .
The proof for Proposition 4.1 can be found in Appendix A.2. Finally, while Kumar et al. [29] also analyze how regularized value function training can provide lower bounds on the value function at each state in the dataset [29] (Proposition 3.1-3.2), our result shows that COMBO is less conservative in that it does not underestimate the value function at every state in the dataset like CQL (Remark 1) and might even overestimate these values. Instead COMBO penalizes Q-values at states generated via model rollouts from ρ(s,a). Note that in general, the required value of β may be quite large similar to prior works, which typically utilize a large constant β, which may be in the form of a penalty on a regularizer [36, 29] or as constants in theoretically optimal algorithms [23, 49]. While it is challenging to argue that that either COMBO or CQL attains the tightest possible lower-bound on return, in our final result of this section, we discuss a sufficient condition for the COMBO lower-bound to be tighter than CQL.
Proposition 4.2. Assuming previous notation, let ∆πCOMBO := Es,a∼dM(s),π(a|s) [ Q̂π(s,a) ] and
∆πCQL := Es,a∼dM(s),π(a|s) [ Q̂πCQL(s,a) ] denote the average values on the dataset under the Qfunctions learned by COMBO and CQL respectively. Then, ∆πCOMBO ≥ ∆πCQL, if:
Es,a∼ρ(s,a) [ π(a|s) πβ(a|s) ] − Es,a∼dM(s),π(a|s) [ π(a|s) πβ(a|s) ] ≤ 0. (∗)
Proposition 4.2 indicates that COMBO will be less conservative than CQL when the action probabilities under learned policy π(a|s) and the probabilities under the behavior policy πβ(a|s) are closer together on state-action tuples drawn from ρ(s,a) (i.e., sampled from the model using the policy π(a|s)), than they are on states from the dataset and actions from the policy, dM(s)π(a|s). COMBO’s objective (Eq. 2) only penalizes Q-values under ρ(s,a), which, in practice, are expected to primarily consist of out-of-distribution states generated from model rollouts, and does not penalize the Q-value at states drawn from dM(s). As a result, the expression (∗) is likely to be negative, making COMBO less conservative than CQL.

heading: 4.2 Safe Policy Improvement Guarantees

text:
Now that we have shown various aspects of the lower-bound on the Q-function induced by COMBO, we provide policy improvement guarantees for the COMBO algorithm. Formally, Proposition 4.3 discuss safe improvement guarantees over the behavior policy. building on prior work [46, 31, 29]. Proposition 4.3 (ζ-safe policy improvement). Let π̂out(a|s) be the policy obtained by COMBO. Then, if β is sufficiently large and ν(ρπ, f)− ν(ρβ , f) ≥ C for a positive constant C, the policy π̂out(a|s) is a ζ-safe policy improvement over πβ in the actual MDPM, i.e., J(π̂out,M) ≥ J(πβ ,M) − ζ, with probability at least 1− δ, where ζ is given by,
O ( γf
(1− γ)2 ) E s∼dπ̂outM [√ |A| |D(s)|DCQL(π̂out, πβ) ] ︸ ︷︷ ︸
:= (1)
+O ( γ(1− f) (1− γ)2 ) DTV(M,M̂)︸ ︷︷ ︸
:= (2)
−β C (1− γ)︸ ︷︷ ︸ := (3) .
The complete statement (with constants and terms that grow smaller than quadratic in the horizon) and proof for Proposition 4.3 is provided in Appendix A.4. DCQL denotes a notion of probabilistic distance between policies [29] which we discuss further in Appendix A.4. The expression for ζ in Proposition 4.3 consists of three terms: term (1) captures the decrease in the policy performance due to limited data, and decays as the size of D increases. The second term (2) captures the suboptimality induced by the bias in the learned model. Finally, as we show in Appendix A.4, the third term (3) comes from ν(ρπ, f)− ν(ρβ , f), which is equivalent to the improvement in policy performance as a result of running COMBO in the empirical and model MDPs. Since the learned model is trained on the dataset D with transitions generated from the behavior policy πβ , the marginal distribution ρβ(s,a) is expected to be closer to d(s,a) for πβ as compared to the counterpart for the learned policy, ρπ . Thus, the assumption that ν(ρπ, f)− ν(ρβ , f) is positive is reasonable, and in such cases, an appropriate (large) choice of β will make term (3) large enough to counteract terms (1) and (2) that reduce policy performance. We discuss this elaborately in Appendix A.4 (Remark 3).
Further note that in contrast to Proposition 3.6 in Kumar et al. [29], note that our result indicates the sampling error (term (1)) is reduced (multiplied by a fraction f ) when a near-accurate model is used to augment data for training the Q-function, and similarity, it can avoid the bias of model-based methods by relying more on the model-free component. This allows COMBO to attain the best-of-both model-free and model-based methods, via a suitable choice of the fraction f .
To summarize, through an appropriate choice of f , Proposition 4.3 guarantees safe improvement over the behavior policy without requiring access to an oracle uncertainty estimation algorithm.

heading: 5 Experiments

text:
In our experiments, we aim to answer the follow questions: (1) Can COMBO generalize better than previous offline model-free and model-based approaches in a setting that requires generalization to tasks that are different from what the behavior policy solves? (2) How does COMBO compare with prior work in tasks with high-dimensional image observations? (3) How does COMBO compare to prior offline model-free and model-based methods in standard offline RL benchmarks?
To answer those questions, we compare COMBO to several prior methods. In the domains with compact state spaces, we compare with recent model-free algorithms like BEAR [28], BRAC [63], and CQL [29]; as well as MOPO [67] and MOReL [26] which are two recent model-based algorithms. In addition, we also compare with an offline version of SAC [17] (denoted as SAC-off), and behavioral cloning (BC). In high-dimensional image-based domains, which we use to answer question (3), we compare to LOMPO [48], which is a latent space offline model-based RL method that handles image inputs, latent space MBPO (denoted LMBPO), similar to Janner et al. [20] which uses the model to generate additional synthetic data, the fully offline version of SLAC [32] (denoted SLAC-off), which only uses a variational model for state representation purposes, and CQL from image inputs. To our knowledge, CQL, MOPO, and LOMPO are representative of state-of-the-art model-free and model-based offline RL methods. Hence we choose them as comparisons to COMBO. To highlight the distinction between COMBO and a naïve combination of CQL and MBPO, we perform such a comparison in Table 8 in Appendix C. For more details of our experimental set-up, comparisons, and hyperparameters, see Appendix B.

heading: 5.1 Results on tasks that require generalization

text:
To answer question (1), we use two environments halfcheetah-jump and ant-angle constructed in Yu et al. [67], which requires the agent to solve a task that is different from what the behavior policy solved. In both environments, the offline dataset is
collected by policies trained with original reward functions of halfcheetah and ant, which reward the robots to run as fast as possible. The behavior policies are trained with SAC with 1M steps and we take the full replay buffer as the offline dataset. Following Yu et al. [67], we relabel rewards in the offline datasets to reward the halfcheetah to jump as high as possible and the ant to run to the top corner with a 30 degree angle as fast as possible. Following the same manner, we construct a third task sawyer-door-close based on the environment in Yu et al. [66], Rafailov et al. [48]. In this task, we collect the offline data with SAC policies trained with a sparse reward function that only gives a reward of 1 when the door is opened by the sawyer robot and 0 otherwise. The offline dataset is similar to the “medium-expert“ dataset in the D4RL benchmark since we mix equal amounts of data collected by a fully-trained SAC policy and a partially-trained SAC policy. We relabel the reward such that it is 1 when the door is closed and 0 otherwise. Therefore, in these datasets, the offline RL methods must generalize beyond behaviors in the offline data in order to learn the intended behaviors. We visualize the sawyer-door-close environment in the right image in Figure 3 in Appendix B.4.
We present the results on the three tasks in Table 1. COMBO significantly outperforms MOPO, MOReL and CQL, two representative model-based methods and one representative model-free methods respectively, in the halfcheetah-jump and sawyer-door-close tasks, and achieves an approximately 8%, 4% and 12% improvement over MOPO, MOReL and CQL respectively on the ant-angle task. These results validate that COMBO achieves better generalization results in practice by behaving less conservatively than prior model-free offline methods (compare to CQL, which doesn’t improve much), and does so more robustly than prior model-based offline methods (compare to MOReL and MOPO).

heading: 5.1.1 Empirical analysis on uncertainty estimation in offline model-based RL

text:
To further understand why COMBO outperforms prior model-based methods in tasks that require generalization, we argue that one of the main reasons could be that uncertainty estimation is hard in these tasks where the agent is required to go further away from the data distribution. To test this intuition, we perform empirical evaluations to study whether uncertainty quantification with deep neural networks, especially in the setting of dynamics model learning, is challenging and could cause problems with uncertainty-based model-based offline RL methods such as
MOReL [26] and MOPO [67]. In our evaluations, we consider maximum learned variance over the ensemble (denoted as Max Var) maxi=1,...,N ∥Σiθ(s,a)∥F (used in MOPO). We consider two tasks halfcheetah-jump and ant-angle. We normalize both the model error and the uncertainty estimates to be within scale [0, 1] and performs linear regression that learns the mapping between the uncertainty estimates and the true model error. As shown in Figure 2, on both tasks, Max Var is unable to accurately predict the true model error, suggesting that uncertainty estimation used by offline model-based methods is not accurate and might be the major factor that results in its poor performance. Meanwhile, COMBO circumvents challenging uncertainty quantification problem and achieves better performances on those tasks, indicating the effectiveness and the robustness of the method.

heading: 5.2 Results on image-based tasks

text:
To answer question (2), we evaluate COMBO on two image-based environments: the standard walker (walker-walk) task from the the DeepMind Control suite [61] and a visual door opening environment with a Sawyer robotic arm (sawyer-door) as used in Section 5.1.
For the walker task we construct 4 datasets: medium-replay (M-R), medium (M), medium-expert (ME), and expert, similar to Fu et al. [12], each consisting of 200 trajectories. For sawyer-door task we use only the medium-expert and the expert datasets, due to the sparse reward – the agent is rewarded only when it successfully opens the door. Both environments are visulized in Figure 3 in Appendix B.4. To extend
COMBO to the image-based setting, we follow Rafailov et al. [48] and train a recurrent variational model using the offline data and use train COMBO in the latent space of this model. We present
results in Table 2. On the walker-walk task, COMBO performs in line with LOMPO and previous methods. On the more challenging Sawyer task, COMBO matches LOMPO and achieves 100% success rate on the medium-expert dataset, and substantially outperforms all other methods on the narrow expert dataset, achieving an average success rate of 96.7%, when all other model-based and model-free methods fail.

heading: 5.3 Results on the D4RL tasks

text: Finally, to answer the question (3), we evaluate COMBO on the OpenAI Gym [6] domains in the D4RL benchmark [12], which contains three environments (halfcheetah, hopper, and walker2d) and four dataset types (random, medium, medium-replay, and medium-expert). We include the results in Table 3. The numbers of BC, SAC-off, BEAR, BRAC-P and BRAC-v are taken from the D4RL paper, while the results for MOPO, MOReL and CQL are based on their respective papers [67, 29]. COMBO achieves the best performance in 9 out of 12 settings and comparable result in 1 out of the remaining 3 settings (hopper medium-replay). As noted by Yu et al. [67] and Rafailov et al. [48], model-based offline methods are generally more performant on datasets that are collected by a wide range of policies and have diverse state-action distributions (random, medium-replay datasets) while model-free approaches do better on datasets with narrow distributions (medium, medium-expert datasets). However, in these results, COMBO generally performs well across dataset types compared to existing model-free and model-based approaches, suggesting that COMBO is robust to different dataset types.

heading: 6 Related Work

text:
Offline RL [10, 50, 30, 34] is the task of learning policies from a static dataset of past interactions with the environment. It has found applications in domains including robotic manipulation [25, 38, 48, 54], NLP [21, 22] and healthcare [52, 62]. Similar to interactive RL, both model-free and model-based algorithms have been studied for offline RL, with explicit or implicit regularization of the learning algorithm playing a major role.
Model-free offline RL. Prior model-free offline RL algorithms have been designed to regularize the learned policy to be “close“ to the behavioral policy either implicitly via regularized variants of importance sampling based algorithms [47, 58, 35, 59, 41], offline actor-critic methods [53, 45, 27, 16, 64], applying uncertainty quantification to the predictions of the Q-values [2, 28, 63, 34], and learning conservative Q-values [29, 55] or explicitly measured by direct state or action constraints [14, 36], KL divergence [21, 63, 69], Wasserstein distance, MMD [28] and auxiliary imitation loss [13]. Different from these works, COMBO uses both the offline dataset as well as model-generated data.
Model-based offline RL. Model-based offline RL methods [11, 9, 24, 26, 67, 39, 3, 60, 48, 33, 68] provide an alternative approach to policy learning that involves the learning of a dynamics model using techniques from supervised learning and generative modeling. Such methods however rely either on uncertainty quantification of the learned dynamics model which can be difficult for deep network models [44], or on directly constraining the policy towards the behavioral policy similar to model-free algorithms [39]. In contrast, COMBO conservatively estimates the value function by penalizing it in out-of-support states generated through model rollouts. This allows COMBO to
retain all benefits of model-based algorithms such as broad generalization, without the constraints of explicit policy regularization or uncertainty quantification.

heading: 7 Conclusion

text:
In the paper, we present conservative offline model-based policy optimization (COMBO), a modelbased offline RL algorithm that penalizes the Q-values evaluated on out-of-support state-action pairs. In particular, COMBO removes the need of uncertainty quantification as widely used in previous model-based offline RL works [26, 67], which can be challenging and unreliable with deep neural networks [44]. Theoretically, we show that COMBO achieves less conservative Q values compared to prior model-free offline RL methods [29] and guarantees a safe policy improvement. In our empirical study, COMBO achieves the best generalization performances in 3 tasks that require adaptation to unseen behaviors. Moreover, COMBO is able scale to vision-based tasks and outperforms or obtain comparable results in vision-based locomotion and robotic manipulation tasks. Finlly, on standard D4RL benchmark, COMBO generally performs well across dataset types compared to prior methods Despite the advantages of COMBO, there are few challenges left such as the lack of an offline hyperparameter selection scheme that can yield a uniform hyperparameter across different datasets and an automatically selected f conditioned on the model error. We leave them for future work.
Acknowledgments and Disclosure of Funding We thank members of RAIL and IRIS for their support and feedback. This work was supported in part by ONR grants N00014-20-1-2675 and N00014-21-1-2685 as well as Intel Corporation. AK and SL are supported by the DARPA Assured Autonomy program. AR was supported by the J.P. Morgan PhD Fellowship in AI.

title: COMBO: Conservative Offline Model-Based Policy Optimization


summaries: This paper proposed a stochastic differential equation model based on the local elasticity assumption of features of neural networks.  The local elasticity phenomenon says samples have greater influence to other samples in the same class than those in different classes. Starting from this assumption, a system of linear SDEs is proposed for features corresponding to each data point, with the coefficients encoding the difference of inter and intra classes impact. Two types of features are studied---the isotropic feature learning model, and the logics-as-features model. For the two models, dynamics of the class means are derived as a system of ODEs and the solutions are found, depending on the coefficients about local elasticity. From the solutions, the separability of classes is theoretically connected with local elasticity. Specifically, the classes are separable if the local elasticity effect is nonzero, while not separable otherwise. The transition between the two cases is sharp.   Numerical experiments are conducted to justify the effectiveness of the proposed model to characterize the actual dynamics of features. Experiments are done on deep convolutional networks on the GeoMNIST and CIFAR10 dataset. Firstly, the coefficients in the SDEs (and the ODEs for the class means) are estimated using real feature dynamics given by the training process (I understand this step as a kind of "fitting"). Then, the ODEs with estimated coefficients are simulated, and the results are compared with the real feature dynamics. It shows that in many cases (but not all) the trajectories produced by the SDEs recover the real trajectories given by the training.

summaries: The paper proposes a proxy tractable dynamics to study how the latent features of neural networks evolve inter- and intra-class during training. The core of this dynamics is a linear ODE, which is shown to exhibit inter-class separability once a “locally elastic” condition is met. The paper then examines this dynamics for two choices of the H matrix (which is a parameter of the dynamics), and concludes with experiments to demonstrate the relevance of the proposed theory with the actual simulation.   ===================== After rebuttal: Thanks for the detailed reply!  Overall while I am pleased that the authors made a lot of efforts to clarify, I'm not quite convinced. Some notes: - I agree that time-dependence of E(t) is a good point; my concern about linearization (which I believe is also the concern of some other reviewers) is actually about why at each time t, F(X(t),t), in the authors' response, is modeled by a linear relationship with X(t). Why ignoring higher-order terms? Are they much smaller than the linear term? I do not see why the higher-order terms can be absorbed into the noise term. - I do not agree with the argument that $c_1\neq c_2$ assumption can be alleviated by taking the results "in probability". Clearly the derivation in the paper requires taking a law-of-large-numbers average.  I agree (and well understood) that the paper does not aim to mimic exactly the dynamics; the aim is to provide a reasonable proxy dynamics that can explain interesting properties. However I feel the paper needs a careful justification to convince readers that this proxy is not oversimplifying.  Given that the paper may attract interests, I'm willing to raise the score, but not fully convinced of the quality to recommend acceptance. If a reviewer is willing to champion the paper, accepting the paper would not be a bad outcome to me.

summaries: This paper builds off a recently proposed phenomenon of deep learning models, local elasticity, which, says the that impact in feature space of a gradient update from an input is in general larger for feature of data in the same class as the input rather than a different class. In this work the authors propose a set of SDE's modeling SGD that captures the inter and intra class effect of back propagation. They find that there is a sharp phase transition in the dynamics governed by these SDEs when they do or don't display local elasticity. In particular, when the SDEs are locally elastic then the features of the training data are guaranteed to become linearity separable, while this is not the case otherwise. They demonstrate this theoretical analysis empirically with CIFAR-10 and connect there results to another recently proposed phenomenon, neural collapse.

summaries: The paper provides a simple phenomenological model of neural networks for K-class classification. Each sample is assigned a feature modeling NN features, and these features are evolved via SDEs in which the drift term of a sample is a linear transformation of the average features of each class. Separation, which means the case in which samples from different classes can be separated by hyperplanes, holds w.h.p. when the intra-class drift coefficients larger than the inter-class coefficients in a certain sense, and it does not hold otherwise. Two specific choices of features are considered: isotropic features and logit features. The latter are shown to be better than the former at modeling the feature evolution for NN on CIFAR and a dataset of their own creation.

abstractText: Understanding the training dynamics of deep learning models is perhaps a necessary step toward demystifying the effectiveness of these models. In particular, how do data from different classes gradually become separable in their feature spaces when training neural networks using stochastic gradient descent? In this study, we model the evolution of features during deep learning training using a set of stochastic differential equations (SDEs) that each corresponds to a training sample. As a crucial ingredient in our modeling strategy, each SDE contains a drift term that reflects the impact of backpropagation at an input on the features of all samples. Our main finding uncovers a sharp phase transition phenomenon regarding the intra-class impact: if the SDEs are locally elastic [19] in the sense that the impact is more significant on samples from the same class as the input, the features of the training data become linearly separable, meaning vanishing training loss; otherwise, the features are not separable, regardless of how long the training time is. Moreover, in the presence of local elasticity, an analysis of our SDEs shows that the emergence of a simple geometric structure called the neural collapse of the features. Taken together, our results shed light on the decisive role of local elasticity in the training dynamics of neural networks. We corroborate our theoretical analysis with experiments on a synthesized dataset of geometric shapes and CIFAR-10.

authors:
Z. Allen-Zhu
Y. Li
Z. Song

title: A convergence theory for deep learning via over-parameterization

venue: International Conference on Machine Learning, pages 242–252

authors:
S. Arora
N. Cohen
N. Golowich
W. Hu

title: A convergence analysis of gradient descent for deep linear neural networks

venue: International Conference on Learning Representations

authors:
S. Arora
N. Cohen
E. Hazan

title: On the optimization of deep networks: Implicit acceleration by overparameterization

venue: International Conference on Machine Learning, pages 244–253

authors:
S. Arora
S. Du
W. Hu
Z. Li
R. Wang

title: Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks

venue: International Conference on Machine Learning, pages 322–332

authors:
P. Chaudhari
S. Soatto

title: Stochastic gradient descent performs variational inference

venue: converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1–10. IEEE

authors:
S. Chen
H. He
W. Su

title: Label-aware neural tangent kernel: Toward better generalization and local elasticity

venue: Advances in Neural Information Processing Systems, volume 33, pages 15847–15858

authors:
Z. Chen
Y. Cao
Q. Gu
T. Zhang

title: A generalized neural tangent kernel analysis for two-layer neural networks

venue: Advances in Neural Information Processing Systems, 33

authors:
J. Cohen
S. Kaur
Y. Li
J.Z. Kolter
A. Talwalkar

title: Gradient descent on neural networks typically occurs at the edge of stability

authors:
A. Damianou
N.D. Lawrence

title: Deep Gaussian processes

venue: Artificial intelligence and statistics, pages 207–215

authors:
Z. Deng
H. He
W. Su

title: Toward better generalization bounds with locally elastic stability

venue: Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2590–2600

authors:
J. Devlin
M. Chang
K. Lee
K. Toutanova

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

venue: Conference of the North American Chapter of the Association for Computational Linguistics, pages 4171–4186. Association for Computational Linguistics

authors:
S. Du
W. Hu

title: Width provably matters in optimization for deep linear neural networks

venue: International Conference on Machine Learning, pages 1655–1664

authors:
S. Du
J. Lee
H. Li
L. Wang
X. Zhai

title: Gradient descent finds global minima of deep neural networks

venue: International Conference on Machine Learning, pages 1675–1685

authors:
S.S. Du
K. Hou
R.R. Salakhutdinov
B. Poczos
R. Wang
K. Xu

title: Graph neural tangent kernel: Fusing graph neural networks with graph kernels

venue: Advances in Neural Information Processing Systems, volume 32

authors:
S.S. Du
X. Zhai
B. Póczos
A. Singh

title: Gradient descent provably optimizes over-parameterized neural networks

venue: 7th International Conference on Learning Representations. OpenReview.net

authors:
C. Fang
H. He
Q. Long
W.J. Su

title: Exploring deep neural networks via layer-peeled model: Minority collapse in imbalanced training

venue: Proceedings of the National Academy of Sciences

authors:
S. Goldt
M. Advani
A.M. Saxe
F. Krzakala
L. Zdeborová

title: Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup

authors:
T. Hazan
T. Jaakkola

title: Steps toward deep kernel methods from infinite neural networks

venue: arXiv preprint arXiv:1508.05133

authors:
H. He
W. Su

title: The local elasticity of neural networks

authors:
K. He
X. Zhang
S. Ren
J. Sun

title: Deep residual learning for image recognition

venue: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 770–778

authors:
S. Ioffe
C. Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 448–456

authors:
A. Jacot
F. Gabriel
C. Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

venue: Advances in Neural Information Processing Systems, volume 31

authors:
S. Jastrzebski
Z. Kenton
D. Arpit
N. Ballas
A. Fischer
Y. Bengio
A. Storkey

title: Three factors influencing minima in SGD

venue: arXiv preprint arXiv:1711.04623

authors:
Z. Ji
M. Telgarsky

title: Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks

title: Deep learning without poor local minima

venue: Advances in Neural Information Processing Systems, volume 29

authors:
N.S. Keskar
D. Mudigere
J. Nocedal
M. Smelyanskiy
P.T.P. Tang

title: On large-batch training for deep learning: Generalization gap and sharp minima

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization

venue: 3rd International Conference on Learning Representations

title: Learning multiple layers of features from tiny images

authors:
A. Krizhevsky
I. Sutskever
G.E. Hinton

title: Imagenet classification with deep convolutional neural networks

venue: Advances in Neural Information Processing Systems, 25:1097–1105

authors:
A.K. Lampinen
S. Ganguli

title: An analytic theory of generalization dynamics and transfer learning in deep linear networks

authors:
Y. LeCun
Y. Bengio
G. Hinton

title: Deep learning

authors:
J. Lee
Y. Bahri
R. Novak
S.S. Schoenholz
J. Pennington
J. Sohl-Dickstein

title: Deep neural networks as gaussian processes

venue: arXiv preprint arXiv:1711.00165

authors:
Q. Li
C. Tai
E. Weinan

title: Stochastic modified equations and adaptive stochastic gradient algorithms

venue: International Conference on Machine Learning, pages 2101–2110

authors:
Y. Li
Y. Liang

title: Learning overparameterized neural networks via stochastic gradient descent on structured data

authors:
Z. Li
S. Malladi
S. Arora

title: On the validity of modeling sgd with stochastic differential equations (sdes)

venue: arXiv preprint arXiv:2102.12470

authors:
S. Mandt
M.D. Hoffman
D.M. Blei

title: Stochastic gradient descent as approximate bayesian inference

venue: The Journal of Machine Learning Research, 18(1):4873–4907

authors:
A.W. Marshall
I. Olkin
B.C. Arnold

title: Inequalities: Theory of Majorization and its Applications

venue: volume 143. Springer, second edition

authors:
V. Papyan
X.Y. Han
D.L. Donoho

title: Prevalence of neural collapse during the terminal phase of deep learning training

venue: Proceedings of the National Academy of Sciences, 117(40):24652–24663

authors:
D. Saad
S.A. Solla

title: Dynamics of on-line gradient descent learning for multilayer neural networks

venue: Advances in Neural Information Processing Systems, pages 302–308

authors:
A.M. Saxe
J.L. McClelland
S. Ganguli

title: Exact solutions to the nonlinear dynamics of learning in deep linear neural networks

title: A mathematical theory of semantic development in deep neural networks

venue: Proceedings of the National Academy of Sciences, 116(23):11537–11546

authors:
B. Shi
W.J. Su
M.I. Jordan

title: On learning rates and Schrödinger operators

venue: arXiv preprint arXiv:2004.06977

authors:
D. Silver
A. Huang
C.J. Maddison
A. Guez
L. Sifre

title: G

venue: Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Bature, 529(7587):484–489

authors:
S.L. Smith
P.-J. Kindermans
Q.V. Le

title: Don’t decay the learning rate

venue: increase the batch size. In International Conference on Learning Representations

authors:
D. Soudry
E. Hoffer
M.S. Nacson
S. Gunasekar
N. Srebro

title: The implicit bias of gradient descent on separable data

venue: The Journal of Machine Learning Research, 19(1):2822–2878

authors:
S. Särkkä
A. Solin

title: Applied Stochastic Differential Equations

venue: Institute of Mathematical Statistics Textbooks. Cambridge University Press

title: Scaling limits of wide neural networks with weight sharing: Gaussian process behavior

venue: gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760

authors:
T. Young
D. Hazarika
S. Poria
E. Cambria

title: Recent trends in deep learning based natural language processing

venue: IEEE Computational Intelligence Magazin, 13(3):55–75

authors:
C. Zhang
S. Bengio
M. Hardt
B. Recht
O. Vinyals

title: Understanding deep learning requires rethinking generalization

heading: 1 Introduction

text:
Deep learning models have achieved significant empirical success over the past decade across a wide spectrum of domains spanning computer vision, natural language processing, and reinforcement learning [31, 43, 48]. Despite these remarkable achievements at the empirical level, there is still much to learn about deep neural networks, as evidenced by the fact that almost all important advances concerning architecture design and optimization for deep learning are based on heuristics, without much input from a theoretical perspective [20, 11, 21, 27].
An important step toward opening these black-box models and unveiling their formidable details is to quantitatively understand the impact of backpropagation in deep learning training. While there has been a continued effort to demystify how simple optimization methods give rise to impressive generalization performance, for example, [49, 26, 4], this is by no means an easy problem, perhaps because of the daunting nonconvex nature of neural networks. Accordingly, for near-term purposes, a more practical approach is to take a phenomenological viewpoint by relating simple empirical patterns to the effectiveness of deep learning models.
In this spirit, we are interested in how data from different classes gradually become separable in their feature space by repetitively calling backpropagation. From a phenomenological viewpoint, this question can be addressed by first analyzing the impact of a single update using a stochastic gradient on the performance of the neural networks. More precisely, imagine that the gradient is evaluated in an image of a cat, how does the hidden representation of another image—say, an image of another
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
X̄ k (T )−X̄ l(T ) for a large T , where k, l are chosen heuristically. The dashed lines in (b) and (d) are
simulated paths from Equation (9) using estimated α̂(t) and β̂(t). More details are given in Section 4 and Appendix D.2.
cat or an image of a plane—evolves because of the backpropagation? Recent studies answer this question by introducing a phenomenon called local elasticity, which, roughly speaking, means that the impact is generally larger on a similar sample (an image of another cat) than on a dissimilar sample (an image of a plane) [19].
Motivated by the phenomenon of local elasticity, we propose a model that captures the interaction between different training samples during deep learning training using a set of stochastic differential equations (SDEs) that reflect local elasticity in neural networks. Characterizing the intra-class and inter-class effects is an essential component of our modeling strategy, each of which contains a drift term that imitates the impact of backpropagation on specific training data of all samples.
Our main finding uncovers a sharp phase transition phenomenon regarding the intra-class and interclass impact; if the SDEs are locally elastic in the sense that the impact is more significant on samples from the same class as the input (the intra-class effect is strictly greater than the inter-class effect), the features of the training data are guaranteed to be linearly separable, meaning vanishing training loss; otherwise, the features are not separable, no matter how long the training time is. This result provides convincing theoretical evidence for the presence of local elasticity in deep learning [19]. Our model is also quite accurate in simulating the feature dynamics of deep learning. As shown in Figure 1 and detailed in Section 4, the dynamics of the simulated logits are reasonably close to the real dynamics of deep learning on both synthetic and real datasets, indicating a well-suited model for theoretical and practical purposes. Moreover, in the presence of local elasticity, our SDEs also predict the emergence of a simple geometric structure called neural collapse of features [38].
Taken together, our results shed light on the decisive role of local elasticity in the training dynamics of neural networks. We corroborate our theoretical analysis with experiments on a synthetic dataset of geometric shapes, as well as on CIFAR-10. The experimental evidence consistently supported our model, which provides new insights into the dynamics of deep learning training.

heading: 1.1 Related Work

text:
Dynamics in Deep Neural Nets. Many properties of linear deep neural nets are relatively well understood, such as the loss landscapes [25], trajectory-based convergence [2, 12], and implicit acceleration [3]. Exact solutions of the training dynamics can be obtained in certain initialization schemes [40, 41, 30]. In the presence of non-linearity, various assumptions are generally made. [39, 17] studied the dynamics of shallow neural nets with non-linearity and the neural tangent kernel (NTK) literature [22, 4, 14] linearizes the network function of an infinitely wide neural net at initialization, which is similar to that of the deep Gaussian process literature [9, 18, 32] (a treatise comparing and contrasting them can be found in [47]). Although as approximations, NTKs are generally used when studying optimization trajectories of neural nets such as [34], which also appears implicitly in many works when studying the optimization trajectories of neural net training [45, 1, 13, 24, 7].
SGD as SDEs in Neural Nets. The study of dynamics or trajectories of weights in deep neural nets via SDEs relies on the more precise characterization of stochasticity [33, 36]. Built on top of
this formalism, [44] studied the trade-off between batch size and learning rate, [23], analyzed factors influencing the quality of local minima, [5] studied the behavior of the SGD near local minima, and [42] studied the effect of learning rates. Although SGD-SDE approximation requires an infinitesimal learning rate, [35] verified that the SDE approximation can be meaningful in practical settings and obtained necessary conditions for the validation of such approximation.
Local Elasticity and Phenomenological Models. Local elasticity is proposed in [19] as a phenomenological approach to reasoning the behaviors of neural networks. This phenomenon has inspired several works on generalization bounds [10] and an improvement on the NTK [6].

heading: 2 Binary Separation via LE-SDE

heading: 2.1 Setup, Notations and Assumptions

text:
Throughout the paper, we work with the following setup and assumptions. For ease of reading, vectors and matrices are written in boldface and we denote by [n] the set {1, . . . , n}. When there is no ambiguity, we will write both X(t) and Xt for a continuous-time (possibly stochastic) process.
Classification Problem. Consider a K-class classification problem with K ≥ 2, with each class having n training examples. We denote by zki ∈ Rd the i-th sample of the k-th class, and yki ∈ [K] its label, where i ∈ [n], k ∈ [K]. A neural net is a function f : Rd → RK that maps the samples to logits ( pre-activation of the softmax).
Feature Vectors. We denote byXki (m) ∈ Rp a p-dimensional feature of the i-th sample in the k-th class learned by the neural net at iteration m. For example, it can be the logits or the output of the second-to-last layer. Assume that the initial valuesXki (0) are i.i.d. samples from some distribution for each i ∈ [n], k ∈ [K]. We use i, j ∈ [n] as indices for an individual sample, k, l ∈ [K] for classes, and capital letters J and L to indicate random samples from Unif([n]) and Unif([K]), respectively.
Training Dynamics. We model the training dynamics in neural nets under SGD with an emphasis on local elasticity. At the m-th iteration, the Jm-th sample is sampled from the Lm-th class, where Jm ∼ Unif([n]) and Lm ∼ Unif([K]). Training on zJmIm affects the features of another data sample zki in the form of
Xki (m)−Xki (m− 1) = h · Ek,Lm(m)XLmJm (m− 1) + √ hζki (m− 1) (1)
where i ∈ [n], k ∈ [K], h is the step size, and ζki (m) is the noise term that is modeled as Gaussian noise. The scalar Ek,Lm measures the strength of local elasticity that z Lm Jm
exerts on zki at iteration m. We assumeXki (0), ζ k i (m) are jointly independent.
Local Elasticity. Clearly, by writing Ek,l(m), we assume that this effect depends only on the class l, k, and time m. We write the effect matrix as E(m) = (Ek,l(m)) K k,l=1. For ease of exposition, we assume E only consists of two values α(m) and β(m), with α(m) representing the intra-class effect and β(m) the inter-class effect. To this end, we assume the effective training assumption, that is, as training progresses, the features become more discriminative: features from the same class are more similar, whereas those from different classes are more distinct, as measured by some similarity measure in the feature space. We also assume that the LE effect is “proportional” to the feature XLmJm (m) itself. We generalize this point in Section 3 by introducing a transformation matrixH on the features.

heading: 2.2 Binary LE-SDE

text:
Our construction of equation (1) emphasizes the effect of intra- and inter-class effects on the dynamics of features, and thus differs from the usual weight dynamics that is common in the literature. Before deriving the general form of our locally elastic SDE (LE-SDE), we shall familiarize the reader with our model by demonstrating this in the case of binary classification (K = 2) with a one-dimensional features (p = 1) — the output of the model to be fed into the softmax function, also called the logit.
Let the intra-class effect be E11 = E22 = α, and the inter-class effect is E12 = E21 = β, both of which are time-independent. Expanding equation (1), for 1 ≤ i, j ≤ n and m ≥ 0, when we train the
model on the Jm-th training example from the Lm-th class, we have{ X1i (m) = X 1 i (m− 1) + h · αXLmJm (m− 1) + √ h · ζLmi (m− 1),
X2j (m) = X 2 j (m− 1) + h · βXLmJm (m− 1) + √ h · ζLmj (m− 1).
In the limit of h → 0, we can show that Xki (m) approximates some continuous-time stochastic processes Xki (t) (under the identification of t = mh) governed by the set of stochastic differential equations as follows:
dXki (t) =
( α
2 X̄k(t) +
β 2 X̄3−k(t)
) dt+ σ dW ki (t), t ≥ 0, k ∈ [K], i ∈ [n] (2)
where X̄k(t) := ( Xk1 (t) + · · ·+Xkn(t) ) /n, and W ki are independent standard Wiener processes. The detailed derivation is given in Appendix A.
Now averaging over i for each k in equation (2), we obtain the following set of two ordinary differential equations (ODEs) governing the per-class means that X̄k(t) for k = 1, 2:
dX̄k(t) =
( α
2 X̄k(t) +
β 2 X̄3−k(t)
) dt+ σ d
W ki (t) + · · ·+W ki (t) n .
Taking the limit of n → ∞, we observe that σdW 1(t)+···+Wn(t)
n ⇒ 0. Thus, the above display converges weakly to the following ODE:
dX̄k(t)
dt = α 2 X̄k(t) + β 2 X̄3−k(t). (3)
With the initial conditions EXki (0) = ck, the solution to the above ODE is
X̄1(t) = c1 − c2
2 e α−β 2 t + c1 + c2 2 e α+β 2 t, X̄2(t) = −c1 − c2 2 e α−β 2 t + c1 + c2 2 e α+β 2 t.
In the finite-sample setting, we may replace X̄1, X̄2 in the SDE (2) by their deterministic solutions and obtain {
X1i (t) = c1−c2 2 e α−β 2 t + c1+c22 e α+β 2 t − c1 +X1i (0) + σW 1i (t), X2j (t) = − c1−c22 e α−β 2 t + c1+c22 e α+β 2 t − c2 +X2j (0) + σW 2j (t).
We are now ready to derive the condition under which these 2n feature vectors become asymptotically separable, that is, miniX1i (t) > maxj X 2 j (t) or maxiX 1 i (t) < minj X 2 j (t) as t→∞.
Theorem 2.1 (Separation in Binary Classification). Given the feature vectors X1i (t), X2j (t) for i, j ∈ [n], as t→∞ and large n,
1. if α > β, they are asymptotically separable with probability tending to one,
2. if α ≤ β, they are asymptotically separable with probability tending to zero. This result indicates a sharp phase transition when α is just above β, that is, in the regime of local elasticity. As long as the intra-class effect is slightly greater than the inter-class effect, separation is guaranteed. This simple model already captures local elasticity and reveals the important role it plays in the perfect separation of training samples. We can generalize this model to more realistic settings: when there are multiple classes, when features are high-dimensional, and when the LE matrix E is time-dependent. In the next section, we discuss each of these three generalizations in more depth.

heading: 3 General LE-SDE Model

text:
Now, we consider the general case where K ≥ 2 and the feature vectors are p-dimensional with p ≥ K. Inquisitive readers may have already noticed that Theorem 2.1 only asserts the emergence of the separation of features, while being inconclusive to their relative orders at separation, that is, which class converges to where? This drawback is intrinsic to the toy model as neither intra-class nor inter-class effect identifies different classes. In this section, we introduce the general LE-SDE model that alleviates this difficulty with the help of an extra block matrixH with the (i, j)-th blockHi,j models how features in the j-th class affect those in the i-th class, which also partially defines how
classes are separated in higher dimensions. In the local elasticity formalism,Hi,j can be viewed as inducing a metric on the feature space under which local elasticity manifests.
As hinted before, in the case of multiple-class features in higher dimensions, we want to guarantee a stronger separation: to know which class converges to where, thus incorporating supervision from label information. For example, when the features are logits (outputs of the neural nets) and the model is trained under the softmax cross-entropy loss, previous work suggests they separate according to specific geometric structures [38]. To this end, we need to adjust the raw feature vectorsXki with a proper transformation that incorporates the label information into the dynamics. This motivates the following modification of the dynamics (1) by adding an extra transformationHk,Lm ∈ Rp×p to the features. For k ∈ [K], i ∈ [n], and at iteration m, we have the following:
Xki (m) = X k i (m− 1) + h · Ek,Lm(m)Hk,Lm(m)XLmJm (m− 1) + √ hζki (m− 1). (4)
TheHk,Lm(m) term models the LE effect as proportional to a linear “transformation” of the features. The dynamics in Equation (1) are special cases whenHk,l(m) ≡ Ip for all k, l ∈ [K]. By specifying a properH , we can overcome the limitation in our toy example of not knowing which class converges to where. We specify interesting choices ofH in Section 3.2.
A further step of abstraction is to write X̃ k (m) instead ofXki (m), to indicate one generic sample from the distribution Dk(m) of all the features of class k at iteration m. As in Section 2.2, we can derive the continuous dynamics of equation (4) in the limit of h → 0 in the same way as equation (2). Similar to writing X̃ = (X̃
k )Kk=1 ∈ RKp for the concatenation of per-class features,
X̄ = (X̄ k )Kk=1 ∈ RKp is the concatenation of per-class mean features. Our model (4) approximates the following SDE with identification t = mh as h→ 0. We term this model the LE-SDE:
dX̃t = M tX̄t dt+ Σ 1 2 t dW t, (5)
whereW t is the standard Wiener process in RKp, Σt is the covariance matrix, andM t ∈ RKp×Kp is aK×K block matrix, with each block of size p×p. The (k, l)th block ofM t isEk,l(t)Hk,l(t)/K when l 6= k, and El,l(t)H l,l(t)/K when l = k. The rationale for dividing K is that we assume that the data are balanced; therefore, each of the K possible classes has an equal chance of being sampled, as proved in Equation (2), where K = 2. In Appendix E, we discuss how we can generalize this to model SGD with mini-batches, imbalanced data, and label corruptions.
Taking expectation with respect to the randomness arising from sampling X̃t from its distribution, the per-class mean X̄t satisfies the following system, which we term the LE-ODE:
X̄ ′ t = M tX̄t. (6)
Under the assumptions in Section 2.1, we define γ(t) = min {α(t)− β(t), α(t) + (K − 1)β(t)} , A(t) = ´ t 0 α(τ) dτ , B(t) = ´ t 0 β(τ) dτ , and Γ(t) = min {A(t)−B(t), A(t) + (K − 1)B(t)}.

heading: 3.1 The Separation Theorem

text:
Similar to the discussions in Theorem 2.1, the LE-SDE allows us to derive the separability result for a general K and p ≥ K. We say the feature vectors { (Xki )i∈[n] } k∈[K] are separable if for any two classes k 6= l, there exists a hyperplane in Rp that linearly separates the features of the two classes. To characterize the separation as in Theorem 2.1, we need conditions on α(t), β(t) as therein. Intuitively, when γ(t) ≤ 0, the classes cannot be separated, even in a pairwise manner. Therefore, we focus on a more interesting case when γ(t) > 0. We now state the following characterization theorem of separability for general LE-SDE dynamics:
Theorem 3.1 (Separation of LE-SDE). Under our working assumptions in Section 2.1, and in the case of local elasticity (i.e., γ(t) > 0), assume H = (Hij)ij is positive semi-definite (PSD) with positive diagonal entries. As t→∞, we have1:
1. if γ(t) = ω (1/t), the features are separable with probability tending to 1;
1Here, γ(t) = ω (1/t) stands for γ(t) 1/t as t→∞. For example, 1/t0.5 = ω (1/t) and (t ln t)−1 = o (1/t) as t→∞.
2. if γ(t) = o (1/t), and the number of per-class-feature n tending to∞ at an arbitrarily slow rate, the features are asymptotically pairwise separable with probability 0.
This theorem sheds light on the crucial impact of the local elasticity effect for separation in a general case. The proof to Theorem 3.1 as well as discussions on the empirically best ways of choosing the universal direction ν (i.e., a direction that does not depend on the class index) are detailed in Appendix C.2.

heading: 3.2 Two Specific Models

text: We next discuss two specific choices of theH matrix that allows us to analyze X̃ precisely.

heading: 3.2.1 Isotropic Feature Learning Model

text:
As a straightforward extension to Section 2.2, we can simply choose H lk = Ip to be the identity matrix. This choice ofH is PSD, and thus, we can apply Theorem 3.1 to obtain the conditions for asymptotic separation. In this case, the solution X̄(t) to the LE-ODE can be computed analytically as given in the following proposition.
Proposition 3.2 (I-model). LetHk,l = Ip, then the solution to the LE-ODE (6) is given by
X̄(t) = ce 1 KA(t)− 1 KB(t) + (1K ⊗ c0) e 1 KA(t)+ K−1 K B(t), (7)
where c = (ck)Kk=1 ∈ RKp and c0 ∈ Rp are constants with ∑K k=1 ck = 0 ∈ Rp and X̄(0) = c+c0.
The derivation of Equation (7) is deferred to Appendix C.2.1. From equation (7), we can easily reconstruct Theorem 3.1 in this special case. The difference between a feature vector from class k and that from class l at time t is given by (ck − cl)e 1 KA(t)− 1 KB(t) + Σ1/2(W kt −W lt), provided that the first deterministic term dominates the random second term, thus ensuring separation, which are precisely the conditions specified in Theorem 3.1. We term this model as the isotropic feature model, or I-model for short; as theH matrix has identity matrices as its blocks and consequently the dynamics do not prescribe any preferred directions for each class.

heading: 3.2.2 Logits-as-Features Model

text:
An important type of features in neural nets is the logits, the outputs of the neural net before the softmax layer. A logit vector (or logits) is K-dimensional, and in this model we identify X̃ k (t) as the logits at time t of a generic sample from the k-th class. In a well-trained neural net, the logits of a learned data instance from the k-th class should have its k-th logit being the largest, and heuristically, the other coordinates should be approximately equal and negative. As we shall detail in Appendix B, the exact dynamics of neural net training pushes the logits X̃ k by its margin, dk := ek − softmax(X̃ k ), which roughly aligns with the direction of ek − 1p/K. This suggests us how to choose the metric under which local elasticity acts: we can chooseH l,k such that it always aligns X̃ k in the direction of dk, that is,
Hij = H̄ j := djd
> j
‖dj‖22 ∈ Rp×p, dj := ej −
1
K 1p ∈ Rp, j ∈ [K]. (8)
Roughly speaking, the map x 7→ H̄jx projects x in the direction of dj and ideally aligns x with dj after iterative applications; hence, H̄
j can be viewed as an approximation of the nonlinear transformation in the exact dynamics in the sense that the direction of their stationary point coincides. Furthermore, H̄j thus defined has operator norm 1; thus, it does not affect the magnitudes, but only directions. Note thatH does not satisfy the condition in Theorem 3.1 as it is not symmetric; yet the separation theorem can be easily extended in light of the following proposition.
Proposition 3.3 (L-model). LetH be the same as in equation (8), then the solution to the LE-ODE (6) is given by
X̄(t) = c0 + C1de 1 KA(t)− 1 KB(t) + ( K−1∑ l=1 C2lf l ) e 1 KA(t)+ 1 K(K−1)B(t), (9)
where f l’s are fixed vectors in RK 2 , c0 ∈ RK 2 is a constant vector with K(K − 1) degrees of freedom, and C1, C2l ∈ R, for l ∈ [K − 1] are constants.
The specific form of f l is not the focus here; equation (9) allows us to prove the statement of Theorem 3.1 under this choice of H . The proof of Proposition 3.3 is deferred to Appendix C.2.2, where we also provide the analytical solution of f l’s when K = 3.
We term this model as the logits-as-features model, or L-model for short, because it is an elaborate model specifically for logits. In Section 4, we provide concrete demonstrations of our abstract feature vector X̃ as logits under L-model. We numerically simulated our LE-ODE and compared its predicted dynamics with real deep learning training dynamics. The experimental results provide strong empirical support for the validity of L-model.

heading: 3.3 Connection with Neural Collapse

text:
Neural collapse is a recent phenomenological finding on the geometry of the logits learned by deep neural nets at convergence with the cross-entropy loss [38] (see an explanation of neural collapse in [16]). Simply speaking, taking our L-model as an example, with balanced training samples, this model asserts that the logit vectors from different classes at convergence form an equiangular tight frame (ETF). ETFs are the best configuration to spread K unit vectors in an ambient space of p dimensions. Formally, we say a set of vectors {si}Ki=1 form an ETF in Rp if they are the columns of a matrix
S =
√ K K − 1Q ( IK − 1 K 1K1 > K ) , (10)
where Q ∈ Rp×K , and Q>Q = IK . As a direct corollary of Proposition 3.3, when H is set according to equation (8), we find that our L-model also predicts the existence of neural collapse from the local elasticity point of view.
Proposition 3.4 (Neural Collapse of the LE-ODE). Under L-model and the same setup as in Theorem 3.1, if γ(t) > 0 and there exists some T > 0 such that B(t) < 0 for t ≥ T , then{ X̄ k (t)/‖X̄k(t)‖ }K k=1 forms an ETF as t→∞.

heading: 4 Experiments

text: We perform various experiments to test our theory, where we choose logits as our protagonist2.

heading: 4.1 Setup

text:
Datasets and Models. We perform experiments on a synthesized dataset called GEOMNIST containing K = 3 types of geometric shapes (RECTANGLE, ELLIPSOID, and TRIANGLE) and on CIFAR-10 ([28], denoted by CIFAR) with K ∈ [2, 3] classes. A few samples from GEOMNIST are shown in Figure 2. We vary the number of training samples per class and label pollution ratio perr and use variants of the ALEXNET ([29]) model. More details can be found in the Appendix.
Training Configurations. All models are trained for T = 105 iterations (for GEOMNIST) or T = 3 × 105 iterations (for CIFAR) with a learning rate of 0.005 and a batch size of 1 under the softmax cross-entropy loss. Models on GEOMNIST converged with training and validation losses to zero, and those on CIFAR to validation accuracies greater than 90%.
Estimation Procedures. Each experiment is repeated for ntrial = 100 independent runs to estimate X̄(t). We use both the isotropic feature learning model (Section 3.2.1) and the logits-as-features model (Section 3.2.2), denoted by L-model and I-model respectively, to estimate α(t) and β(t). The
2Code for reproducing our experiments is publicly available at github.com:zjiayao/le_sde.git.
0 2000 4000 6000 8000 10000
t
0
5
10
15
20
Â (t
), B̂
(t )
Â(t)
B̂(t)
(a) GEOMNIST (I-model).
0 2000 4000 6000 8000 10000
t
−0.005
0.000
0.005
0.010
0.015
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(b) GEOMNIST (I-model).
0 5000 10000 15000 20000 25000 30000
t
−2.5
0.0
2.5
5.0
7.5
10.0
12.5
15.0
Â (t
), B̂
(t )
Â(t)
B̂(t)
(c) CIFAR (I-model).
0 5000 10000 15000 20000 25000 30000
t
−0.0010
−0.0005
0.0000
0.0005
0.0010
0.0015
0.0020
0.0025
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(d) CIFAR (I-model).
0 2000 4000 6000 8000 10000
t
−20
−15
−10
−5
0
5
10
15
20
Â (t
), B̂
(t )
Â(t)
B̂(t)
(e) GEOMNIST (L-model).
0 2000 4000 6000 8000 10000
t
−0.010
−0.005
0.000
0.005
0.010
0.015
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(f) GEOMNIST (L-model).
0 5000 10000 15000 20000 25000 30000
t
−30
−20
−10
0
10
Â (t
), B̂
(t )
Â(t)
B̂(t)
(g) CIFAR (L-model).
0 5000 10000 15000 20000 25000 30000
t
−0.003
−0.002
−0.001
0.000
0.001
0.002
0.003
0.004
α̂ (t
), β̂
(t )
α̂(t)
β̂(t)
(h) CIFAR (L-model).
Figure 3: Estimated Â(t), B̂(t), α(t), and β(t). The first row was estimated using I-model and the second L-model; the first two columns are on GEOMNIST and the last two on CIFAR. The first and third rows show Â(t) and B̂(t) and the other two rows α̂(t) and β̂(t).
L-model is used only when K = 3. To estimate α(t) and β(t), we first estimate A(t) and B(t) by
(I-model) Â(t) = avg avgk log ∣∣∣|X(X̄k−|X)K−1 c0c K−1 k ∣∣∣, B̂(t) = − avg avgk log ∣∣∣ c0ck X̄k−|X|X ∣∣∣, |Xt := avgl X̄ l t,
(L-model)
{ Â(t) = A′(t) + 2B′(t),
B̂(t) = 2(B′(t)−A′(t)), A ′(t) := log ∣∣∣〈X̄>v1 − 1〉∣∣∣ , B′(t) := log ∣∣∣〈X̄> (v2 − 43v1)〉∣∣∣ , (11)
where vector division is interpreted entry-wise. We write avgl for averaging over the class index, avg for averaging over the coordinates, and define 〈X〉(t) := X(t)/X(0). We explain how and why to choose the vectors v1 and v2 in Appendix D.1. The main idea is to view the eigenvectors of the Kp-by-Kp drift matrix as a concatenation of K vectors of dimension p and construct their linear combinations such that one or more independent components in the solution vanishes. With A(t) and B(t) estimated, we use the Savitzky - Golay filter to obtain α̂(t) and β̂(t) through numerical differentiation. For GEOMNIST and CIFAR datasets, we choose window sizes of this filter as 191 and 551, respectively, in Figure 3, and 21 and 21, respectively, in Figure 5.
We assess the tail of α̂(t) and β̂(t) by a tail index defined as rα := sups {s : limt→∞ α(t) · ts <∞} and rβ is defined similarly. We estimate r̂α by fixing an interval [T1, T2] with T1 < T2 sufficiently large such that we may ignore terms with smaller order and have r̂α = 1− avgT1≤t≤T2 logα(t) log(1+t) , and similarly for r̂β . We use the estimates from the last 1000 iterations for averaging in our experiments.

heading: 4.2 Results

text:
Local Elasticity in Neural Net Training. Local elasticity manifests from our model as the heaviness of the tail of γ(t) = α(t) − β(t), and in Figure 3, we plot the estimations Â(t), B̂(t), α̂(t), and β̂(t) using both I-model and L-model. We note that (i) The estimations from the two models are visually similar, especially in the late stage of training when t is large; (ii) The major difference lies in the initial stage, where the estimates from L-model behave slightly wilder. This is not surprising because of the effect of the unknown constant offset c0 in the L-model; (iii) Both α̂(t) and β̂(t) behave similarly on both datasets.
Phase Transition of Separability. Theorem 3.1 states that separation of features under the LE-SDE takes place when γ(t) = α(t)− β(t) = ω(1/t), or roughly speaking, when rγ = min{rα, rβ} < 1. Although we cannot directly control α(t) and β(t), we can bias them by tuning the label corruption ratio perr. When perr ≈ p∗err := 2/3, we are in effect assigning labels completely at random and thus we expect a phase transition of separability should happen around p∗err. This is indeed the story depicted in Figure 4: Figures 4a and 4b show that the validation loss and accuracy for p ≥ perr
0 2000 4000 6000 8000 10000
t
−15
−10
−5
0
5
10
X̄ (t
)
X̄00 X̄01 X̄02
(a) GEOMNIST (I-model).
0 2000 4000 6000 8000 10000
t
−7.5
−5.0
−2.5
0.0
2.5
5.0
7.5
10.0
X̄ (t
)
X̄00 X̄01 X̄02
(b) GEOMNIST (L-model).
0 5000 10000 15000 20000 25000 30000
t
−2
−1
0
1
2
X̄ (t
)
X̄00 X̄01 X̄02
(c) CIFAR (I-model).
0 5000 10000 15000 20000 25000 30000
t
−2
−1
0
1
2
X̄ (t
)
X̄00 X̄01 X̄02
(d) CIFAR (L-model).
Figure 5: Simulated LE-ODE solutions versus genuine dynamics. We use α̂(t) and β̂(t) estimated from I-model ((a) and (c)) or L-model, ((b) and (d)) and numerically simulate the solution under the L-model. The results were overlaid with true dynamics from neural nets. We note L-model in general imitated true dynamics reasonably well.
are not increasing over time and in Figure 4c we observe the minimum tail index of α(t) and β(t) crosses 1 from below around perr = 0.6, entering the non-separable regime (shaded in red) from the separable regime (shaded in green), given in Theorem 3.1.
Simulating DNN Dynamics via LE-ODE. Having estimated α(t) and β(t), it is natural to ask, to what capacity can our LE-ODE models recover the real dynamics of deep neural nets? We use the forward Euler method to simulate the L-model using α̂(t) and β̂(t) estimated from either I-model or L-model, We choose K = 3 and show in Figure 5 the simulated solution (solid line) with error bars depicting one standard deviation over 500 independent runs, overlaying on the real dynamics from DNNs in the background (shaded transparent markers). As the moving average may reduce the magnitudes of α(t) and β(t), we rescale the simulated paths such that its first coordinate is approximately equal to the ground truth at convergence. Note that estimations from L-model can faithfully recover the genuine dynamics from neural nets, whereas those from I-model fail, notably in Figure 5a, where the simulated paths preserve the relative magnitude but fail to identify the correct order of three logits.

heading: 5 Discussion and Future Works

text: In this study, we introduce LE-SDE/ODE models that draw inspiration from the local elasticity phenomenon. Conditions for sharp phase transition of separability of features are derived. We also show that once the elasticity strengths α(t) and β(t) are well estimated, our model can faithfully simulate the dynamics of neural nets. We outline a few interesting problems for future research while leaving the details in the Appendix. (i) General LE Matrix. A similar result as in Theorem 3.1 may be expected for symmetric but no necessarily semi-definite LE matrices E(t). (ii) Minibatch Training, Imbalanced Datasets, and Label Corruptions. Generalizing the drift matrix to M t = (Et ⊗ P ) ◦H/K for a K-by-K doubly stochastic matrix P can be used to model various sampling effects. (iii) Beyond L-model for Imitating Genuine Dynamics of DNNs. Although the L-model is shown to be able to mimic the real dynamics reasonably well, we postulate that a more precise model might have its (i, j)-th block encode the other directions other than dj .

heading: Acknowledgements

text: This work was supported in part by NSF through CCF-1934876, an Alfred Sloan Research Fellowship, the Wharton Dean’s Research Fund, and ONR Contract N00014-19-1-2620. We would like to thank Dan Roth and the Cognitive Computation Group at the University of Pennsylvania for stimulating discussions and for providing computational resources.

title: Imitating Deep Learning Dynamics via Locally Elastic Stochastic Differential Equations