[name=Theorem,numberwithin=section]thm
[name=Corollary,numberwithin=section]cor
[name=Observation,numberwithin=section]obs University College Doctor of
Philosophy

# Generalization Through the Lens of Learning Dynamics

Clare Lyle

###### Abstract

A machine learning (ML) system must learn not only to match the output of a
target function on a training set, but also to generalize to novel situations
in order to yield accurate predictions at deployment. In most practical
applications, the user cannot exhaustively enumerate every possible input to
the model; strong generalization performance is therefore crucial to the
development of ML systems which are performant and reliable enough to be
deployed in the real world. While generalization is well-understood
theoretically in a number of hypothesis classes, the impressive generalization
performance of deep neural networks has stymied theoreticians. In deep
reinforcement learning (RL), our understanding of generalization is further
complicated by the conflict between generalization and stability in widely-
used RL algorithms. This thesis will provide insight into generalization by
studying the learning dynamics of deep neural networks in both supervised and
reinforcement learning tasks.

We begin with a study of generalization in supervised learning. We propose new
PAC-Bayes generalization bounds for invariant models and for models trained
with data augmentation. We go on to consider more general forms of inductive
bias, connecting a notion of training speed with Bayesian model selection.
This connection yields a family of marginal likelihood estimators which
require only sampled losses from an iterative gradient descent trajectory, and
analogous performance estimators for neural networks. We then turn our
attention to reinforcement learning, laying out the learning dynamics
framework for the RL setting which will be leveraged throughout the remainder
of the thesis. We identify a new phenomenon which we term capacity loss,
whereby neural networks lose their ability to adapt to new target functions
over the course of training in deep RL problems, for which we propose a novel
regularization approach. Follow-up analysis studying more subtle forms of
capacity loss reveals that deep RL agents are prone to memorization due to the
unstructured form of early prediction targets, and highlights a solution in
the form of distillation. We conclude by calling back to a different notion of
invariance to that which started this thesis, presenting a novel
representation learning method which promotes invariance to spurious factors
of variation in the environment.

###### Acknowledgements.

First and foremost, this thesis would not have been possible without the
invaluable guidance and support of my supervisors, Yarin Gal and Marta
Kwiatkowska. Their advice, mentorship, and insight over the past almost four
years has been vital to my development as a scientist. I am deeply grateful to
both for being so incredibly generous with their time, particularly in the
early years of the DPhil, and for giving me the freedom to set my own research
direction and explore a broad range of topics. Thanks are also due to Marc
Bellemare, Pablo Samuel Castro, and Prakash Panangaden for their mentorship
prior to the start of my DPhil. I learned many invaluable lessons about how to
do good research from their example, which have served me well throughout my
DPhil. I have had the good fortune to engage in a number of collaborations
while at Oxford. Many thanks go to Amy Zhang, with whom I wrote the first
published paper of my DPhil, for modelling how to organize and execute a
research project. I am also indebted to Mark Rowland and Will Dabney, who have
been a joy to work with on many of the papers that appear in this document.
This thesis would not have been possible without valuable discussions and
collaborations with Benjamin Bloem-Reddy, Mark van der Wilk, Benjie Wang,
Angelos Filos, Natasha Jaques, Greg Farquhuar, Lisa Schut, Robin Ru, Aidan
Gomez, Lorenz Kuhn, Jannick Kossen, Neil Band, Georg Ostrovski, Shagun
Sodhani, and Andreas Kirsch. Beyond direct collaborations, I also benefited
immensely from being part of the broader machine learning community at Oxford.
My perspective as a researcher was enriched by having a group of brilliant
people with whom I could run over nascent ideas, debug a tricky experiment,
and gripe about Reviewer 2. Thanks go in particular to Joost, Milad, Seb,
Panos, Tim, Jan, Freddie, Pascal, Matt, Pascale, Rhiannon, Luca, Andrea, Emi,
Michael, Sahra, Charline, and Jakob. Thanks as well go to the many people I
interacted with at DeepMind, including Daniel, Dave, Remi, Diana, Anna, Bilal,
Bernardo, and Mo. Finally, I would be remiss to omit my family, whose
unconditional support has provided the solid foundation on which I have been
able to take risks and and explore as a researcher. Thanks go as well to the
friends I’ve made in Oxford: Alex, Caitlin, Nayani, Colin, Matt, Rahul, Imi,
Becca, and my teammates from OUBbC and UCBC. Thanks in particular to the
individual who did not want to be named in this document for challenging me to
address meaningful problems and for proofreading many chapters of this thesis.

###### Contents

  1. 1 Introduction
    1. 1.1 Learning to generalize
    2. 1.2 Learning to act
    3. 1.3 Understanding the learning process
    4. 1.4 Thesis contributions and structure
  2. 2 Background & literature review
    1. 2.1 Learning frameworks
    2. 2.2 Generalization in supervised learning
    3. 2.3 Generalization in reinforcement learning
  3. 3 The role of invariance in generalization
    1. 3.1 Introduction
    2. 3.2 Background
    3. 3.3 Invariance and optimization dynamics
    4. 3.4 Invariance and generalization bounds
    5. 3.5 Conclusions
  4. 4 Training speed and model selection
    1. 4.1 Introduction
    2. 4.2 Background on Bayesian inference
    3. 4.3 Marginal likelihood estimation with training statistics
    4. 4.4 Training speed and interference
    5. 4.5 Conclusions
  5. 5 Dynamics of reinforcement learning
    1. 5.1 Introduction
    2. 5.2 Mathematical framework
    3. 5.3 Value function & representation dynamics
    4. 5.4 Auxiliary tasks
    5. 5.5 Experiments
    6. 5.6 Conclusions
  6. 6 Capacity loss
    1. 6.1 Introduction
    2. 6.2 Learning capacity in neural networks
    3. 6.3 InFeR: mitigating capacity loss with feature regularization
    4. 6.4 Conclusions
  7. 7 Interference and generalization
    1. 7.1 Introduction
    2. 7.2 Learning dynamics and smoothness
    3. 7.3 Generalization and interference in deep RL
    4. 7.4 Post-training distillation and generalization
    5. 7.5 Conclusions
  8. 8 Generalization across environments
    1. 8.1 Introduction
    2. 8.2 Background on causality
    3. 8.3 Problem setting
    4. 8.4 State abstractions and causal feature sets
    5. 8.5 State abstractions for rich observations
    6. 8.6 Conclusions
  9. 9 Conclusion
  10. 10 The role of invariance in generalization
    1. 10.1 Proofs
    2. 10.2 Examples, counterexamples, tighter bounds
    3. 10.3 Computation details for PAC-Bayes bounds
  11. 11 Training speed and model selection
    1. 11.1 Linear model combination
    2. 11.2 Proofs of theoretical results
    3. 11.3 Experiments
  12. 12 Dynamics of reinforcement learning
    1. 12.1 Auxiliary results
    2. 12.2 Proofs
    3. 12.3 Additional results from Table 5.1
    4. 12.4 Beyond diagonalisability assumptions
    5. 12.5 Extensions beyond one-step temporal difference learning
    6. 12.6 Bayes-optimality of RSBFs
    7. 12.7 Further discussion of features and operator decompositions
    8. 12.8 Learning dynamics for ensemble prediction
    9. 12.9 Experimental details
  13. 13 Capacity loss
    1. 13.1 Proofs
    2. 13.2 Sequential supervised learning
    3. 13.3 Atari evaluations
  14. 14 Interference and generalization
    1. 14.1 Proofs
    2. 14.2 Experiment details
    3. 14.3 Additional numerical evaluations
    4. 14.4 Additional empirical results
  15. 15 Generalization across environments
    1. 15.1 Implementation details

###### List of Figures

  1. 1.1 A visualization of a fully connected multi-layer perceptron.
  2. 1.2 A figure visualizing the RL problem; featuring the environment (pictured on the left) and agent (right) dichotomy.
  3. 2.1 State abstractions allow an agent to treat similar states as though they were equivalent, potentially simplifying planning and generalization.
  4. 2.2 Visualization of multi-environment generalization framework with coinrun levels sampled from the ProcGen benchmark (Cobbe et al., 2020).
  5. 3.1 A visualization of a fully connected multi-layer perceptron.
  6. 3.2 Measurements over the course of training a CNN using different DA and FA approaches.
  7. 3.3 Comparing the ranks of gradients computed on minibatches consisting of 32 independent random input samples, or 8 random samples from the training set augmented with all 90 degree rotations (resulting in 32 non-independent samples).
  8. 4.1 Visualization of the relationship between the Bayesian marginal likelihood and training speed.
  9. 4.2 Evaluation of our estimators.
  10. 4.3 Illustration of marginal likelihood estimation via our method on NTK models.
  11. 4.4 Effect of minibatch gradient steps on train, test, and minibatch loss over the course of training a variety of neural network architectures on CIFAR-10.
  12. 4.5 Linear combinations of DNNs on FashionMNIST.
  13. 4.6 Rankings of a randomly-selected subset of models from the NASBench-201 dataset given by training speed estimators.
  14. 5.1 An example of qualitatively different value function dynamics for a two-state MDP for 1-step temporal difference learning and Monte Carlo learning, with fixed point $V^{\pi}$ in red.
  15. 5.2 Visualizing the Grassmann distance between two 2-dimensional subspaces.
  16. 5.3 Study of feature dynamics.
  17. 5.4 Transfer of EBFs, RSBFs, and RFs across the value-improvement path of a chain MDP, with and without the value function as an additional feature.
  18. 5.5 Learning curves for DDQN, DDQN+RC, DDQN+Ensemble, DDQN+REM, and QR-DQN agents on several dense reward ALE environments.
  19. 5.6 Sweep over number of auxiliary heads for RC, REM and Ensemble.
  20. 6.1 Deeper analysis on a sequential MNIST prediction setting showing that target-fitting capacity sees the greatest reduction in low capacity networks with ReLU units.
  21. 6.2 Neural networks exhibit a decline in their ability to fit random network outputs over the course of training in two demonstrative Atari environments.
  22. 6.3 Feature rank and performance over the course of training.
  23. 6.4 Evaluation of the correlation between feature rank and performance.
  24. 6.5 Visualization and evaluation of InFeR.
  25. 6.6 Deeper investigation into the effect of InFeR.
  26. 7.1 Value function evolution and eigendecomposition on the MountainCar environment.
  27. 7.2 Networks trained to fit high-frequency target functions exhibit pathological interpolation properties when later fine-tuned on a value function.
  28. 7.3 Update matrices for a toy kernel regression problem.
  29. 7.4 Evolution of interference in DQN agents trained on Atari.
  30. 7.5 Update dimension of actor-critic methods in ProcGen.
  31. 7.6 Performance and robustness to perturbations of different distillation approaches in games from the Atari suite.
  32. 7.7 Effect of policy distillation on generalization in environments from the Procgen suite.
  33. 8.1 Graphical model of a block MDP with emissions $x$ from a state which can be decomposed into a causal component $s$ and a spurious component $\eta$.
  34. 8.2 Illustration of a causal state abstraction where the observation is a concatenation of variables from a causal graph.
  35. 8.3 The presence of spurious uncorrelated variables in the state can still lead to poor generalization of linear function approximation methods.
  36. 8.4 Evaluation of our nonlinear MISA algorithm.
  37. 8.5 The Cheetah Run environment from Deepmind Control with different camera angles. The first two images are from the training environments and the last image is from evaluation environment.
  38. 8.6 Performance of nonlinear MISA on RL tasks.
  39. 11.1 Relative rankings given by optimize-then-prune, ML, and estimated $\mathcal{L}(\mathcal{D})$. Left: feature selection. Middle: prior variance selection. Right: RFF frequency selection. Rankings are consistent with what our theoretical results predict. Results are averaged over $5$ runs.
  40. 11.2 Weight assigned to subnetwork by SGD in a deep neural network (x-axis) versus the subnetwork performance (estimated by the sum of cross-entropy, on the y-axis) for different FashionMNIST classes. The light blue ovals denote depict $95\%$ confidence intervals, estimated over 10 seeds (i.e. 2$\sigma$ for both the weight and SOTL). The orange line depicts the general trend.
  41. 11.3 Weight assigned to subnetwork by SGD in a deep neural network (x-axis) versus the subnetwork performance (estimated by the sum of cross-entropy, on the y-axis) for different CIFAR-10 classes. The light blue ovals denote depict $95\%$ confidence intervals, estimated over 10 seeds (i.e. 2$\sigma$ for both the weight and SOTL). The orange line depicts the general trend.
  42. 11.4 Linear combinations of DNNs on FashionMNIST. Left: ensemble weights versus the test loss for parallel training; we observe a negative correlation. Middle: SOTL (standardized by the number of training samples) versus test loss for concurrent and concurrent training. We observe a strong correlation indicating that the SOTL generalizes well. Right: training curves for the different models in concurrent training schemes. All results are averaged over $10$ runs, and standard deviations are shown by the shaded regions around each observation. The model parameters, given in the parentheses, are the number of layers ($l$), nodes per layer ($n$) and kernel size ($k$), respectively.
  43. 12.1 Value functions learned by the ensemble head at index 0 for different training regimes. Plot titles of form (feature initialization scheme, train/fix weight matrix, number of heads in ensemble). Observe that the representation learned with fixed weights tends to converge to smoother eigenfunctions than those learned with weights that are also allowed to train. f
  44. 12.2 Values of ensemble feature at index 0 for different training regimes. Plot titles of form (feature initialization scheme, train/fix weight matrix, number of heads in ensemble). Observe that the representation learned with fixed weights tends to converge to smoother eigenfunctions than those learned with weights that are also allowed to train.
  45. 12.3 Learning curves on sparsified MsPacman (left), sparsified Seaquest (centre), and sparisifed Q*bert (right).
  46. 12.4 Results of hyper-parameter sweep for Random Cumulant (RC) method, where each row is for a different value of multiplicative scale applied to the auxiliary losses and each bar corresponds to the number of auxiliary heads ($M$). Note that the first row of results corresponds to initializing a network with the auxiliary heads, but setting the weight to zero, effectively disabling the auxiliary task.
  47. 12.5 Results of hyper-parameter sweep for the Ensemble method, where each row is for a different learning rate and each bar corresponds to the number of auxiliary heads ($M$).
  48. 12.6 Results of hyper-parameter sweep for the REM method, where each row is for a different learning rate and each bar corresponds to the number of auxiliary heads ($M$).
  49. 13.1 Mean squared error at the end of training on each iteration of the hash-MNIST task. Target-fitting error increases over time in smaller networks, but increasing the depth or width of the network slows down capacity loss, enabling positive transfer in the largest networks we studied.
  50. 13.2 Mean squared error after 2e3 training steps on the random-MNIST task. Target-fitting error increases over time in under-parameterized networks, but increasing the depth or width of the network slows down capacity loss, enabling positive transfer in the largest network we studied.
  51. 13.3 Mean squared error after 2e3 training steps on the threshold-MNIST task. Target-fitting error increases over time in under-parameterized networks, but increasing the depth or width of the network slows down capacity loss, enabling positive transfer in the largest network we studied.
  52. 13.4 Effect of adding InFeR to the regression objective in a random reward prediction problem on the non-stationary MNIST environment studied previously. We see that the InFeR objective produces networks that can consistently outperform those trained with a standard regression objective, exhibiting minimal capacity loss in comparison to the same network architecture trained on the same sequence of targets.
  53. 13.5 Hyperparameter sweeps for the DDQN+InFeR agent. Each contour plot shows average capped human-normalized score at the end of training marginalized over all hyperparameters not shown on its axes.
  54. 13.6 Hyperparameter sweeps for the Rainbow+InFeR agent. Each contour plot shows average capped human-normalized score at the end of training marginalized over all hyperparameters not shown on its axes.
  55. 13.7 Feature rank and performance of RL agents on demonstrative Atari environments.
  56. 13.8 feature rank of agent representations over the course of training on all 57 games in the Atari benchmark. We compare Rainbow against Rainbow+InFeR.
  57. 13.9 Mean squared error, after $10000$ training steps for the target-fitting on random network targets. We also show the corresponding feature rank of the pre-trained neural network (before target-fitting).
  58. 13.10 Full evaluation of capped human-normalized performance on Atari benchmarks for the default Rainbow architecture.
  59. 13.11 Full evaluation of capped human-normalized performance on Atari benchmarks in the double-width Rainbow architecture.
  60. 13.12 Full evaluation of raw scores on Atari benchmarks for the default Rainbow architecture.
  61. 13.13 Full evaluation of raw scores on Atari benchmarks for the double-width Rainbow architecture.
  62. 13.14 Evaluations of the effect of InFeR on performance of a Double DQN agent. Overall we see a slight improvement in average performance over all games, alongside a significant improvement in Montezuma’s Revenge.
  63. 14.1 Example levels from the dodgeball environment.
  64. 14.2 Fourier decomposition of Atari value functions when viewed as a function of time. We sample $k$ consecutive states from the replay buffer and compute the predicted value on each state (fixing an arbitrary action) to get a function $V:\\{1,\dots,k\\}\rightarrow\mathbb{R}$. We then compute the Fourier decomposition of this function. The top row shows indices $k=0\dots 50$, while the bottom row omits the $k=0$ index (the constant function) to better illustrate the rate of decay of the spectrum of each function.
  65. 14.3 Numerical evaluations of kernel gradient descent with an RBF kernel. The MDP in question is a "circle MDP" whose states are integers $n\in\\{1,\dots,50\\}$. We assume the agent is ‘trained’ on states 1 to 40, and does not perform value function updates on the final ten states, use the policy which always takes the agent from state $n$ to $n+1\mod 50$, and set a single reward at state 25. Each row corresponds to a different value of the discount factor $\gamma$: the top corresponds to $\gamma=0.5$, and the bottom to $\gamma=0.99$. Each column corresponds to the lengthscale which parameterizes the kernel, going left to right: 0.01, 1.0, and 100. The left hand side and right hand side are distinguished by the number of update steps which the TD dynamics are evaluated for. The LHS runs TD for only 20 steps, while the RHS runs it for 100 steps. MC updates are run for 1500 steps on both figures. We see that for $\gamma=0.99$, the larger-lengthscale kernel predictions diverge under TD dynamics, though not Monte Carlo. The Monte Carlo dynamics further nicely illustrate the trade-off between generalizing out of the training set and ability to fit the discontinuities of the value function on the training set. The larger lengthscale has lower MSE from the value function on the test set, but fails to fit the discontinuity of the value function at the reward state. Meanwhile, the smaller lengthscales easily fit the value function on the training set but predict zero for all over states.
  66. 14.4 Generalization of predicted function under kernel regression using $n$-step return targets evaluated on a random subset of states (left), and projecting value function onto top or bottom eigenvectors of $P^{\pi}$ (right). We see a similar trend where for larger $n$ (corresponding to smoother targets), the kernel regression method generalizes better with increasing dataset sizes. For smaller $n$ and for the projection of $V^{\pi}$ onto non-smooth eigenvectors, adding additional data points does not improve generalization performance.
  67. 14.5 Results from post-training distillation on a variety of objectives. We note that advantage regression tends to exhibit the lowest update rank, with the qr agent tending to exhibit the highest update rank and the q-regression objectives falling somewhere in between. Because the behaviour cloning objective minimizes a cross-entropy loss rather than a regression loss, further investigation is required to understand how the trajectory of its update dimension differs from those of the regression objectives.
  68. 14.6 Update matrices for distributional and DQN agents on four games from the Atari suite, chosen to represent a range of reward densities and difficulties. Each iteration corresponds to 1e5 training frames.

###### List of Tables

  1. 3.1 Generalization performance for a permutation-invariant point cloud classification task (see text for details).
  2. 5.1 Summary of dynamics and limiting solutions under some common auxiliary tasks in the limit of infinitely-many prediction outputs.
  3. 12.1 Summary of decompositions of various matrices associated with MDP transition operators, and associated features.
  4. 15.1  A complete overview of used hyper parameters.

##  Chapter 1 Introduction

###  1.1 Learning to generalize

The ability to generalize a lesson from the classroom to the real world is
what separates learning from memorization. In a range of tasks ranging from
mathematics to language, humans are remarkably skilled at identifying abstract
patterns, and applying these patterns to novel contexts. The reader is
unlikely to have previously encountered the sentence ‘the armadillo tipped its
blue hat and returned to the game of marbles’, and yet would likely have no
trouble interpreting it, or answering questions concerning the colour of the
armadillo’s hat. However, the range of domains in which humans exhibit this
ability to learn and generalize is limited. A human can easily parse natural
language, but will struggle to identify structure in strings of base pairs
arising from a genome. In these settings, we can benefit from computational
tools. Machine learning approaches, in particular the training of deep neural
networks on large datasets, present a promising direction towards the
development of general algorithms which can identify patterns in data and
solve a range of problems.

In a typical machine learning pipeline, the practitioner collects data (a
training set) which is then fed into a machine learning algorithm, with the
hope that a system which can accurately model this data will have captured the
underlying structure of the task. The training set will not encompass the set
of all possible inputs a model may receive at deployment; a learned predictor
must extrapolate from the data it was trained on in order to make useful
predictions. Generalization is crucial both to obtain good performance and to
ensure the safety and reliability of these systems when they encounter data
that was not seen during training. Yet how do we ensure that these highly
expressive systems are learning and not memorizing? This question is a central
concern of a long line of literature, and of this thesis.

####  1.1.1 Defining generalization

The machine learning community broadly distinguishes between two classes of
generalization: within-distribution generalization, and out of distribution
(OOD) generalization. While both types of generalization concern the
performance of a predictor on data not seen during training, they differ in
their structural assumptions on how this data is generated. Within-
distribution generalization assumes that the process by which the training
data was collected will also generate the data on which we will evaluate the
trained model. This assumption is used in a rich theoretical literature which
provides provable guarantees on the generalization performance of certain
classes of learning algorithms. However, it is not reflected in many of the
settings found in practice, where the procedure by which the training data is
collected differs from how data is generated at evaluation; this difference is
widely referred to as distribution shift.

The real world is full of nonstationarities which can induce distribution
shift: a user’s taste in films will evolve as they age, slang terms enter and
leave common usage, and interest in advertisements for winter boots may vary
with the seasons. Data collected one month may quickly fall out of date and
cease to be reflective of the phenomenon being modelled within a matter of
weeks. Even the ways in which data is collected may introduce biases into the
training set. Many image datasets, for example, include the main subject
centred nicely in the middle of the frame, signalling to a learning algorithm
that the corners of an image do not contain relevant information. These
nonstationarities and biases can result in large distribution shifts when the
model is deployed. The technical difficulty of overcoming such distribution
shifts will depend on their structure and magnitude. Work on developing models
which are robust to distribution shifts typically must make explicit
assumptions on the type of shift being considered.

####  1.1.2 The importance of generalization

While we often marvel at great feats of memorization, such as memory champions
who can recite the order of a shuffled deck of cards after a minute’s
concentration or taxi drivers who can recall the street maps of large cities
by heart, memorization is discouraged when it occurs in place of
understanding. The difference between the two is one of generalization:
memorizing multiplication tables enables quick recall, but learning how to
multiply numbers together algorithmically enables generalization to previously
unseen number pairs. While we hesitate to anthropomorphize a linear regression
model in saying that it ‘understands’ the relationship between inputs and
outputs, we nonetheless seek out a similar phenomenon in our machine learning
systems. It is not sufficient to make perfect predictions on the training
data: the model must be able to apply the relationship between input and
output to new contexts, and make accurate predictions in these contexts.

This generalization is crucial if we hope to see the power of machine learning
deployed in real-world settings, where mistakes on novel inputs can have
catastrophic consequences. Unforeseen distribution shifts in medical data,
such as replacing an imaging device in a hospital, can result in a decline in
performance that has serious ramifications for patients’ health. Cross-
validation, the canonical approach to evaluating generalization performance,
will not identify the model’s robustness to distribution shifts that the
developer did not already foresee and include in the evaluation set. This
presents a particular challenge for practitioners; even if a decline in
performance is detected, the best we can hope to do is re-train or fine-tune
the model to improve its performance on the new data. This can become much
more expensive than if the model had learned relationships that generalized
well off the bat.

####  1.1.3 Deep learning

Figure 1.1: A visualization of a fully connected multi-layer perceptron. The
neural network consists of a series of layers, where each layer applies a non-
linear activation to a linear combination of activations from the previous
layer. Each circle in this figure corresponds to an activation unit, while the
lines correspond to weights connecting the left hand side layer to the right.
The output of the network in this case corresponds to a probability
distribution indicating the likelihood that the input image corresponds to a
given digit.

The deep learning revolution has spurred the growth of several international
conferences, the creation of multiple industrial AI labs, and the allocation
of three Turing awards. This is due to the success of deep neural networks
(DNNs, see Figure 1.1) at modelling a wide range of data modalities. The
applications of DNNs range from the benign, such as helping people with visual
impairments navigate a street and translating text from one language to
another, to those with the potential for malicious use, such as identifying
human faces in surveillance footage. Strikingly, DNNs often obtain impressive
generalization performance and are considered robust enough to be used in many
commercial applications.

The success of these models has brought attention to the discipline, but also
stymied theoreticians. Compared to the expressivity of deep neural networks,
traditional learning algorithms seek to model data by searching over a
relatively small class of functions. Theoretical analysis of these algorithms
crucially depends on the size of the function class in order to provide
guarantees on the expected error of the function found by the algorithm on new
data. In contrast, the set of functions expressible by a given neural network
architecture is so large as to result in vacuous results when traditional
analysis is applied to deep learning. This has opened up a number of exciting
approaches to study the generalization of DNNs which often have a more
experimental flavour than prior work on learning theory.

###  1.2 Learning to act

We will be particularly interested in studying machine learning systems which
are capable of doing things in the world, rather than passively outputting
predictions. This is captured by the reinforcement learning (RL) framework
(see Figure 1.2). At its core, a reinforcement learning problem consists of an
agent which can interact with an environment with the goal of maximizing the
cumulative reward signal that it receives. Just as a trainer can teach a dog
to sit by providing suitable rewards to reinforce the desired behaviour, we
can apply the power of machine learning algorithms to maximize a prespecified
reward function in the RL framework. We use the terminology behaviour policy
to refer to the distribution over actions that the agent takes in each state
of the environment. The optimal policy is the action-selection rule which
maximizes the expected cumulative reward from each state.

####  1.2.1 The reinforcement learning problem

Learning how to behave optimally is often aided by learning to predict the
expected cumulative reward the agent will receive after it visits a state and
then follows some behaviour policy. Our usage of the word ‘learn’ differs from
that used in other areas of machine learning, where it refers to the
identification of a relationship between inputs and outputs from a data set.
An RL agent does not receive explicit information about the optimal policy
from the environment; this policy must be deduced from the reward and
transition structure via planning. Reinforcement learning is thus closely
related the problem of optimal control. Control problems assume an
environment, modelled as a Markov Decision Process (MDP), with known
transition and reward structure and seek to identify an optimal behaviour
policy. Reinforcement learning also seeks to obtain an optimal policy, but
does not assume that the structure of the MDP is known in advance. Instead,
the agent must interact with the MDP in order to obtain information about the
reward and transition structure, and use this information to identify an
approximately optimal policy. The number of interactions with the environment
needed for an agent to identify an approximately optimal policy, its sample
complexity, is a key criterion by which RL algorithms are evaluated, and which
distinguishes reinforcement learning from optimal control, where the dynamics
of the world are known a priori.

Figure 1.2: A figure visualizing the RL problem; featuring the environment
(pictured on the left) and agent (right) dichotomy.

####  1.2.2 Generalization in reinforcement learning

Many problems of interest in reinforcement learning do not require
generalization, i.e. it is assumed that the agent will only encounter states
at deployment that it encountered during training. Indeed, the field contains
a rich literature on the analysis of tabular problems, whereby the states of
the MDP are simply an enumeration of integers and learning consists of
updating a lookup table. In many rich observation settings, however, there may
be an interesting functional relationship between the state observations
emitted by the environment and the reward and transition structure at that
state. While lookup tables are sufficient for small state spaces, most
applications of RL to real-world data necessitate generalization either due to
the magnitude of the state space or in order to be robust to distribution
shifts. For example, the angles and torques of a robot’s actuators can take on
a continuum of values, and the set of possible image inputs far exceeds the
size that can be fit into computer memory as a lookup table. Further, even if
such a representation of the value function were possible it is not clear
whether that would be desirable. Visual similarity between states can provide
information about the optimal policy and value function which may be useful to
the agent. A function approximator with an appropriate inductive bias will be
able to leverage this similarity to accelerate the learning process.

However, a dark side of generalization arises in reinforcement learning
problems: instability. Instability is particularly problematic in some of the
most popular algorithms in the deep RL literature, where careful hyper-
parameter tuning and engineering tricks are needed to prevent the network
parameters from diverging to infinite values. Excessive generalization can
also slow down learning if the inductive bias encoded by the function
approximator is not aligned with the structure of the environment. A bias
towards smooth functions might thus result in a network that fails to
accurately distinguish between states with large differences in value. The
tension between generalization and stability in deep reinforcement learning
presents an additional layer of difficulty to deep RL, as compared to
supervised deep learning, and is a problem we will explore in later chapters.

###  1.3 Understanding the learning process

Throughout this thesis, we will seek to understand how a model will generalize
by studying the optimization trajectory it took during training. In contrast,
most theoretical results in the literature characterize generalization using
only properties of the final outputs of a learning algorithm, i.e. the neural
network’s final trained parameters. Studying the trajectory of a learning
algorithm (which we will refer to as its learning dynamics) gives us the
opportunity to gain insights into a model that cannot be obtained by only
considering the final trained parameters. We will leverage these insights in
later chapters to obtain novel estimators with significant predictive power
over the ranking of a model’s final generalization performance.

####  1.3.1 Learning to generalize between data points

Many neural network training procedures leverage large datasets which cannot
fit onto a single GPU. To accelerate training, learning algorithms often
partition the data into subsets called minibatches. The learner then updates
its predictions for each minibatch iteratively. This procedure provides
extremely useful information about whether the agent is learning to generalize
or to memorize, by revealing whether the learner’s update based on one
minibatch has generalized to improve its predictions on the other minibatches.

A key intuition throughout this thesis is that generalization between disjoint
subsets of the training set can be indicative of generalization to the test
set. If an update to the network intended to improve its predictions for one
minibatch also improves the network’s predictions on many other data points,
this is likely to result in an improvement to the learner’s predictions on
novel inputs drawn from the same process that generated the training data. In
contrast, if an update does not improve the learner’s predictions on the other
data points in the training set, it is unlikely to improve the agent’s
predictions on the data it will see at deployment. This relationship will be
explored in greater depth in Chapter 4.

####  1.3.2 Stability vs extrapolation

Our study of learning dynamics in RL agents will be particularly enlightening,
as these dynamics are much more complex than their supervised counterparts.
Supervised learning algorithms tend to induce well-behaved dynamics that, even
if they correspond to a non-convex loss surface, nonetheless at least come
with reasonable guarantees on the convergence of learning algorithms to local
minima. Reinforcement learning, in contrast, incorporates nonstationarity as
part of the learning process. This nonstationarity has many forms: the
distribution of states that the agent visits will change as its policy
improves, but so will the target values that the agent is trying to predict.
This results in a dynamical system that superficially resembles the stable
gradient descent regime of supervised learning algorithms, but lacks
convergence guarantees many problem settings of interest.

Even more pernicious, reinforcement learning agents must also face the
challenge that the very properties of a function approximator which are
associated with better generalization in the supervised learning setting are
precisely those which can cause divergence in RL algorithms, as we will see in
later sections. This means that deep RL agents must overcome two distinct
hurdles in order to learn an optimal behaviour policy which also generalizes
to new settings. First, they must learn to behave optimally in the training
environment, while avoiding issues of divergence and other pathologies of
function approximation in RL. Second, having achieved a high-performing
policy, they must then ensure that the policy is robust to superficial changes
to the observations they receive from the environment.

###  1.4 Thesis contributions and structure

####  1.4.1 Contributions

Broadly speaking, this thesis presents a set of novel empirical and
theoretical tools to predict, understand, and improve generalization in deep
neural networks in a range of problem settings. Crucial to these results will
be an analysis of the dynamics of learning algorithms in various settings. The
primary contributions of this thesis are enumerated as follows:

  1. 1.

A characterization of the relationship between invariance, training speed, and
generalization, with theoretical results complemented by empirical validation
of our main findings in practically relevant settings.

    1. (a)

A theoretical analysis of the effect of invariance on generalization via PAC-
Bayes bounds, allowing an explicit characterization of the role of symmetries
in generalization bounds via a quantity we term the symmetrization gap. This
theoretical analysis is complemented by an empirical study which highlights
the limitations of the types of approximate invariance promoted by data
augmentation to generalize to out of distribution inputs.

    2. (b)

A new estimator of the marginal likelihood which complements the analysis of
upper bounds on generalization error to provide useful rankings of models for
architecture search and hyperparameter selection. The analysis of this
estimator reveals a deep connection between training speed and Bayesian model
selection, and yields a novel performance estimator for architecture search in
deep neural networks.

  2. 2.

A theoretical framework for the study of representation dynamics in deep
reinforcement learning along with practical insights derived thereof. These
insights concern both the ability of a learned feature representation to
linearly approximate a prediction target, and its ability to generalize to
novel prediction objectives over the course of training.

    1. (a)

A theoretical model and analysis of the representation learning dynamics of
value-based reinforcement learning algorithms, yielding an analytic
characterization of the effect of auxiliary tasks on agents’ learned
representations.

    2. (b)

The identification of the ‘capacity loss’ phenomenon in deep RL which
characterizes the tendency of neural networks to catastrophically overfit to
early prediction targets in sparse-reward environments.

    3. (c)

A new regularization method based on the insights from this previous analysis
that improves generalization to new prediction objectives even after long
training periods.

  3. 3.

Theoretical analysis leveraging the above framework to provide insight into
and algorithmic improvements to generalization to novel observations and
environments.

    1. (a)

A theoretically grounded explanation for prior empirical observations of
overfitting in the broader deep RL literature, including dense-reward
problems, and a set of recommendations for principled approaches to reduce
memorization and improve generalization between observations.

    2. (b)

A representation-learning objective which goes beyond the single-environment
setting to enable zero-shot generalization to novel environments sharing
underlying structure with the training environments.

####  1.4.2 Warmup: supervised learning

The first two content chapters of this thesis will lay the groundwork for our
understanding of generalization in deep neural networks. Chapter 3 presents a
novel PAC-Bayes generalization bound for invariant models, characterizing the
effect of invariance on generalization through a quantity which we term the
symmetrization gap. We empirically verify that the symmetrization gap appears
in computations of PAC-Bayes bounds for invariant deep neural networks, and
further show a strong correlation between the rankings of these upper bounds
and the rankings given by final generalization performance. This chapter
further contributes new empirical analysis of the optimization trajectories of
DNNs trained to exhibit approximate invariances via data augmentation,
illustrating the limitations of these trained invariances to extrapolate
beyond their training distribution. The importance of invariance to
generalization in reinforcement learning will be revisited in Chapter 8.

A failing of PAC-Bayesian generalization bounds is that they do not offer
predictions about which of a set of models will generalize best. In order to
achieve correctness, these bounds pay the price of predictive power. Chapter 4
shifts focus to consider performance estimators which can predict the relative
performance of different neural network architectures. It presents a novel
marginal likelihood estimator which can be used to enable Bayesian model
selection for a broader class of models from which we only require accurate
posterior samples, overcoming the normative limitations of generalization
bounds and computational challenges of the exact marginal likelihood. This
estimator has a number of appealing properties, chief among which is that it
can be applied to a subset of training losses from a gradient descent
trajectory. This illustrates a deep relationship between generalization and
learning dynamics, as a model’s generalization performance can thus be said to
depend on its training speed. These chapters will principally be based on the
following papers: *

  * •

lyle2020bayesian

  * •

lyle2020benefits

with supporting evidence drawn from

  * •

ru2020revisiting

####  1.4.3 Generalization in reinforcement learning

The remaining chapters present an analysis of the learning dynamics of RL
algorithms and explore a number of applications of this analysis to
representation learning and generalization in deep RL. Chapter 5 lays down a
novel learning dynamics framework which will be leveraged throughout the
remainder of the thesis, and reveals that the dynamics followed by temporal
difference methods can cause the learned representation to reflect aspects of
the transition structure of the environment (Lyle et al., 2021a). In Chapter
6, we explore one application of this analysis which yields a novel
regularization approach, InFeR (Lyle et al., 2022a). This method enables deep
RL agents to attain nontrivial return in the notoriously difficult Montezuma’s
Revenge game using only a naive $\epsilon$-greedy exploration algorithm. This
section is based on the following papers:

  * •

lyle2021effect

  * •

lyle2021understanding

Chapter 7 explores implications of the previous chapters on generalization,
providing novel analysis and insight into why value-based deep reinforcement
learning often produces highly brittle agents (Lyle et al., 2022b), as well as
identifying principled approaches to remedy the tendency of deep RL methods to
overfit. Chapter 8 revisits the discussion of invariance from Chapter 3 with a
new perspective grounded in causal inference, presenting a novel
representation learning method to encourage generalization in some classes of
MDPs by promoting invariance to spurious factors of variation in the
environment (Zhang et al., 2020a; Lyle et al., 2021b). These chapters are
based on the following papers.

  * •

lyle2022generalization

  * •

zhang2020invariant

A number of papers that I worked on during my PhD did not fit into this
thesis, including the following.

  * •

wang2021provable

  * •

filos2021psiphi

  * •

kossen2021self

  * •

bellemare2019geometric

  * •

lyle2019comparative

A discussion of contributions to joint work can be found at the end of the
thesis.

The key idea driving this thesis is that properties of a network’s training
trajectory can tell us a great deal about how it will generalize to novel
inputs. Chapters 3 and 4 ground this idea in the supervised learning regime by
studying generalization between minibatches and validate its utility by
developing both novel generalization bounds and practical model selection
tools. Chapters 5 and 6 identify key properties of the training dynamics of
reinforcement learning agents that differ from the supervised setting, and
show how these properties can be both beneficial and detrimental to
representation learning. Finally, Chapters 7 and 8 apply the notions of
invariance and within-training-set generalization from Chapters 3 and 4 to the
RL problem, leveraging the theoretical and empirical tools presented in
Chapters 5 and 6 to quantify memorization and improve generalization in
reinforcement learning. We will conclude with a discussion of how these ideas
can be (and in some cases have already been) further leveraged in the pursuit
of learning systems which can effectively learn and generalize across a range
of tasks.

##  Chapter 2 Background & literature review

This chapter will provide the high-level background and literature necessary
to contextualize the contributions of this thesis, and set out a standard set
of notation that will be used in the following chapters. Where necessary,
individual chapters may also contain a background section; these sections will
relay information that pertains only to the contents of the chapter that
contains them.

###  2.1 Learning frameworks

A learning algorithm can be applied to a wide variety of problems, and it is
often useful to categorize learning problems based on the information
available to the algorithm. We might task a learning system with identifying a
mapping between input-label pairs (supervised learning), with constructing an
embedding of inputs that captures relevant structure (unsupervised learning),
with generating samples from some distribution (generative modelling), or with
maximizing a reward signal via interaction with an environment (reinforcement
learning). The types of algorithms we can deploy in each of these situations
differ not only in terms of the types of outputs they produce, but also in the
stability of their learning dynamics. This thesis will focus primarily on the
distinction between supervised learning, where learning dynamics of gradient
descent algorithms are relatively straightforward to analyze, and
reinforcement learning, wherein analogues of even simple methods such as
linear regression can suffer from instability and divergence.

####  2.1.1 Supervised learning

Supervised learning is concerned with characterizing a functional relationship
between inputs and labels. This framework assumes that the data takes the form
of input-label pairs $(\mathbf{x},y)$ generated by sampling from some
distribution $P_{\mathcal{D}}$. The objective of the learning algorithm is to
find a function $f$ such that $f(\mathbf{x})=y$. The task of finding such a
function is nontrivial: one must both propose a suitable class of functions
over which to search, and an effective means of identifying functions from
this class which are likely to capture the target relationship.

##### Empirical risk minimization

We first consider the task of identifying a candidate function $f$ with the
property that $f(\mathbf{x})=y$ on $(\mathbf{x},y)$ pairs sampled from
$P_{\mathcal{D}}$, including those not seen during training. We use a loss
function $\ell$ to quantify the quality of $f(\mathbf{x})$ as an approximator
to $y$; we will also refer to the expectation of this quantity as the risk.
When $y$ belongs to a continuous space, we call this a regression problem.
When $y$ belongs to a finite set, we have a classification problem.

For a class of functions $F=\\{f:\mathcal{X}\to\mathcal{Y}\\}$, a set
$\mathcal{D}^{n}=(\mathbf{x}_{i},y_{i})_{i=1}^{n}\sim P_{\mathcal{D}}$, and a
loss function $\ell:\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}$, we define the
expected risk $R_{\ell}$ as

$\displaystyle R_{\ell}(f)$ $\displaystyle=\mathbb{E}_{(\mathbf{X},Y)\sim
P_{\mathcal{D}}}[\ell(f(\mathbf{X}),Y)]\;.$ (2.1) Similarly, we define the
empirical risk $\widehat{R}_{\ell}$ as
$\displaystyle\widehat{R}_{\ell}(f,\mathcal{D}^{n})$
$\displaystyle=\textstyle\frac{1}{n}\textstyle\sum_{i=1}^{n}\ell(f(\mathbf{x}_{i}),y_{i})\;.$
(2.2)

In regression problems, $\ell$ is typically set to be the squared error
$(f(\mathbf{x})-y)^{2}$. In classification, it is usually the cross-entropy
loss between a categorical distribution $p(\cdot|\mathbf{x})$ and the Dirac
delta distribution at the label $y$. In this case, the hypothesis class $F$
will consist of mappings from inputs to distributions over labels in
$\mathcal{Y}$.

When the hypothesis class $F$ contains the true functional relationship $f$,
the learning problem is realizable. However, most settings of interest are not
realizable, and so we seek instead a function $f^{*}$ which minimizes the
expected loss over the data-generating distribution, which can be expressed
formally as follows.

$f^{*}=\mathrm{argmin}_{f\in F}R_{\ell}(f)=\mathrm{argmin}_{f\in
F}\mathbb{E}_{(\mathbf{X},Y)\sim P_{\mathcal{D}}}[\ell(f(\mathbf{X}),Y)]$
(2.3)

In practice, computing this expectation is impossible and we instead use
samples to estimate its true value. The empirical risk minimization framework
assumes a finite sample $(\mathbf{x}_{i},y_{i})_{i=1}^{n}=\mathcal{D}^{n}\sim
P_{\mathcal{D}}$, and proposes to find a function $f\in F$ that minimizes the
empirical expectation of the loss over this sample, i.e. the empirical risk.
Concretely, the function $\hat{f}$ is called the empirical risk minimizer if
the following holds:

$\hat{f}=\mathrm{argmin}_{f\in F}\widehat{R}_{\ell}(f)=\mathrm{argmin}_{f\in
F}\frac{1}{n}\sum_{i=1}^{n}\ell(f(\mathbf{x}_{i}),y_{i})\;.$ (2.4)

The empirical risk minimization principle has seen widespread application in
the machine learning literature (Vapnik, 1991; Donini et al., 2018). Vapnik
(1968) characterizes a number of appealing asymptotic properties of the
empirical risk minimizer under certain conditions on the function class $F$;
this analysis has formed the basis for the work on generalization bounds which
will be discussed in Section 2.2.1. This principle is agnostic to the choice
of function class, and does not give guidance on how to obtain a minimizer
$f^{*}$ when such classes are too large for exhaustive search. The following
discussion will focus on one such class and search procedure: neural networks
trained with gradient-based optimization.

##### Deep learning

Deep neural networks (DNNs) form a powerful and expressive class of function
approximators (Raghu et al., 2017). A DNN is a parameterized function
$f_{\theta}$ which consists of layer-wise computations going from input to
output. Some neural architectures include recurrent connections, where the
output of the network is fed back into itself as input (Hochreiter and
Schmidhuber, 1997b); this thesis will focus exclusively on feedforward
architectures, where only a single pass through the network is executed in
order to obtain the function outputs. Feedforward neural networks constitute a
rich and widely-used class of models whose dynamics are more amenable to
analysis, including fully-connected multi-layer perceptrons (MLPs),
convolutional neural networks, transformers (Vaswani et al., 2017), and
ResNets (He et al., 2016). The function $f^{k}$ computed by each layer $k$ of
a feedforward neural network typically consists of a linear transformation of
the output of the previous layer, followed by a non-linear activation
function. A variety of activations have been used in DNNs; one popular example
is the Rectified Linear Unit (ReLU), of the form $\sigma(x)=\max(0,x)$. The
output of a neural network can thus be expressed as a composition of layer-
wise operations,

$f_{\theta}(\mathbf{x})=f_{\theta}^{L}\circ\dots\circ
f_{\theta}^{1}(\mathbf{x})\;.$ (2.5)

Deep neural networks are trained using gradient-based optimization algorithms.
The most fundamental of these is stochastic gradient descent. In this setting,
the data $\mathcal{D}^{n}$ is uniformly at random divided into minibatches of
size $k$,
$((\mathbf{x}_{b_{1}},y_{b_{1}}),\dots,(\mathbf{x}_{b_{k}},y_{b_{k}}))_{b=1}^{\lfloor
n/k\rfloor}$. At each minibatch, we compute the gradient of the loss to obtain
an update direction $g$ as follows,

$g(\theta,\mathcal{D}_{b})=-\nabla_{\theta}\frac{1}{k}\sum_{i=1}^{k}\ell(f_{\theta}(\mathbf{x}_{b_{i}}),y_{b_{i}})\;.$

This yields an iterative algorithm where the parameters $\theta$ are updated
according to the gradient for each successive minibatch. In most cases, we use
a step-size parameter $\alpha\in(0,1]$ to improve the stability and
convergence properties of the algorithm. The iteration step typically takes
the following general form:

$\theta_{t+1}\leftarrow\theta_{t}+\alpha_{t}g(\theta_{t},\mathcal{D}_{t})\;.$
(2.6)

Many formulations of gradient descent allow the step size to depend on the
iteration $t$; such dependence on the iteration is crucial to obtain
convergence guarantees (Robbins and Monro, 1951). Many adaptive optimization
schemes (Duchi et al., 2011; Kingma and Ba, 2015) further accumulate
parameter-dependent learning rates, and allow the update direction
$g(\theta_{t},\mathcal{D}_{t})$ to also depend on prior gradients. While these
optimizers will feature in the empirical analysis that appears later, their
precise form is not important to our discussion.

####  2.1.2 Reinforcement learning

Whereas the supervised learning framework seeks to model a functional
relationship between two variables, the reinforcement learning framework
models an agent’s interaction with an environment with the goal of identifying
a behaviour policy which maximizes some reward signal. We model the
environment as a Markov Decision Process (MDP)
$\mathcal{M}=(\mathcal{X},\mathcal{A},R,P,\gamma)$, where $\mathcal{X}$
denotes the state space, $\mathcal{A}$ the action space,
$R:\mathcal{X}\rightarrow\mathbb{R}$ the reward function,
$P:\mathcal{X}\times\mathcal{A}\rightarrow\mathscr{P}(\mathcal{X})$ the
transition probability function, and $\gamma$ the discount factor. The agent
obtains observation $x$ which indicates the environment’s state. It may then
take an action $a\in\mathcal{A}$, after which the environment outputs a new
observation $x^{\prime}$ and a reward $r$. The agent’s objective is to
maximize the cumulative discounted reward it receives over time. In finite-
horizon environments, the agent may only take a finite number of steps; in
continuing or infinite-horizon environments, on which this thesis will
predominantly focus, the agent may take an unlimited number of steps in the
environment. The discount factor $\gamma$ determines the degree to which near-
term rewards are preferred, resulting in the maximization target
$\sum_{k=0}^{\infty}\gamma^{k}R_{k}$, called the return. The return is a
random variable that depends on the sequence of states and actions taken by
the agent. The value of a state-action pair under some action-selection policy
$\pi:\mathcal{X}\rightarrow\mathscr{P}(\mathcal{A})$ is equal to the expected
value of the return starting at some state-action pair $(x,a)$ and following
the policy $\pi$. This can be expressed by the action-value function
$Q^{\pi}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$, defined as

$Q^{\pi}(x,a)=\mathbb{E}_{\pi,\mathcal{P}}[\sum_{k=0}^{\infty}\gamma^{k}R(x_{k},a_{k})|x_{0}=x,a_{0}=a]\;.$
(2.7)

Value-based methods seek to learn the (resp. action-) value function
$V^{\pi}:\mathcal{X}\rightarrow\mathbb{R}$ (resp.
$Q^{\pi}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$) associated with
some policy $\pi$ (Sutton and Barto, 2018). In particular, we are interested
in learning the value function associated with the optimal policy $\pi^{*}$
which maximizes the expected discounted sum of rewards from any state. Such a
value function is then straightforward to translate into an optimal behaviour
policy: the agent need only take the action with the highest predicted value
at each state.

At the core of value-based RL are the _Bellman operators_ (Puterman, 1994).
The Bellman policy evaluation operator
$T^{\pi}:\mathbb{R}^{\mathcal{X}}\rightarrow\mathbb{R}^{\mathcal{X}}$ is
defined with respect to a policy $\pi$ and provides a method to update a
predicted value function $V$ to more closely resemble the value $V^{\pi}$ of
the policy $\pi$. It is defined as

$\displaystyle(T^{\pi}V)(x)=\mathbb{E}_{X_{1}\sim
P(\cdot|x,\pi(x))}[R(x)+\gamma V(X_{1})]\,.$

Introducing a matrix notation of the transition operator
$P^{\pi}\in\mathbb{R}^{\mathcal{X}\times\mathcal{X}}$ defined by
$P^{\pi}[x,x^{\prime}]=\sum_{a\in\mathcal{A}}\pi(a|x)P(x^{\prime}|x,a)$, and
the expected reward vector $R^{\pi}\in\mathbb{R}^{\mathcal{X}}$ defined by
$R^{\pi}(x)=\mathbb{E}_{\pi}[R_{0}|X_{0}=x]$, this can be expressed even more
succinctly in matrix notation as

$\displaystyle T^{\pi}V=R^{\pi}+\gamma P^{\pi}V\,.$

$T^{\pi}$ is a contraction on the space of value functions (Puterman, 1994),
and so repeated application of $T^{\pi}$ to any initial value function
converges to $V^{\pi}$ (Bertsekas and Tsitsiklis, 1996). For control problems,
where we seek to obtain the value of an unknown optimal policy, learning
requires not only estimating the value of a policy but also improving that
policy to increase its expected return. In this setting, we leverage an
analogous operator termed the Bellman optimality operator
$T^{*}:\mathbb{R}^{\mathcal{X}\times\mathcal{A}}\rightarrow\mathbb{R}^{\mathcal{X}\times\mathcal{A}}$.
The action of $T^{*}$ on value functions is defined by

$\displaystyle(T^{*}Q)(x,a)\\!=\\!\mathbb{E}_{x^{\prime}\sim
P(x,a)}[R(x)\\!+\\!\gamma\max_{a^{\prime}\in\mathcal{A}}Q(x^{\prime},a^{\prime})]\,.$

The Bellman optimality operator attains similar convergence guarantees as the
policy evaluation operator, and can be shown in tabular state spaces to
converge to the value of the optimal policy. In principle both $T^{*}$ and
$T^{\pi}$ can be defined over action-value or state-value functions, but this
thesis will predominantly consider the application of $T^{\pi}$ to value
functions and $T^{*}$ to action-value functions We will refer to the value
$T^{\pi}V$ or $T^{*}Q$ as the Bellman target associated with the (resp.
action-) value function $V$ (resp. Q), where the operator in use will be clear
from context.

In most settings of interest, we do not have access to the expected reward
vector $R^{\pi}$ or the environment transition matrix $P$. As a result, value-
based RL agents must use sampled transitions of the form
($x_{t},a_{t},r_{t},x_{t+1},a_{t+1}$) to approximate these updates. In the
case of the policy evaluation operator, this sample-based approximation takes
the form of the SARSA update, so called due to its use of (State-Action-
Reward-State-Action) transitions. The SARSA algorithm assumes that the
transition has been sampled from some fixed behaviour policy $\pi$, and
estimates $Q^{\pi}$ by iteratively applying the update rule

$Q_{t+1}(x_{t},a_{t})=Q_{t}(x_{t},a_{t})+\alpha[r_{t}+\gamma
Q_{t}(x_{t+1},a_{t+1})-Q_{t}(x_{t},a_{t})]\;.$ (2.8)

Sample-based methods will use some step size $0<\alpha<1$ in order to average
out noise in the Bellman targets due to stochasticity in the environment. The
seminal Q-learning algorithm (Watkins and Dayan, 1992), which forms the basis
of many deep RL agents (Mnih et al., 2015), can similarly be viewed as
approximating the iterative application of $T^{*}$ and related operators
(Tsitsiklis, 1994; Jaakkola et al., 1994; Bertsekas and Tsitsiklis, 1996).
Q-learning is based on the following update

$Q_{t+1}(x_{t},a_{t})=Q_{t}(x_{t},a_{t})+\alpha[r_{t}+\gamma\max_{a^{\prime}\in\mathcal{A}}Q_{t}(x_{t+1},a^{\prime})-Q_{t}(x_{t},a_{t})]\;.$
(2.9)

Policy gradient methods (Sutton et al., 2000) operate directly on a
parameterized policy $\pi_{\theta}$. We let $d^{\pi}$ denote the stationary
distribution induced by a policy $\pi$ over states in the MDP. Rather than
first going to the trouble of learning a value function, policy gradient
methods directly optimize the parameters $\theta$ of the policy so as to
maximize the expected return $J(\pi_{\theta})=\mathbb{E}_{s_{0}\sim
P_{\mathcal{M}}(s_{0})}V^{\pi_{\theta}}(s_{0})$. The gradient of this loss can
be estimated from sampled trajectories when expressed as follows:

$\nabla_{\theta}J(\pi_{\theta})=\mathbb{E}_{x_{t},a_{t}\sim
P(\cdot|\pi_{\theta})}[\nabla_{\theta}\log\pi_{\theta}(a_{t}|x_{t})Q^{\pi_{\theta}}(x_{t},a_{t})]\;.$
(2.10)

Variations on this learning rule include actor-critic methods (Konda and
Tsitsiklis, 2000), which use a baseline given by a value-based learner to
reduce update variance, and trust-region based methods, such as Trust Region
Policy Optimization (Schulman et al., 2015) and Proximal Policy Optimization
(PPO) (Schulman et al., 2017).

Function approximation schemes enable RL agents to generalize their knowledge
about the value or policy at one state to other states in the environment.
Linear function approximation assumes state-action pairs are embedded as
features $\phi(x,a)\in\mathbb{R}^{d}$, and some linear map
$\mathbf{w}\in\mathbb{R}^{d}$ is used to approximate the value function
$Q^{\pi}(x,a)=\langle\phi(x,a),\mathbf{w}\rangle$. This regime has been the
study of a rich literature exploring the stability and sample complexity of RL
in the presence of function approximation (Jin et al., 2020; Wang et al.,
2020; Precup et al., 2001; Tsitsiklis and Van Roy, 1996).

A second, more widely-used family of algorithms involve the use of neural
networks as function approximators of the form
$Q_{\theta}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$. This is the
deep RL regime. This approach is well-suited to an array of complex tasks,
ranging from systems control problems (Degrave et al., 2022) to video games
(Mnih et al., 2015). At its core, value-based deep RL involves training a
neural network to approximate the Bellman targets of a predicted value
function using (semi-, see e.g. (Sutton and Barto, 2018)) gradient descent.
This approach computes a semi-gradient update direction $f(\theta)$ of the
following form, where $a^{*}=\max_{a}(Q_{\theta}(x_{t+1},a))$.

$f(\theta)=(\nabla_{\theta}Q_{\theta})[r_{t}+\gamma
Q_{\theta}(x_{t+1},a^{*})-Q_{\theta}(x_{t},a_{t})]$ (2.11)

Because the final layer of a neural network is usually linear, it is possible
to express the output $Q_{\theta}(x,a)$ in the form
$Q_{\theta}(x,a)=\langle\phi_{\theta^{\prime}}(x,a),\mathbf{w}\rangle$, and
$\theta=\theta^{\prime}\oplus\mathbf{w}$ where $\oplus$ denotes concatenation.
Under this parameterization, the map $\phi_{\theta^{\prime}}$ is referred to
as the feature map. This framework has been used to study the learned
representations of deep reinforcement learning agents in a number of recent
works (Kumar et al., 2021; Lan et al., 2022; Lyle et al., 2019a).

###  2.2 Generalization in supervised learning

We now turn our attention to the principal object of interest in this thesis:
generalization. The study of generalization in supervised learning problems
spans decades, from the seminal work of Vapnik (1968) to recent exciting
developments in the kernel analysis of deep networks (Jacot et al., 2018) and
the interpolation regime (Bartlett et al., 2020). We will begin by presenting
classical bounds on the generalization error of a learning algorithm’s output.
However, these classical approaches, based on quantifying the complexity of an
algorithm’s hypothesis class, fail to account for the generalization
performance of deep neural networks, motivating more recent empirical
approaches. We will conclude with an overview of the current state of the art
of our understanding of generalization in deep learning.

####  2.2.1 Generalization bounds

The empirical risk minimization (ERM) framework puts forward the maxim:
‘always pick a hypothesis which minimizes the risk on the training set’. This
hypothesis will not in general attain the lowest risk on the underlying data-
generating distribution, however. A long line of work (Vapnik, 1968, 1999;
Bousquet and Elisseeff, 2002) has characterized upper bounds on the gap
between the empirical risk of a hypothesis and its true risk, yielding
generalization bounds which provide high-probability guarantees on the
expected risk of a hypothesis. Of particular interest to us will be the
application of such bounds to hypothesis classes generated by neural networks
(Baum and Haussler, 1989; Dziugaite and Roy, 2017; Bartlett et al., 2017).
While generalization bounds for any hypothesis class are almost always looser
than the upper bound on the expected risk given by computing the model’s
validation loss on held-out data, the study of generalization bounds continues
to thrive as a means of developing theoretical insight into learning
algorithms, motivating their inclusion in our discussion.

##### Formalism

Most generalization bounds in the literature share a similar structure,
consisting of the sum of a hypothesis’ empirical risk and a complexity measure
scaled by a function (typically the inverse square root) of the number of
samples. We use the notation of Section 2.1.1, and let $f$ be a function
output by some learning algorithm with access to data $\mathcal{D}^{n}$ of
size $n$, drawn from hypothesis class $F$. We let $R_{\ell}$ and
$\widehat{R}_{\ell}$ be defined as in Equations 2.1 and 2.2 respectively.
Letting $C(F)$ denote some complexity measure (for example, the logarithm of
the number of functions in the hypothesis class) and $g(\cdot)$ some non-
negative function defined over $\mathbb{R}$, usually $g(x)=\log(\frac{1}{x})$
or something similar, we obtain the generic form

$R_{\ell}(f)\leq\widehat{R}_{\ell}(f)+\sqrt{\frac{C(F)+g(\delta)}{n}}\text{
with probability $1-\delta$}.$ (2.12)

The magnitude of the complexity measure $C(F)$ is crucial to the tightness of
a bound. In some cases the complexity measure which bounds the expected risk
may take a value so large that it dwarfs the maximal value the risk can
obtain, resulting in bounds that are vacuous. For example, a vacuous bound
would guarantee that a neural network’s probability of making a classification
error on new samples from the the MNIST dataset will be less than 500%. In
neural networks, bounds based on the margin around the decision boundary (Wei
and Ma, 2019), the VC dimension (Harvey et al., 2017), and the spectral norm
of the network weights (Bartlett et al., 2017) all become vacuous for network
architectures of the scale typically applied to popular benchmarks such as
CIFAR-10 or ImageNet (Bartlett, 1997; Dziugaite and Roy, 2017). However,
recent work has found complexity measures which can capture a reasonable
notion of simplicity on neural networks – at least to the point where bounds
are non-vacuous (Dziugaite and Roy, 2017).

##### Complexity and Occam’s razor

The term $C(F)$ can be interpreted as a form of Occam’s razor applied to the
hypothesis class: given hypotheses drawn from two classes which attain the
same empirical risk, we should prefer the hypothesis drawn from the simpler
class. This simple idea drives most results in model selection and
generalization in machine learning (Rasmussen and Ghahramani, 2001). However,
while generalization bounds of the flavour shown above characterize the
complexity of the entire class of hypotheses, one might also prefer to use
Occam’s razor as a tool to select hypotheses within a single class. This is
the core idea behind PAC-Bayes generalization bounds (McAllester, 1999;
Langford and Shawe-Taylor, 2003; Lever et al., 2013; Catoni, 2007), which
allow us to assign a complexity penalty that distinguishes between different
hypotheses within a function class for randomized predictors. Leveraging a
notion of hypothesis-level complexity is a powerful tool (Bartlett, 1997),
yielding some of the first non-vacuous generalization bounds for over-
parameterized neural networks (Dziugaite and Roy, 2017).

Occam’s razor does not give us an out-of-the-box definition of simplicity,
however. Two such definitions are favoured by the machine learning community
to explain generalization of deep networks. The first of these is flatness of
the local minima to which gradient-based optimization tends to converge
(Hochreiter and Schmidhuber, 1997a). Much of the folk wisdom surrounding
generalization in deep learning hypothesizes that the reason neural networks
generalize is because stochastic gradient descent drives the learned weights
towards flat regions of the loss landscape (Keskar et al., 2017). This picture
is slightly complicated by the observation of Dinh et al. (2017) showing that
sharp minima can also generalize, but nonetheless demonstrates strong
predictive power in empirical evaluations (Jiang et al., 2020). Optimization
algorithms which deliberately add noise to the training process, such as
entropy-SGD and the direct optimization of a PAC-Bayes bound (Chaudhari et
al., 2017; Dziugaite and Roy, 2018a, b), have been shown to increase flatness
and improve the tightness of some generalization bounds on neural networks.
PAC-Bayes risk bounds are also deeply connected to the Bayesian notion of
marginal likelihood (Germain et al., 2016a).

A second widely-used notion of simplicity is model compressibility, whereby
models which can be compressed to a smaller length are said to be simpler.
This is also referred to as the minimum description length (MDL) principle
(Akaike, 1998; Hinton and Van Camp, 1993). Compressibility and flatness
overlap significantly: parameters in a flat region of the loss landscape can
be perturbed, for example by a quantization step of a compression algorithm,
without significantly affecting the loss. However, the techniques used to
measure the two quantities differ, with compression approaches typically being
more compute-intensive (Zhou et al., 2019; Ullrich et al., 2017). PAC-Bayes
bounds are appealing as either the flatness of minima (Neyshabur et al., 2019,
2017) or the compressibility (Zhou et al., 2019) notion of simplicity can be
used to define the model complexity term in the bound.

##### The overparameterized regime

At the core of the challenge of applying results from statistical learning
theory to neural networks is the expressiveness of neural network function
classes; the overparameterized models of recent years are capable of
memorizing even uniform random labels of the data (Zhang et al., 2017).
Uniform convergence bounds which depend on worst-case analysis are thus
limited in the guarantees they can provide these networks. Indeed, Nagarajan
and Kolter (2019) argue that certain forms of uniform convergence results may
be fundamentally incapable of explaining generalization in deep learning,
presenting a simple example of a learning problem where it is impossible to
use uniform convergence guarantees to characterize the generalization
performance of a neural network function class. While Negrea et al. (2020) and
Bartlett et al. (2020) show that modified analysis can still yield uniform
convergence results for overparameterized predictors in similar contexts,
these results require studying either a surrogate predictor, or a restricted
problem setting.

Even more intriguing has been the observation that, contrary to the received
wisdom in learning theory that overparameterized models will overfit to their
training data, increasing overparameterization can lead to improved
generalization (Neyshabur et al., 2015a). These empirical observations have
prompted the study of interpolating predictors (Bartlett et al., 2020), which
seeks to outline conditions by which a model which can attain zero empirical
risk may still obtain near-optimal risk or robustness properties on the data-
generating distribution (Bubeck and Sellke, 2021; Koehler et al., 2021). This
work is intricately connected to the double descent phenomenon, whereby the
risk of a model, when plotted against the number of parameters, exhibits two
descent regions: one in the underparameterized regime, and one in the
overparameterized regime (Advani et al., 2020; Belkin et al., 2019; Nakkiran
et al., 2021).

A gap remains, however, in leveraging these insights to attain tight bounds on
the generalization error of modern neural network architectures trained on
real-world datasets. Even worse, Jiang et al. (2020) show that many of the
complexity measures that appear in generalization bounds are negatively
correlated with generalization. This presents a double blow to the argument
that generalization bounds might provide insight into what makes neural
networks generalize well: not only are such bounds too loose to give
practically relevant information, but they do not even offer an accurate
ranking of models.

####  2.2.2 Explaining generalization without bounds

Why does a given neural network generalize well or poorly? This is a question
that generalization bounds cannot (currently) answer, yet is crucial to the
principled development and application of neural networks to real-world
datasets. A large value of a complexity measure does not entail that a
learning algorithm will not generalize; it simply states that there is not
sufficient information to guarantee a small generalization gap. To address
questions of why particular models generalize well, we must look to the
particulars of the network architectures, training procedures, and datasets
that arise from the natural world. The works discussed in this section will
present such an empirical study. While sometimes similar notions of complexity
to those leveraged in generalization bounds may be used, the philosophical
underpinnings of this research differ fundamentally from the formal guarantees
of learning theory.

##### Scientific vs mathematical rigor

The scientific study of any phenomenon depends on a mixture of empirical
experiments, during which data is gathered, and theory-building, where an
explanation of the phenomenon is proposed (Popper, 1968). A scientific theory
is one that makes falsifiable predictions, which can then be tested by
experiments. The generalization bounds of Section 2.2.1 can be viewed as
theory-building in a loose sense; however, the goal of a generalization bound
is to make a statement that will be guaranteed to hold with high probability,
accepting that this may result in a large gap between the predicted upper
bound and the empirical realization of the generalization error in some
settings. This runs counter to the qualities of a good scientific theory,
whose goal is to explain a phenomenon with as simple a model as possible, and
to expose itself to the risk of falsification in doing so.

One notable step towards this framework of formulating and falsifying
hypotheses arises from the work of Jiang et al. (2020) and Dziugaite et al.
(2020), who conduct a large-scale empirical study of the correlation between
complexity measures and generalization in deep neural networks. This line of
work seeks to translate the formal guarantees of generalization bounds into
testable scientific theories which make predictions about which models will
generalize best. The key to doing so is to treat the ranking over models given
by a complexity measure as a prediction, rather than an upper bound. If the
generalization measure $C$ is higher for network A than for network B, we
interpret this as a prediction that network A should generalize worse than
network B stemming from the theory that a model which generalizes well must do
so because $C$ is low.

Dziugaite et al. (2020) focus on identifying experimental settings where a
measure fails to predict generalization, seeking to identify instances where
the theory can be falsified. Both works show that complexity measures based on
flatness, such as PAC-Bayes generalization bounds (Dziugaite and Roy, 2017),
are highly predictive of generalization across a range of experimental
settings. However, none of the generalization measures studied by Dziugaite et
al. (2020) robustly predicts generalization in neural networks under all
experimental conditions, motivating further investigation of quantites which
can predict – and ideally also explain – generalization in neural networks.

##### The loss landscape of neural networks

Much theoretical work studying neural networks has considered questions on the
existence of parameters that allow the network to represent a function (Hornik
et al., 1989; Raghu et al., 2017; Dong et al., 2020), but less attention has
been paid to the search process by which such parameters might be found. To
this end, recent work studying the loss landscape of neural network seeks to
understand the properties of the optimization problem faced by neural networks
in practice. These results have found that the initialization schemes used
today (He et al., 2015; Glorot and Bengio, 2010) push the parameters towards a
well-behaved region of the loss landscape that is relatively convex (Fort and
Scherlis, 2019) and induces stable learning dynamics in suitable architectures
(Yang and Schoenholz, 2017). Indeed, many architecture design choices such as
residual connections (He et al., 2016) and batch normalization (Ioffe and
Szegedy, 2015) that improve performance can be shown to increase the
smoothness of the loss landscape (Li et al., 2018; Santurkar et al., 2018),
leading to minima which are flatter and should thus generalize better.

Much recent work has focused in particular on the geometry of the loss
landscape around the minima found by gradient descent (Maddox et al., 2020).
One intriguing property identified recently is that of linear mode
connectivity, whereby the convex combination of two locally optimal parameter
vectors attains a similarly low loss as the two minima from which it was
derived (Benton et al., 2021). Frankle et al. (2020a) show that stochasticity
in the optimization process early in training leads the trajectory towards
disconnected regions of parameter space, while stochasticity late in training
results in a perturbation within a linearly connected loss basin. This
geometric perspective is consistent with prior work studying the importance of
the early learning period in determining the network’s ability to represent
certain functions (Achille et al., 2018). Flatness properties of local minima
have also been studied in the context of Bayesian deep learning (Izmailov et
al., 2018), highlighting the connection between flat minima and approximations
of the Bayesian marginal likelihood (Daxberger et al., 2021).

##### Correlates of generalization in deep learning

Much empirical work on generalization in deep learning seeks to identify
measurable quantities which correlate with generalization in neural networks
(Neyshabur et al., 2017). This work has found such quantities in the
robustness of minima to perturbations (Keskar et al., 2017) and to pruning
(Bartoldson et al., 2020). Relatedly, Morcos et al. (2018) identify a
relationship between the dependence of a network on single dimensions of
activation space and generalization. Neyshabur et al. (2015b) propose that
implicit norm-based regularization may contribute to the generalization
performance of deep networks; some evidence suggests, however, that the
mechanism by which overparameterization regularizes norm may not apply to
fully general problem constructions (Hanin and Rolnick, 2019).

Further attention has been paid to regularization methods which improve
generalization such as dropout (Srivastava et al., 2014), data augmentation
(Wu et al., 2020), batch normalization (Ioffe and Szegedy, 2015), and weight
decay (Arpit et al., 2017). The study of early stopping (Duvenaud et al.,
2016) has further revealed intriguing connections between variational
inference and flatness of the loss landscape. However, the precise
relationship between training speed, early stopping, and generalization has
yet to be fully understood: Hardt et al. (2016) provide a theoretical and
empirical analysis of the relationship between training time and
generalization error; in contrast, Hoffer et al. (2017) observe that longer
training can improve regularization provided a suitable optimization procedure
is used. This is emblematic of a broader lack of granular understanding of
generalization in deep neural networks. While many correlates of
generalization and techniques to improve generalization are widely used, the
precise mechanisms by which they improve generalization have yet to be
elucidated.

##### Studying the learning trajectory

The study of how networks evolve over the course of training has revealed a
number of intriguing insights which promise to shed light on some of these
mechanisms. A number of works have endeavoured to explain the efficacy of
early stopping by arguing that neural networks learn functions of increasing
complexity over the course of training (Rahaman et al., 2019; Arpit et al.,
2017; Kalimeris et al., 2019), with similar biases towards other notions of
simplicity also widely observed in the mapping between parameters and the
resulting output functions of neural networks (Valle-Pérez et al., 2018; De
Palma et al., 2019).

Approximating a discrete gradient descent trajectory as a continuous-time
process in the limit of infinite layer width has facilitated closed-form
analysis of the trajectory of deep neural networks (Jacot et al., 2018; Lee et
al., 2019), which has sparked a number of intriguing insights into
generalization (Arora et al., 2019), posterior sampling (He et al., 2020), and
feature learning (Yang and Hu, 2021). The analysis of Jacot et al. (2018)
allows the optimization trajectory of a neural network to be modelled by a
kernel gradient descent procedure with respect to a specific kernel called the
neural tangent kernel (NTK). Smith et al. (2020) and Barrett and Dherin (2021)
analyze a similar continuous-time approximation of gradient descent to
identify an explicit regularization term that biases gradient descent
algorithms towards flatter regions of the parameter space, providing insight
into the relationship between the flatness of the region of the loss landscape
traversed by gradient descent and sources of stochasticity such as minibatch
sizes and finite learning rates.

####  2.2.3 Summary

As this section has shown, the question of why neural networks generalize has
sparked a wide-ranging literature characterizing generalization and
optimization dynamics in DNNs. In spite of this rich literature, the answer to
this question remains open. Generalization bounds present a domain of great
theoretical interest, but fail to explain generalization in DNNs. Meanwhile,
empirical work has identified many quantities corresponding to model
complexity terms in generalization bounds that correlate with generalization
(Jiang et al., 2020). However, no single quantity studied thus far has been
found to predict generalization performance under all possible experimental
conditions (Dziugaite et al., 2020). Complementary lines of work have shed
significant insight into the optimization landscape and generalization
properties of DNNs, but this picture is far from complete. A number of recent
theoretical results, such as the analytic form of optimization dynamics in the
infinite-width limit of DNNs (Jacot et al., 2018), the double descent
phenomenon (Belkin et al., 2019), and benign overfitting (Bartlett et al.,
2020), hint at a theoretical explanation of the benefits of
overparameterization for generalization, but these findings do not currently
apply to modern deep learning training regimes. Similarly, empirical
investigation into the loss landscapes and gradient structure in DNNs has
yielded principled insights into the optimization dynamics of deep networks,
but still admits a gap between theory and practice. One contribution of this
thesis will be to propose a family of performance estimators based on a
network’s training speed which are predictive of generalization performance
rankings in a number of experimental settings. These estimators will be
motivated by a theoretical and empirical analysis linking generalization
bounds, invariance, and training speed, providing a potential mechanism to
explain their success.

###  2.3 Generalization in reinforcement learning

The supervised learning regime has proven to be fertile ground for the study
of generalization. However, reinforcement learning agents also benefit from
accurate extrapolation to unseen inputs, and moreover existing approaches are
particularly prone to overfitting (Zhang et al., 2018c). Though many of the
subproblems of generalization in reinforcement learning mirror those found in
supervised learning (Cobbe et al., 2019), the nonstationarity of the training
objectives used in deep RL adds additional challenges (Igl et al., 2021) which
we will discuss in this section.

####  2.3.1 Single environment

Generalization is not always necessary for reinforcement learning: in small
state spaces, for example, one may assume that all states the agent will
encounter are known in advance and that the value function can be represented
as a lookup table; the scope for generalization in this setting is limited to
the context of learning an optimal policy under new reward functions (Dayan,
1993). However, in large spaces generalization is crucial for sample
efficiency, enabling faster convergence to an optimal policy and, in cases
where the state space is so large that not all states will be visited during
training, ensuring robust performance in previously-unseen states at
evaluation time. This thesis will explore two different perspectives on
generalization in a single environment: state abstractions, which can be
leveraged to develop theoretical insights into the complexity of certain
classes of environments, and interference, which studies how deep neural
networks disentangle and distinguish between states.

Figure 2.1: State abstractions allow an agent to treat similar states as
though they were equivalent, potentially simplifying planning and
generalization. Here, we see an example of grouping together states with
similar transition dynamics, reducing the size of the state space from 15
states to 3.

##### State abstractions

State abstractions are used in reinforcement learning to simplify planning and
to enable generalization in large or infinite state spaces; see Figure 2.1 for
an illustration. Mathematically, a state abstraction is a mapping
$\phi:\mathcal{X}\rightarrow\bar{\mathcal{X}}$, where $\bar{\mathcal{X}}$ is
some simplified space (for example, $\mathbb{R}^{d}$). A state abstraction is
useful for reinforcement learning if it captures the similarity between states
in the environment. That is, $\phi(s)=\phi(s^{\prime})$ if and only if $s$ and
$s^{\prime}$ have similar values or optimal policies. Li et al. (2006)
characterize a hierarchy of such abstractions, with the finest being a
bisimulation of the original MDP and the coarsest being a partition of states
according to the optimal action at each state. We will be particularly
interested in the former, as our study of representation learning in Chapter 8
will focus on this class of state abstractions, known as model-irrelevance
state abstractions. The following definition is for finite state spaces, but
can be generalized to continuous state spaces.

######  Definition 2.3.1.

A _model-irrelevance state abstraction_ on a Markov Decision Process
$\mathcal{M}=(\mathcal{X},\mathcal{A},R,P,\gamma)$ is a mapping
$\phi:\mathcal{X}\rightarrow\bar{\mathcal{X}}$ such that if
$\phi(x_{1})=\phi(x_{2})=\bar{x}$, the following holds, where
$x^{\prime}\in\mathcal{X}$, and $\bar{x^{\prime}}\in\bar{\mathcal{X}}$.

  1. 1.

$R(x_{1},a)=R(x_{2},a)$ $\forall a\in\mathcal{A}$

  2. 2.

$\sum_{x^{\prime}\in\phi^{-1}(\bar{x}^{\prime})}P(x^{\prime}|x_{1},a)=\sum_{x^{\prime}\in\phi^{-1}(\bar{x}^{\prime})}P(x^{\prime}|x_{2},a)$,
$\forall a\in\mathcal{A}$

There are a range of approaches to discovering and selecting state
abstractions: Jong and Stone (2005) eliminate irrelevant variables from inputs
where $\mathcal{X}\subset\mathbb{R}^{n}$; other approaches leverage
statistical (Jiang et al., 2015) and information-theoretic (Abel et al., 2019)
tools to aggregate states, while Taylor et al. (2008) construct state
abstractions based on state similarity in the MDP. State aggregation is
further leveraged in algorithms for RL in rich-observation settings (Misra et
al., 2020).

In deep reinforcement learning, the state abstraction framework begins to
merge with the feature-learning and representation-learning frameworks of deep
learning. One can interpret the latent feature representation of the neural
network as a state abstraction, and the remaining layers as a value function
approximator. Gelada et al. (2019) and Zhang et al. (2020b) use ideas from the
state abstraction literature to construct representation-learning objectives
for DNNs which allow them to capture the structure of the environment. Such
structure can be made explicit in the architecture of the neural network used
as a function approximator to obtain guarantees on the equivariance of the
learned value function to symmetries that arise in the MDP (van der Pol et
al., 2020b, a).

##### Representations in deep RL

A major challenge present in deep reinforcement learning is the sparsity of
the reward signal relative to the rich structure of the environment. A broad
range of auxiliary tasks have been developed over the years to overcome this
sparsity in order to improve generalization and representation learning in
deep RL (Jaderberg et al., 2017; Veeriah et al., 2019; Gelada et al., 2019;
Machado et al., 2018b). These contributions have in common the goal of
enriching the scalar reward signal to encourage the network to pick up
additional structure in the environment, either with the downstream objective
of using this information later for policy improvement, or as a more stable
learning signal for the network. Additional work has analyzed the geometry
(Bellemare et al., 2019) and stability (Ghosh and Bellemare, 2020) of the
learned features in deep RL agents, along with their linear algebraic
properties (Kumar et al., 2021; Gogianu et al., 2021).

A second challenge in the application of deep neural networks to RL arises
from the nonstationarity of the learning objective. The challenge of training
a network on a sequence of tasks has been studied in great depth in both
reinforcement learning (Schaul et al., 2019; Teh et al., 2017; Igl et al.,
2021) and supervised learning settings (Sharkey and Sharkey, 1995; Ash and
Adams, 2020; Beck et al., 2021). Of particular interest has been the problem
of catastrophic forgetting, with prior work proposing novel training
algorithms using regularization (Kirkpatrick et al., 2017; Bengio et al.,
2014; Lopez-Paz and Ranzato, 2017) or distillation (Schwarz et al., 2018;
Silver and Mercer, 2002; Li and Hoiem, 2017) approaches. Benjamin et al.
(2019), apply a function-space regularization approach, but require saving
input-output pairs into a memory bank. However, in reinforcement learning the
converse problem must also be considered: forward-interference, whereby
networks may see deteriorating ability to solve later tasks encountered during
training. Methods which involve re-initializing a new network have seen
particular success at reducing interference between tasks in deep
reinforcement learning (Igl et al., 2021; Teh et al., 2017; Rusu et al., 2016;
Fedus et al., 2019b).

##### Interference

Even in the absence of an explicit abstraction-learning objective,
generalization between states in deep RL will naturally arise from the
network’s inductive bias. In many problem settings, such as in large
observation spaces or procedurally generated environments, some degree of
generalization is necessary in order to obtain good performance at test time
(Kirk et al., 2021). However, when an update to the predicted value of one
state exerts excessive influence on the network’s output at other states, this
form of generalization (which we will also refer to as interference) can lead
to instability (Jiang et al., 2021b; Van Hasselt et al., 2018). The
introduction of instability as a result of interference between states is not
unique to deep RL; even in the case of linear function approximation, off-
policy temporal difference learning is not guaranteed to converge to the
optimal value function (Baird, 1993; Tsitsiklis and Van Roy, 1996).
Interference has been identified as a barrier to learning progress in both
policy-based and value-based problems (Schaul et al., 2019; Fedus et al.,
2019b).

There are two principal approaches by which we will quantify this notion of
interference. To a first-order approximation, interference is equivalent to
gradient alignment. To formalize this concept, we let $f$ be a function which
takes as input an observation $\mathbf{x}$ and a set of parameters $\theta$
(for example, $f$ might be the loss function of a neural network, or one
dimension of its output). We define gradient alignment
$\mathrm{I}_{\nabla}(\mathbf{x},\mathbf{y};\theta)$ (where $\nabla$ indicates
that we are studying the gradient of the loss) with respect to some loss
function and network architecture parameterized by $\theta$ and taking inputs
$\mathbf{x}$ and $\mathbf{y}$ as follows.

$\mathrm{I}_{\nabla}(\mathbf{x},\mathbf{y};\theta)\overset{\text{def}}{=}\langle\nabla_{\theta}f(\mathbf{x};\theta),\nabla_{\theta}f(\mathbf{y};\theta)\rangle$
(2.13)

We may instead measure the effect of a gradient step computed on data point
$(\mathbf{x}_{1},y_{1})$ on the loss of another data point
$\mathbf{x}_{2}\neq\mathbf{x}_{1}$ with corresponding target $y_{2}$. Letting
the targets $y$ be left implicit we obtain the following, where $f$ is some
function as before and $\ell$ is the optimization objective (which may or may
not be equal to $f$). We let $g_{\ell}$ be the update computed by some
optimizer which may or may not be equal to the gradient
$\nabla_{\theta}\ell(\mathbf{x}_{1},y_{1};\theta)$, and which may also depend
on the optimizer state $\eta$.

$\mathrm{I}_{\Delta}(\mathbf{x},\mathbf{y};\theta)\overset{\text{def}}{=}f(\mathbf{y};\theta)-f(\mathbf{y},\theta^{\prime})\text{
where }\theta^{\prime}=\theta+\alpha g_{\ell}(\mathbf{x};\theta,\eta)$ (2.14)

Interference has been studied under a variety of names. Fort et al. (2019)
refer to the gradient alignment between data points (a measure related to
$I_{\nabla}$ as ‘stiffness’, while He and Su (2020) call a notion that more
closely resembles $\mathrm{I}_{\Delta}$ ‘elasticity’.

A number of approaches which specifically reduce the effect of a gradient
update for state $s$ on the target $V(s^{\prime})$ have been shown to improve
the stability and robustness of off-policy algorithms (Ghiassian et al., 2020;
Lo and Ghiassian, 2019). Prior works have also endeavoured to rigorously
define and analyze interference in deep RL (Liu et al., 2020b, a; Bengio et
al., 2020), and to study its role in the stability of offline algorithms
(Kumar et al., 2022). Similarly, some recent methods (Shao et al., 2022;
Pohlen et al., 2018) include an explicit penalty which discourages gradient
updates from affecting the target values used in TD updates, though this is
incidental to the principal contribution of these works. Tuning the degree of
interference exhibited by the function approximator to maximize performance
thus remains an open problem.

####  2.3.2 Multiple environments

Figure 2.2: Visualization of multi-environment generalization framework with
coinrun levels sampled from the ProcGen benchmark (Cobbe et al., 2020).

Most standard formulations of generalization in RL do not concern
generalization to unseen states within the training environment, but rather
generalization to novel environments, MDPs with different observations,
transition dynamics, and reward functions than those seen during training
(Kirk et al., 2021). This generalization problem in deep RL can be decomposed
into two categories: generalization of the function approximator to novel
observations from an MDP which follows the same underlying dynamics as the
training environment, and generalization to novel tasks, where some
combination of the reward function, observations, and transition dynamics are
allowed to vary. Many techniques that benefit one of these categories also
benefit the other (Packer et al., 2018). Farebrother et al. (2018) argue, for
instance, that an agent which has learned a policy which generalizes well will
be robust to changes in the difficulty of a game in addition to more
superficial changes in the observation space.

##### Formalism

We will be concerned with the generalization gap incurred by a policy trained
on a set of training environments $\mathcal{E}_{{\mathrm{train}}}$ when
deployed to a novel test environment.

$\mathbb{E}_{\mathcal{E}_{{\mathrm{train}}},\pi}[\sum_{t=0}^{\infty}\gamma^{t}R(x_{t},a_{t})]-\mathbb{E}_{\mathcal{E}_{{\mathrm{test}}},\pi}[\sum_{t=0}^{\infty}\gamma^{t}R(x_{t},a_{t})]$
(2.15)

Typically $\mathcal{E}_{{\mathrm{test}}}$ will be assumed to share some
structure with $\mathcal{E}_{{\mathrm{train}}}$, but the nature of this shared
structure may vary between problem settings. In large observation spaces,
$\mathcal{E}_{{\mathrm{test}}}$ may be equal to
$\mathcal{E}_{{\mathrm{train}}}$ with a different initial state distribution,
while for multi-environment problems $\mathcal{E}_{{\mathrm{train}}}$ may be
homomorphic to $\mathcal{E}_{{\mathrm{test}}}$ (Zhang et al., 2020a). In the
worst case, the evaluation environments may be chosen adversarially so as to
make generalization difficult; we will not consider this setting, instead
focusing predominantly on test environments that share some characteristics
with the training environment.

Procedurally generated environments (Cobbe et al., 2020; Küttler et al., 2020;
Samvelyan et al., 2021) enable us to sample an almost-unlimited number of
evaluation environments which are visually diverse but which share important
structural characteristics. A visualization of frames from different
procedurally-generated levels of the CoinRun game is given in Figure 2.2.
Being able to generate an unlimited set of environments brings us closer to
the setting of supervised learning, where training and evaluation data points
are sampled from the same distribution, by allowing us to sample training and
evaluation environments from the same distribution. This regime opens
interesting research questions concerning the optimal sampling of environments
during training so as to maximize the agent’s learning progress and
generalization performance (Jiang et al., 2021a; Parker-Holder et al., 2022).
However, the number of training environments in these settings is typically
orders of magnitude smaller than the number of training data points seen in
standard supervised learning benchmarks, making generalization more
challenging.

##### Regularization

Regularization has been widely shown to improve generalization in deep
learning (Srivastava et al., 2014; Krogh and Hertz, 1991), and its application
to reinforcement learning tasks has a rich history extending back to even
before the explosion of the deep RL paradigm (Farahmand, 2011). More recent
works focus specifically on the regularization of neural networks trained on
RL tasks (Cobbe et al., 2019; Farebrother et al., 2018), and show that many
regularization methods used to improve generalization in supervised deep
learning, such as dropout and $\ell_{2}$ regularization, are also applicable
to generalization in reinforcement learning. For example, Schrittwieser et al.
(2020) use $\ell_{2}$ regularization in their model-based algorithm which
attains state-of-the-art results in many game environments. Many recent works
have further sought to leverage the benefits of data augmentation to improve
generalization and robustness to input perturbations in reinforcement learning
agents (Raileanu et al., 2021; Laskin et al., 2020a, b; Hansen and Wang,
2021). Li and Pathak (2021) use low-frequency learned Fourier features to
regularize deep Q-networks towards smooth functions. Trust region-based
methods further leverage a form of explicit regularization on a learned policy
to stabilize learning (Schulman et al., 2015).

The stochasticity induced by some regularization methods serves a dual purpose
as a tool for exploration. Igl et al. (2019) use dropout masks to induce
stochasticity into the exploratory policy of an agent, while Fortunato et al.
(2018) apply a Gaussian perturbation to the network parameters. This approach
is also leveraged by the Rainbow (Hessel et al., 2018) architecture. Goyal et
al. (2019) use an information bottleneck architecture both to improve
generalization and to generate exploration strategies.

##### Transfer learning and meta-RL

In some problem settings, we seek to train an agent that is capable of quickly
adapting to novel environments. In the case where we wish to generalize to
novel rewards under the same transition dynamics as we saw during training,
this can be achieved via the successor representation (Dayan, 1993; Barreto et
al., 2017). In deep RL, we approximate the successor representation using
successor features, which have been shown to improve sample efficiency and
generalization in both single-agent RL and multi-agent settings (Machado et
al., 2018b; Filos et al., 2021). Recently, Abdolshah et al. (2021) present an
approach to learning successor features which can generalize across different
environments.

However, in many cases the novel environment may feature a novels observation
space with different transition dynamics from those seen at training. With the
advent of powerful function approximation architectures and the development of
effective meta-learning algorithms (Finn et al., 2017; Zintgraf et al., 2019),
generalization across multiple environments exhibiting less shared structure
may also be approached using transfer- and meta-learning techniques. Teh et
al. (2017) and Schmitt et al. (2018), for example, use policy distillation to
improve transfer learning to new tasks. Other approaches use context-dependent
policies (Rakelly et al., 2019; Sodhani et al., 2021), or encourage invariance
across training environments (Zhang et al., 2020b).

####  2.3.3 Summary

The study of generalization in the reinforcement learning literature has
historically been restricted to the framework of state abstractions, which
consider arbitrary aggregations of states rather than explicit inductive
biases on some input space (Li et al., 2006). Some attempts have been made to
define generalization with respect to input observations in single-environment
RL, but these definitions run into limitations due to the nonstationary nature
of RL, wherein improving prediction accuracy does not necessarily entail
improved performance (Bengio et al., 2020; Liu et al., 2020b). One
contribution of this thesis will be to present explicit definitions of
interference and learning capacity in deep RL that are less dependent on the
chaotic relationship between performance and prediction error, which will be
used to study the evolution of learned representations in deep RL agents. We
will explore both the benefits, e.g. improved robustness to input distribution
shifts, and the drawbacks, e.g. decreased ability to adapt to new reward
signals, of generalization between inputs. This exploration will result in
algorithmic tools to accelerate learning and improve generalization in deep RL
agents. While recent analysis has studied interference between different
regions of the state space as a source of instability and a cause of
performance plateaus, this line of work largely ignores the beneficial effects
of generalization.

In contrast, the benefits of generalization are widely acknowledged in the
context of multi-environment RL, where an agent is evaluated on previously-
unseen environments at test time. However, the current state of the art in
this regime requires prohibitive sample complexity in order to generalize well
outside of their training set (Cobbe et al., 2019), largely due to the weak
assumptions placed on the generative process from which environments are
sampled. While some approaches from the supervised learning literature, such
as weight decay, dropout, and data augmentation, have been shown to improve
generalization in deep RL, these approaches do not have as great an effect as
in supervised learning (Raileanu et al., 2021). Orthogonal approaches from the
meta- and transfer-learning literature have seen more success as they allow
for fine-tuning on the test set, though these methods encounter similar sample
complexity issues when applied to zero-shot generalization. In Chapter 8 we
will consider a characterization of families of environments in which
generalization from a handful of training environments is theoretically
possible, and present algorithms for identifying or learning state
abstractions for which an optimal policy will generalize zero-shot to any
novel environment from this family.

##  Chapter 3 The role of invariance in generalization

###  3.1 Introduction

The first object of interest in our study of learning dynamics and
generalization will be invariance. Invariance arises naturally as a part of
the data generating process in many real-world datasets, such as point clouds
(invariant to permutation), speech recognition (invariant to the average pitch
of the speaker’s voice), natural language (different sentences may be
semantically equivalent), and objects in natural images (often rotation- and
translation-invariant). In many instances, an invariance can be precisely
characterized in terms of the action of a specified group, meaning that the
input distribution is partitioned into equivalence classes, for example all
rotations of an image or all permutations of elements in a set. This structure
yields straightforward desiderata on the generalization properties of a model
trained to fit such a data-generating distribution: learning about one element
of an equivalence class should be sufficient for the model to correctly
identify any other element in the equivalence class. In this sense invariance
presents an idealized form of interference: updates to an invariant model’s
output on one element of an equivalence class will necessarily have an
identical effect across all inputs in this class.

It is not surprising then that models which capture the invariant structure of
a given distribution tend to perform better than those which do not on a
variety of tasks (Cohen and Welling, 2016; Fawzi et al., 2016; Salamon and
Bello, 2017). Yet while the empirical benefits of invariant models have been
widely corroborated, work which seeks to explicitly quantify and understand
how invariance can benefit generalization is scarce. Part of this challenge
stems from the breadth of approaches by which invariance can be incorporated
into a model class. One can build the invariance into the network as a
convolution or weight-tying scheme, average network predictions over
transformations of the input (feature averaging), or simply train on a dataset
augmented with these transformations (data augmentation). Comparison of these
different approaches requires different tools, as each incorporates
confounding factors into the learning process which are distinct from the
effect of the training method on invariance.

This chapter will employ a diverse range of techniques to gain insight into
how invariance benefits generalization. The analysis presented here will lay
the foundation for our study of more nebulous forms of inductive bias in
Chapter 4, and our study of generalization in reinforcement learning in
Chapters 7 and 8. We will provide both an empirical and theoretical analysis
contrasting neural networks trained to exhibit approximate invariance via data
augmentation and networks which exhibit exact invariance via architectural
design or feature averaging. Our empirical analysis will study the degree to
which data augmentation yields models whose invariances generalize outside of
the training set, evaluating the effect of these methods on interference
between inputs in the same equivalence class. We further characterize
conditions under which data augmentation will provably induce convergence to
an invariant model parameterization in Section 3.3.1, and study the
differential effects of data augmentation and feature averaging on
interference in Section 3.3.2.

In our theoretical analysis, we will shed light onto the role of invariant
structure in PAC-Bayes generalization bounds. Our main contribution will be a
ranking of a set of PAC-Bayes bounds on models trained with data augmentation
and with feature averaging, resulting in the practical conclusion that feature
averaging at test time improves upon data augmentation, which in turn
outperforms training on data sampled independently from $P_{\mathcal{D}}$. We
find that the differential effects of feature averaging and data augmentation
on interference mirror those on the model complexity term in our PAC-Bayes
bound. This analogy between model simplicity and interference in the setting
of group-theoretic invariances grounds our further study of interference in
the learning problems of Chapters 4, 7, and 8 where more nuanced inductive
biases must be identified.

###  3.2 Background

“Invariance” has been used to describe a number of related but distinct
phenomena in the machine learning and statistics literature concerning the
stability of a function’s output under some set of transformations of its
input (Zou et al., 2012; Raj et al., 2017; Jaderberg et al., 2015; van der
Wilk et al., 2018; Chen et al., 2019). We focus on invariance under the action
of a group $\mathcal{G}$, a setting shared by several prior works (e.g., Cohen
and Welling, 2016; Kondor and Trivedi, 2018; Bloem-Reddy and Teh, 2019). The
action of $\mathcal{G}$ on a set $\mathcal{X}$ is a mapping
$\alpha:\mathcal{G}\times\mathcal{X}\to\mathcal{X}$ which is compatible with
the group operation, i.e. ${\alpha(e,x)=x}$ and
${\alpha(g_{1},\alpha(g_{2},x))=\alpha(g_{1}g_{2},x)}$. We will abbreviate
this operation as $\alpha(g,x)=gx$ when the form of $\alpha$ is not important.
The orbit of any $x\in\mathcal{X}$ is the subset $\mathcal{G}_{x}$ of
$\mathcal{X}$ that can be obtained by applying an element of $\mathcal{G}$ to
$x$, $\mathcal{G}_{x}=\\{gx:g\in\mathcal{G}\\}$. We note that the sets
$\mathcal{G}_{x}$ are equivalence classes which partition the underlying set
$\mathcal{X}$, and will refer to them as the orbits of $\mathcal{G}$.

For mathematical simplicity, we assume $\mathcal{G}$ to be compact, with
(unique) normalized Haar measure denoted by $\lambda$.111$\lambda$ is
analogous to the uniform distribution on $\mathcal{G}$ (Haar, 1933).
Critically, any finite group is compact under the discrete topology. For the
purposes of this chapter, the reader may therefore substitute ‘compact’ with
‘finite’ and ‘Haar measure’ with ‘uniform distribution over $\mathcal{G}$’ and
preserve the correctness (though not the full generality) of the theoretical
results.

We denote a random element of $\mathcal{G}$ by $G$. A mapping
$f:\mathcal{X}\to\mathcal{Y}$ is invariant under $\mathcal{G}$ (or
$\mathcal{G}$-invariant) if

$\displaystyle f(gx)=f(x)\;,\quad\forall\,g\in\mathcal{G},\
x\in\mathcal{X}\;.$ (3.1)

The act of transforming an arbitrary function $f:\mathcal{X}\to\mathbb{R}$ to
be $\mathcal{G}$-invariant is known as symmetrization, and can be performed by
averaging the output of $f$ over the orbits of $\mathcal{G}$. We define this
symmetrization operator $S_{\mathcal{G}}$ as follows:

$\displaystyle
f^{\circ}(x):=S_{\mathcal{G}}f(x)=\mathbb{E}_{G\sim\lambda}[f(Gx)]\;,\quad
x\in\mathcal{X}\;.$ (3.2)

While other symmetrization methods exist, such as taking the supremum over the
group’s orbit, we will always use $f^{\circ}$ to refer to the average of $f$
over the orbit of its input. We consider a typical machine learning scenario,
with a training data set $\mathcal{D}^{n}$ of $n$ observations
$(X_{i},Y_{i})_{i=1}^{n}\in(\mathcal{X},\mathcal{Y})^{n}$ sampled i.i.d. from
some (unknown) probability distribution $P_{\mathcal{D}}$. Additionally,
$P_{\mathcal{D}}$ _is assumed to be $\mathcal{G}$-invariant_, i.e.

$\displaystyle P_{\mathcal{D}}(gX,Y)=P_{\mathcal{D}}(X,Y)\;,\quad
g\in\mathcal{G}\;.$ (3.3)

For example, $X$ may be an image of an animal, $Y$ a label of the animal, and
$\mathcal{G}$ the group of two-dimensional rotations.

Given a group action on the space $\mathcal{X}$, we can partition it into
disjoint orbits of the form $\\{g\Phi:g\in\mathcal{G}\\}$, where
$\Phi\in\mathcal{X}$ denotes an orbit representative. The
$\mathcal{G}$-invariance of $P_{\mathcal{D}}$ enables the decomposition of
$P_{\mathcal{D}}$ into a distribution over orbits $P_{\Phi}$ along with a
conditional distribution $P_{X|\Phi}=\lambda(g)$ where $g\Phi=X$ (see, e.g.,
Bloem-Reddy and Teh, 2019). Letting $G\sim\lambda(\mathcal{G})$, this entails
$(X,Y)\overset{D}{=}(G\Phi,Y)$ and $P_{\mathcal{D}}=P_{\Phi}\times
P_{X|\Phi}\times P_{Y|X}$. In other words, the distribution $P_{\mathcal{D}}$
can be interpreted as first sampling an equivalence class $\Phi$, then
sampling an element $X$ uniformly at random from the set $\Phi$, and then
sampling $Y$ based on the conditional distribution $P_{\mathcal{D}}(Y|X)$.

This decomposition will play a key role in our analysis in later sections, as
it will allow us to iterate expectations over $P_{\mathcal{D}}$ as

$\displaystyle\mathbb{E}_{(X,Y)\sim
P_{\mathcal{D}}}[f(X,Y)]=\mathbb{E}_{\Phi\sim
P_{\Phi}}[\mathbb{E}_{G\sim\lambda(\mathcal{G})}[\mathbb{E}_{Y\sim
P_{Y|X}}[f(G\Phi,Y)\mid\Phi,Y]\mid Y]]\;.$

This will allow us to marginalize over and condition on orbits of
$\mathcal{G}$, making it possible to reason about the expected behaviour of a
function within and between orbits.

####  3.2.1 Modes of invariance

When the data-generating distribution and target function are known to be
$\mathcal{G}$-invariant, incorporating this invariance into the model class is
an intuitive approach to accelerate convergence and improve generalization.
This can be done in a variety of ways.

Trained invariance is implemented as data augmentation (DA) (Fawzi et al.,
2016; Cubuk et al., 2019): elements $G_{ij}$ of $\mathcal{G}$ are applied to
each observation $X_{i}$ of the training data, with the label $Y_{i}$ left
unchanged. The result is an augmented dataset
$\mathcal{D}^{n}_{\mathcal{G}}=((G_{ij}X_{i},Y_{i})_{j\leq m})_{i\leq n}$ used
to minimize the augmented empirical risk

$\widehat{R}^{\circ}_{\ell}(f,\mathcal{D}^{n})=\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{G\sim\lambda}[\ell(f(GX_{i}),Y_{i})]$
(3.4)

which can be approximated using samples as follows

$\widehat{R}^{\circ}_{\ell}(f,\mathcal{D}^{n})\approx\frac{1}{nm}\sum_{i=1}^{n}\sum_{j=1}^{m}\ell(f(G_{ij}X_{i}),Y_{i})=\widehat{R}_{\ell}^{\widehat{\circ}}(f,\mathcal{D}^{n})\;.$
(3.5)

DA is now a standard method in the practitioner’s toolkit (Iyyer et al., 2014;
Zhou and Troyanskaya, 2015; Salamon and Bello, 2017; Zhao et al., 2018), and
recent theoretical work has further established connections to variance
reduction methods (Dao et al., 2019). A major benefit of data augmentation is
its versatility: while our study in this chapter focuses on augmentations that
can be expressed as the action of a group, data augmentation can also be used
to promote other properties of the target function such as smoothness by
leveraging a broader class of transformations.

Architectural invariance restricts the function class being learned to contain
only invariant functions, typically through either feature averaging (FA) or
symmetric function composition. FA symmetrizes the output of an arbitrary
neural network by computing an average over $\mathcal{G}$ at one or more
layers, such that the overall network is invariant under the action of
$\mathcal{G}$. In practice, averaging is typically done at the penultimate or
final layer, resulting in a $\mathcal{G}$-invariant network ${f}^{\circ}$. A
network $f$ with $D$ layers is written as the composition of
$h_{D}\circ\dots\circ h_{1}$, with the shorthand $h_{d}^{d^{\prime}}$
referring to the composition of layers $d$ through $d^{\prime}$. The empirical
risk of a network with FA at layer $d$ evaluated on $\mathcal{D}^{n}$ is

$\displaystyle\widehat{R}_{\ell}(f^{\circ},\mathcal{D}^{n})=\frac{1}{n}\sum_{i=1}^{n}\ell\big{(}h_{d}^{D}\circ\mathbb{E}_{G\sim\lambda}[h_{1}^{d-1}(GX_{i})],Y_{i}\big{)}\;.$

The development of invariant or equivariant neural network architectures
enables users to avoid the potentially costly computation of this expression
for a number of different group invariances. A body of literature of varying
degrees of generality has developed to propose new architectures which
characterize invariant function classes under a variety of group actions, and
to study their properties (Wood and Shawe-Taylor, 1996; Ravanbakhsh et al.,
2017; Kondor and Trivedi, 2018; Bloem-Reddy and Teh, 2019; Cohen et al., 2019;
Cohen and Welling, 2016).

####  3.2.2 PAC-Bayes generalizations bounds

As discussed in Section 2.2.1, PAC-Bayes bounds characterize the risk of a
randomized prediction rule; the randomization can be interpreted as a
posterior distribution over functions $Q$ that may depend on
$\mathcal{D}^{n}$, though this distribution need not correspond to the exact
posterior of a Bayesian model. A PAC-Bayes bound on generalization error is
typically expressed in terms of a sum of the empirical risk and the Kullback-
Leibler (KL) divergence (Kullback and Leibler, 1951) between $Q$ and a fixed
prior distribution $P$, which as mentioned in Chapter 2 can be interpreted as
characterizing the complexity of functions in the support of $P$.

We define the risk of a probability distribution $Q$ defined on a function
class $F$ in the natural way,

$R_{\ell}(Q)=\mathbb{E}_{f\sim Q}R_{\ell}(f)\quad\text{ and analogously
}\quad\widehat{R}_{\ell}(Q)=\mathbb{E}_{f\sim Q}\widehat{R}_{\ell}(f)\;.$

PAC-Bayes bounds are highly correlated with generalization performance (Jiang
et al., 2020), and in some cases provide non-vacuous results for even highly
over-parameterized neural networks (Dziugaite and Roy, 2017, 2018b; Zhou et
al., 2019). The following is a standard bound due to Catoni (2007), which
holds for general data generating distributions $P_{\mathcal{D}}$ evaluated on
the 0-1 loss.

######  Theorem 3.2.1 (Catoni (2007)).

Let $\mathcal{D}^{n}$ be sampled i.i.d. from $P_{\mathcal{D}}$, and let $\ell$
be the binary 0-1 loss. For any prior $P$ and any $\delta\in(0,1)$, with
probability $1-\delta$ over samples $\mathcal{D}^{n}$, for all posteriors $Q$
and for all $\beta>0$,

$R_{\ell}(Q)\leq\frac{1-e^{-\beta\widehat{R}_{\ell}(Q,\mathcal{D}^{n})-\frac{1}{n}(\text{\rm{KL}}(Q\;||\;P)+\log\frac{1}{\delta})}}{1-e^{-\beta}}\;.$
(3.6)

For bounded loss functions, analogous bounds rely on a sum of the empirical
risk and a function of the KL divergence (see, e.g., Dziugaite and Roy, 2017,
2018b). We state all results only for variations of Catoni’s bound (3.6), but
analogous findings hold for bounds of the form proposed by McAllester (1999).

###  3.3 Invariance and optimization dynamics

We begin with an analysis of the manner in which data augmentation and feature
averaging influence the optimization dynamics and resulting function output by
gradient descent. This analysis will yield two main insights, both concerning
the setting of convex losses: first, that while data augmentation can cause
gradient descent to converge to a $\mathcal{G}$-invariant solution in certain
realizable settings, in general it leads to networks that are only
approximately invariant. In fact, these networks can overfit to the training
dataset such that their variance over orbits of out-of-distribution data
increases over the course of training. Second, we will show that feature
averaging, in addition to ensuring $\mathcal{G}$-invariance, reduces the
variance of both the loss and its gradients. We will relate this finding to
the notion of gradient interference discussed in Section 2.3.1, illustrating
the effect of trained and architectural invariance on interference within and
between orbits of a group $\mathcal{G}$.

####  3.3.1 Data augmentation and trained invariance

To begin our study of optimization dynamics and invariance, we lay out some
basic properties of the symmetrized risk $\widehat{R}^{\circ}_{\ell}$ and its
interaction with invariant models. When the loss function is convex, Jensen’s
inequality can be applied to the augmented risk to compare DA and FA risk
estimates. We note that the relevant notion of convexity to this line of
argument is with respect to the loss function’s first argument, not the
parameters of the model. Many (though not all) widely used loss functions are
convex in this way (e.g., squared error, cross-entropy). We provide
counterexamples highlighting the necessity of this property in Appendix
10.2.1.

######  Proposition 3.3.1.

Let $\ell:\mathbb{R}\times\mathbb{R}\to\mathbb{R}_{+}$ be a loss function that
is convex in its first argument. Then for any $f:\mathcal{X}\to\mathcal{Y}$,

$\displaystyle\widehat{R}_{\ell}(f^{\circ},\mathcal{D}^{n})=\widehat{R}^{\circ}_{\ell}(f^{\circ},\mathcal{D}^{n})\leq\widehat{R}^{\circ}_{\ell}(f,\mathcal{D}^{n})\;,$

and therefore analogous inequalities hold for
$\widehat{R}_{\ell}(Q^{\circ},\mathcal{D}^{n})$, $R_{\ell}(f)$, and
$R_{\ell}(Q)$. Furthermore, if
$\ell(f({\,\vbox{\hbox{\tiny$\bullet$}}\,}),{\,\vbox{\hbox{\tiny$\bullet$}}\,})\in
L_{2}(P_{\mathcal{D}})$ (i.e., has finite second moment), then

$\displaystyle\text{\rm{Var}}_{\mathcal{D}^{n}\sim
P_{\mathcal{D}}^{n}}\big{[}\widehat{R}_{\ell}(f^{\circ},\mathcal{D}^{n})\big{]}\leq\text{\rm{Var}}_{\mathcal{D}^{n}\sim
P_{\mathcal{D}}^{n}}\big{[}\widehat{R}^{\circ}_{\ell}(f,\mathcal{D}^{n})\big{]}\;.$

The proof of this result can be found in Appendix 10.1; it follows from
Jensen’s inequality and the invariance of the data-generating distribution to
augmentations. The beneficial effect of an invariant function on the expected
risk stems from its reduced variance over orbits. This observation can be
flipped to characterize the detrimental effect of incorporating an invariance
in the model which is not present in the data: by mapping all elements of an
orbit to the same output, the model’s expected risk can be lower bounded by a
function of the average variance of the true labels over orbits.

In principle, symmetrizing the risk function via data augmentation should
regularize the training procedure towards invariant functions; however, in
most settings it is not possible to guarantee that the learned function will
be exactly $\mathcal{G}$-invariant, particularly on inputs that differ from
those seen during training. Neural networks trained with data augmentation
will in general exhibit low but nonzero variance over equivalence classes, and
their failure to generalize this low variance to out-of-distribution is
illustrated in Figure 3.1. One notable setting in which it is possible to
obtain such guarantees is linear function approximation with groups admitting
a linear representation over the input space. To state the result, we define a
data-generating distribution $P_{\mathcal{D}}$ supported on the set
$\mathcal{X}\times\mathcal{Y}$, where $\mathcal{Y}\subseteq\mathbb{R}$ and
$\mathcal{X}\subseteq V\simeq\mathbb{R}^{d}$ is a subset of a $d$-dimensional
vector space over $\mathbb{R}$, $V$, with dual vector space $V^{*}$. We assume
that $\mathcal{X}$ spans $V$. Furthermore, let $\mathcal{G}$ admit a linear
representation (Serre, 1977), $\rho:\mathcal{G}\to GL(V)$, with corresponding
dual $\rho^{*}_{g}=\rho_{g^{-1}}^{\top}$. Finally, we assume that
$P_{\mathcal{D}}$ is invariant to the action of $\mathcal{G}$ and that
$\mathbb{E}[y|\mathbf{x}]=\langle\mathbf{w}^{*},\mathbf{x}\rangle$ for some
unique $\mathbf{w}^{*}$.

######  Proposition 3.3.2.

Let $P_{\mathcal{D}}$ and $\mathbf{w}^{*}$ be defined as above, and let
$\ell(\mathbf{w})=\mathbb{E}_{\mathbf{x},y\sim
P_{\mathcal{D}}}[(\langle\mathbf{w},\mathbf{x}\rangle-y)^{2}]$. Then the
(global) minimizer $\mathbf{w}^{*}$ of $\ell$ satisfies
$\rho_{g}^{*}\mathbf{w}^{*}=\mathbf{w}^{*}$ for $\lambda$-almost all
$g\in\mathcal{G}$. In particular, the function
$f(\mathbf{x})=\langle\mathbf{w}^{*},\mathbf{x}\rangle$ is
$\mathcal{G}$-invariant.

In brief, this result shows that in the case of linear models, data
augmentation with respect to certain types of symmetry can induce exact
invariance in the learned function. The proof of Propsition 3.3.2 can be found
in Appendix 10.1. It follows straightforwardly from the condition that the
group act linearly on the input space along with the linearity of the function
$f(\mathbf{x})=\langle\mathbf{w},\mathbf{x}\rangle$, as in this case the
symmetrization of the parameters $\mathbf{w}$ is equivalent to the
symmetrization of the function, so by Jensen’s inequality any minimizer of
$\ell$ must be invariant to the group action. Under suitable step-size