summaries: This paper presents the idea of self-consistent reinforcement learning.  By enforcing "consistency" between value functions and approximate  (learned) models, a model-based learning algorithm can use synthetic experience generated from the approximate model.  This paper presents a few minor variations on this theme, and explores the results empirically.

summaries: This paper considers the problem of model-based reinforcement learning. The authors introduce a new loss, which drives the learned model and value function to be consistent with each other, rather than the model being consistent with the real environment and the value function being consistent with the real environment and the learned model. Experiments show that the proposed self-consistency loss can improve the sample efficiency and final performance.

summaries: This work proposes to improve model-based RL methods by enforcing the model and value learned to be “self-consistent”, i.e. to satisfy the Bellman eqn where both the Bellman operator and value function are defined in terms of the model. The motivation behind this is that in classic MBRL, the model and value function are learned separately, and the value function is typically learned to be consistent with the model but not vice versa. The authors propose to enforce self-consistency by adding a "self-consistency loss” (L_SC) to the base loss and jointly learning the model and value functions. L_SC is the squared TD error with a stop-gradient on the target and is calculated using rollouts from the model. Although in MuZero/Muesli the model and value functions are learned jointly as well, the targets of the loss are all defined with respect to data from the true environment and there is not term enforcing the “self-consistency” between model and value functions. Experimental evaluations on a selection of Atari games show that jointly updating the model and value functions with the additional L_SC term leads to improved sample efficiency as compared to Dyna and a model-free baseline. The authors also empirically study the search control problem.

summaries: The paper offers a new approach to model-based reinforcement learning with their idea of self-consistency of models and values. In the traditional (Dyna) model-based RL paradigm, an approximate model is learned with ground truth transitions and values (and policies) are updated to be consistent with the learned model. The proposed self-consistency update however, formulates a new update to the approximate model so as to make it consistent with the current value estimates using completely imagined (virtual) experience. This update is inspired by minimizing the bellman residual of the approximate model and value for a given policy (since for the true value and model, the bellman equation is satisfied and the residual is zero). An empirical analysis is presented for a set of variations of the proposed self-consistency objective with two types of (grounded) model updates -- MLE (maximum-likelihood) and value equivalent updates. In both tabular and function approximate settings, improved sample efficiency is demonstrated by the proposed method.

abstractText: Learned models of the environment provide reinforcement learning (RL) agents with flexible ways of making predictions about the environment. In particular, models enable planning, i.e. using more computation to improve value functions or policies, without requiring additional environment interactions. In this work, we investigate a way of augmenting model-based RL, by additionally encouraging a learned model and value function to be jointly self-consistent. Our approach differs from classic planning methods such as Dyna, which only update values to be consistent with the model. We propose multiple self-consistency updates, evaluate these in both tabular and function approximation settings, and find that, with appropriate choices, self-consistency helps both policy evaluation and control.

authors:
A. Abdolmaleki
J.T. Springenberg
Y. Tassa
R. Munos
N. Heess
M. Riedmiller

title: Maximum a Posteriori Policy Optimisation

venue: International Conference on Learning Representations

authors:
T.W. Archibald
K. McKinnon
L.C. Thomas

title: On the generation of markov decision processes

venue: Journal of the Operational Research Society, 46(3):354–361

authors:
I. Babuschkin
K. Baumli
A. Bell
S. Bhupatiraju
J. Bruce
P. Buchlovsky
D. Budden
T. Cai
A. Clark
I. Danihelka
C. Fantacci
J. Godwin
C. Jones
T. Hennigan
M. Hessel
S. Kapturowski
T. Keck
I. Kemaev
M. King
L. Martens
V. Mikulik
T. Norman
J. Quan
G. Papamakarios
R. Ring
F. Ruiz
A. Sanchez
R. Schneider
E. Sezener
S. Spencer
S. Srinivasan
W. Stokowiec

title: and F

venue: Viola. The DeepMind JAX Ecosystem

title: Residual algorithms: Reinforcement learning with function approximation

venue: Machine Learning Proceedings 1995, pages 30–37. Elsevier

authors:
P. Baudiš
J.L. Gailly

title: Pachi: State of the art open source Go program

venue: Advances in computer games, pages 24–38. Springer

authors:
M.G. Bellemare
Y. Naddaf
J. Veness
M. Bowling

title: The arcade learning environment: An evaluation platform for general agents

venue: Journal of Artificial Intelligence Research, 47:253–279

title: A Markovian decision process

venue: Journal of Mathematics and Mechanics, 6(5): 679–684

authors:
O. Bojar
A. Tamchyna

title: Improving translation model by monolingual data

venue: Proceedings of the Sixth Workshop on Statistical Machine Translation, pages 330–336

authors:
J. Bradbury
R. Frostig
P. Hawkins
M.J. Johnson
C. Leary
D. Maclaurin
G. Necula
A. Paszke
J. VanderPlas
S. Wanderman-Milne

title: and Q

venue: Zhang. JAX: composable transformations of Python+NumPy programs

authors:
A. Byravan
J.T. Springenberg
A. Abdolmaleki
R. Hafner
M. Neunert
T. Lampe
N. Siegel
N. Heess
M. Riedmiller

title: Imagined value gradients: Model-based policy optimization with tranferable latent dynamics models

venue: Conference on Robot Learning, pages 566–589. PMLR

authors:
V. Chelu
D. Precup

title: and H

venue: P. van Hasselt. Forethought and Hindsight in Credit Assignment. In Advances in Neural Information Processing Systems, volume 33, pages 2270–2281

authors:
S. Edunov
M. Ott
M. Auli
D. Grangier

title: Understanding back-translation at scale

venue: arXiv preprint arXiv:1808.09381

authors:
L. Espeholt
H. Soyer
R. Munos
K. Simonyan
V. Mnih
T. Ward
Y. Doron
V. Firoiu
T. Harley
I. Dunning

title: et al

venue: Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference on Machine Learning, pages 1407–1416. PMLR

title: Iterative Value-Aware Model Learning

venue: NeurIPS, pages 9090–9101

authors:
A.M. Farahmand
A. Barreto
D. Nikovski

title: Value-aware loss function for model-based reinforcement learning

venue: Artificial Intelligence and Statistics, pages 1486–1494. PMLR

authors:
G. Farquhar
T. Rocktäschel
M. Igl
S. Whiteson

title: Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning

venue: arXiv preprint arXiv:1710.11417

authors:
A. Filos
C. Lyle
Y. Gal
S. Levine
N. Jaques
G. Farquhar

title: PsiPhi-Learning: Reinforcement Learning with Demonstrations using Successor Features and Inverse Temporal Difference Learning

venue: arXiv preprint arXiv:2102.12560

authors:
K. Gregor
D.J. Rezende
F. Besse
Y. Wu
H. Merzic
A. Oord

title: Shaping belief states with generative environment models for rl

venue: arXiv preprint arXiv:1906.09237

authors:
C. Grimm
A. Barreto
S. Singh
D. Silver

title: The Value Equivalence Principle for Model- Based Reinforcement Learning

venue: Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual

authors:
A. Guez
F. Viola
T. Weber
L. Buesing
S. Kapturowski
D. Precup
D. Silver
N. Heess

title: Value-driven hindsight modelling

venue: arXiv preprint arXiv:2002.08329

authors:
D. Ha
J. Schmidhuber

title: World models

venue: arXiv preprint arXiv:1803.10122

authors:
D. Hafner
T. Lillicrap
J. Ba
M. Norouzi

title: Dream to control: Learning behaviors by latent imagination

venue: arXiv preprint arXiv:1912.01603

authors:
D. Hafner
T. Lillicrap
M. Norouzi
J. Ba

title: Mastering Atari with Discrete World Models

venue: arXiv e-prints, art

authors:
J.B. Hamrick
A.L. Friesen
F. Behbahani
A. Guez
F. Viola
S. Witherspoon
T. Anthony
L.H. Buesing
P. Veličković
T. Weber

title: On the role of planning in model-based deep reinforcement learning

authors:
M. Hessel
I. Danihelka
F. Viola
A. Guez
S. Schmitt
L. Sifre
T. Weber
D. Silver

title: and H

venue: van Hasselt. Muesli: Combining Improvements in Policy Optimization. CoRR, abs/2104.06159

authors:
M. Hessel
M. Kroiss
A. Clark
I. Kemaev
J. Quan
T. Keck
F. Viola

title: and H

venue: van Hasselt. Podracer architectures for scalable Reinforcement Learning. CoRR, abs/2104.06272

authors:
G.Z. Holland
E.J. Talvitie
M. Bowling

title: The effect of planning shape on dyna-style planning in high-dimensional state spaces

venue: arXiv preprint arXiv:1806.01825

authors:
M. Jaderberg
V. Mnih
W.M. Czarnecki
T. Schaul
J.Z. Leibo
D. Silver
K. Kavukcuoglu

title: Reinforcement learning with unsupervised auxiliary tasks

venue: arXiv preprint arXiv:1611.05397

authors:
L.P. Kaelbling
T. Lozano-Pérez

title: Hierarchical Task and Motion Planning in the Now

venue: Proceedings of the 1st AAAI Conference on Bridging the Gap Between Task and Motion Planning, AAAIWS’10-01, page 33–42. AAAI Press

authors:
G. Kahn
A. Villaflor
B. Ding
P. Abbeel
S. Levine

title: Self-supervised deep reinforcement learning with generalized computation graphs for robot navigation

venue: 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 5129–5136. IEEE

authors:
L. Kaiser
M. Babaeizadeh
P. Milos
B. Osinski
R.H. Campbell
K. Czechowski
D. Erhan
C. Finn
P. Kozakowski
S. Levine
R. Sepassi
G. Tucker
H. Michalewski

title: Model-Based Reinforcement Learning for Atari

authors:
P.R. Kumar
P. Varaiya

title: Stochastic systems: Estimation

venue: identification, and adaptive control. SIAM

authors:
A.X. Lee
A. Nagabandi
P. Abbeel
S. Levine

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

venue: arXiv preprint arXiv:1907.00953

authors:
M.C. Machado
M.G. Bellemare
E. Talvitie
J. Veness
M. Hausknecht
M. Bowling

title: Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents

venue: Journal of Artificial Intelligence Research, 61:523–562

title: Gradient Temporal-Difference Learning Algorithms

venue: PhD thesis, CAN

authors:
T.M. Moerland
J. Broekens
C.M. Jonker

title: Model-based reinforcement learning: A survey

venue: arXiv preprint arXiv:2006.16712

authors:
R. Munos
T. Stepleton
A. Harutyunyan
M.G. Bellemare

title: Safe and efficient off-policy reinforcement learning

venue: Proceedings of the 30th International Conference on Neural Information Processing Systems, pages 1054–1062

authors:
J. Oh
X. Guo
H. Lee
R. Lewis
S. Singh

title: Action-conditional video prediction using deep networks in atari games

venue: arXiv preprint arXiv:1507.08750

authors:
J. Oh
S. Singh
H. Lee

title: Value prediction network

venue: arXiv preprint arXiv:1707.03497

authors:
Y. Pan
J. Mei
A.-M. Farahmand

title: Frequency-based Search-control in Dyna

authors:
T. Pohlen
B. Piot
T. Hester
M.G. Azar
D. Horgan
D. Budden
G. Barth-Maron
H. Van Hasselt
J. Quan
M. Večerík

title: et al

venue: Observe and look further: Achieving consistent performance on atari. arXiv preprint arXiv:1805.11593

title: Markov Decision Processes: Discrete Stochastic Dynamic Programming

venue: John Wiley & Sons, Inc., USA, 1st edition

authors:
S. Racanière
T. Weber
D.P. Reichert
L. Buesing
A. Guez
D. Rezende
A.P. Badia
O. Vinyals
N. Heess
Y. Li

title: et al

venue: Imagination-augmented agents for deep reinforcement learning. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 5694–5705

authors:
D.J. Rezende
I. Danihelka
G. Papamakarios
N.R. Ke
R. Jiang
T. Weber
K. Gregor
H. Merzic
F. Viola
J. Wang

title: et al

venue: Causally correct partial models for reinforcement learning. arXiv preprint arXiv:2002.02836

authors:
J. Richalet
A. Rault
J. Testud
J. Papon

title: Model predictive heuristic control

venue: Automatica (journal of IFAC), 14(5):413–428

title: An on-line algorithm for dynamic reinforcement learning and planning in reactive environments

venue: 1990 IJCNN international joint conference on neural networks, pages 253–258. IEEE

authors:
J. Schrittwieser
I. Antonoglou
T. Hubert
K. Simonyan
L. Sifre
S. Schmitt
A. Guez
E. Lockhart
D. Hassabis
T. Graepel
T. Lillicrap
D. Silver

title: Mastering Atari, Go, chess and shogi by planning with a learned model

authors:
J. Schrittwieser
T. Hubert
A. Mandhane
M. Barekatain
I. Antonoglou
D. Silver

title: Online and Offline Reinforcement Learning by Planning with a Learned Model

authors:
D. Silver
J. Veness

title: Monte-Carlo Planning in Large POMDPs

venue: Advances in Neural Information Processing Systems, volume 23, pages 2164–2172. Curran Associates, Inc.

title: H

venue: van Hasselt, M. Hessel, T. Schaul, A. Guez, T. Harley, G. Dulac-Arnold, D. Reichert, N. Rabinowitz, A. Barreto, et al. The predictron: End-to-end learning and planning. In International Conference on Machine Learning, pages 3191–3199. PMLR

title: Dyna, an Integrated Architecture for Learning, Planning, and Reacting

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

title: Reinforcement Learning: An Introduction

venue: The MIT Press, Cambridge, MA

authors:
A. Tamar
Y. Wu
G. Thomas
S. Levine
P. Abbeel

title: Value iteration networks

venue: arXiv preprint arXiv:1602.02867

authors:
H. van Hasselt
M. Hessel
J. Aslanides

title: When to use parametric models in reinforcement learning

venue: In Advances in Neural Information Processing Systems,

title: Learning how the world works: Specifications for predictive networks in robots and brains

venue: Proceedings of IEEE International Conference on Systems, Man and Cybernetics, NY

authors:
T. Yu
C. Lan
W. Zeng
M. Feng
Z. Chen

title: PlayVirtual: Augmenting Cycle-Consistent Virtual Trajectories for Reinforcement Learning

venue: arXiv preprint arXiv:2106.04152

authors:
S. Zhang
W. Boehmer
S. Whiteson

title: Deep Residual Reinforcement Learning

venue: In AAMAS 2020: Proceedings of the Nineteenth International Joint Conference on Autonomous Agents and Multi-Agent Systems,

authors:
J.-Y. Zhu
T. Park
P. Isola
A.A. Efros

title: Unpaired image-to-image translation using cycle-consistent adversarial networks

venue: Proceedings of the IEEE international conference on computer vision, pages 2223–2232

heading: 1 Introduction

text:
Models of the environment provide reinforcement learning (RL) agents with flexible ways of making predictions about the environment. They have been used to great effect in planning for action selection [45, 29, 49], and for learning policies or value functions more efficiently [51]. Learning models can also assist representation learning, serving as an auxiliary task, even if not used for planning [25, 48]. Traditionally, models are trained to be consistent with experience gathered in the environment. For instance, an agent may learn maximum likelihood estimates of the reward function and state-transition probabilities, based on the observed rewards and state transitions. Alternatively, an agent may learn a model that only predicts behaviourally-relevant quantities like rewards, values, and policies [47].
In this work, we study a possible way to augment model-learning, by additionally encouraging a learned model m̂ and value function v̂ to be jointly self-consistent, in the sense of jointly satisfying the Bellman equation with respect to m̂ and v̂ for the agent’s policy π. Typical methods for using models in learning, like Dyna [51], treat the model as a fixed best estimate of the environment, and only update the value to be consistent with the model. Self-consistency, by contrast, jointly updates the model and value to be consistent with each other. This may allow information to flow more flexibly between the learned reward function, transition model, and approximate value. Since the true model and value are self-consistent, this type of update may also serve as a useful regulariser.
We investigate self-consistency both in a tabular setting and at scale, in the context of deep RL. There are many ways to formulate a model-value update based on the principle of self-consistency, but we find that naive updates may be useless, or even detrimental. However, one variant based on a semi-gradient temporal difference objective can accelerate value learning and policy optimisation. We evaluate different search-control strategies (i.e. the choice of states and actions used in the selfconsistency update), and show that self-consistency can improve sample efficiency in environments such as Atari, Sokoban and Go. We conclude with experiments designed to shed light on the mechanisms by which our proposed self-consistency update aids learning.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

heading: 2 Background

text:
We adopt the Markov Decision Process (MDP) formalism for an agent interacting with an environment [7, 42]. The agent selects an action At ∈ A in state St ∈ S on each time step t, then observes a reward Rt+1 and transitions to the successor state St+1. The environment dynamics is given by a “true” model m∗ = (r, P ), consisting of reward function r(s, a) and transition dynamics P (s′|s, a). The behaviour of the agent is characterized by a policy π : S → ∆A, a mapping from states to the space of probability distributions over actions. The agent’s objective is to update its policy to maximise the value vπ(s) = Em∗,π [∑∞ j=0 γ jr(St+j,At+j) | St = s ]
of each state. We are interested in model-based approaches to this problem, where an agent uses a learned model of the environment m̂ = (r̂, P̂ ), possibly together with a learned value function v̂(s) ≈ vπ , to optimise π. The simplest approach to learning such a model is to make a maximum likelihood estimate of the true reward and transition dynamics, based on the agent’s experience [32]. The model can then be used to perform updates for arbitrary states and action pairs, entirely in imagination. Dyna [51] is one algorithm with this property, that has proven effective in improving the data efficiency of RL [22, 31].
An alternative model-learning objective is value equivalence1 i.e. equivalence of the true and learned models in terms of the induced values for the agent’s policy. This can be formalised as requiring the induced Bellman operator of the model to match that of the environment:
T πm∗vπ = T πm̂vπ, (1)
where the Bellman operator for a model m̂ is defined by T πm̂v = Em̂,π [r(s, a) + γv(s′)]. Such models do not need to capture every detail of the environment, but must be consistent with it in terms of the induced Bellman operators applied to certain functions (in this case, the agent’s value vπ).
We take a particular interest in value equivalent models for two reasons. First, they are used in stateof-the-art agents, as described in the next section and used as baselines in our deep RL experiments. Second, these algorithms update the model using a value-based loss, like the self-consistency updates we will study in this work. We describe this connection further in Section 3.3.

heading: 2.1 Value equivalent models in deep RL

text:
A specific flavour of value equivalent model is used by the MuZero [47] and Muesli [25] agents. Both use a representation function (encoder) hθ to construct a state zt = hθ(o1:t) from a history of observations o1:t2. We define a deterministic recurrent model m̂ ≡ mθ in this latent space. On each step k the model, conditioned on an action, predicts a new latent state zkt and reward r̂ k t ; the agent also predicts from each zkt a value v̂ k t and a policy π̂ k t . We use superscripts to denote time steps in the model: i.e., zkt is the latent state after taking k steps with the model, starting in the root state z 0 t ≡ zt.
Muesli and MuZero unroll such a model for K steps, conditioning on a sequence of actions taken in the environment. The learned components are then trained end-to-end, using data gathered from interaction with the environment, D∗π , by minimizing the loss function
Lbaset (m̂, v̂, π̂|D∗π) = ED∗π K∑ k=0 [ `r(rtargett+k , r̂ k t ) + ` v(vtargett+k , v̂ k t ) + ` π(πtargett+k , π̂ k t ) ] . (2)
This loss is the sum of a reward loss `r, a value loss `v and a policy loss `π . Note that there is no loss on the latent states zkt : e.g. these are not required to match zt+k = hθ(o1:t+k). The targets r target t+k for the reward loss are the true rewards Rt+k observed in the environment. The value targets are constructed from sequences of rewards and n-bootstrap value estimates vtargett+k = ∑n j=1 γ j−1rt+k+j + γ nṽt+k+n. The bootstrap value estimates ṽ, as well as the policy targets, are constructed using the model m̂. This is achieved in MuZero by applying Monte-Carlo tree search, and, in Muesli, by using one-step look-ahead to create MPO-like [1] targets. Muesli also employs a policy-gradient objective which does not update the model. We provide some further details about Muesli in Appendix B.2, as we use
1We borrow the terminology from Grimm et al. [19], who provided a theoretical framework to reason about such models. Related ideas were also introduced under the name ‘value aware’ model learning [15, 14]
2We denote sequences of outcomes of random variables O1 = o1, . . . , Ot = ot as o1:t for simplicity.
Algorithm 1: Model-based RL with joint grounded and self-consistency updates. input : initial m̂, v̂, π output :estimated value v̂ ≈ vπ and/or optimal policy π∗ repeat
Collect D∗π from m∗ following π Compute grounded loss: Lbase(m̂, v̂, π|D∗π) // e.g. Eq.(2) Generate D̂µ from m̂ following µ Compute self-consistency loss: Lsc(m̂, v̂|D̂µ) // see Eqs.(4,5,6) Update m̂, v̂, π by minimising (e.g., with SGD): L = Lbase + Lsc
until convergence;
it as a baseline for our deep RL experiments. For comprehensive descriptions of MuZero and Muesli we refer to the respective publications [47, 25].
In both cases, the model is used to construct targets based on multiple imagined actions, but the latent states whose values are updated by optimising the objective (2) always correspond to real states that were actually encountered by the agent when executing the corresponding action sequence in its environment. Value equivalent models could also, in principle, update state and action pairs entirely in imagination (similarly to Dyna) but, to the best of our knowledge, this has not been investigated in the literature. Our proposed self-consistency updates provide one possible mechanism to do so.

heading: 3 Self-consistent models and values

text:
A particular model m̂ and policy π induce a corresponding value function vπm̂, which satisfies a Bellman equation T πm̂vπm̂ = vπm̂. We describe model-value pairs which satisfy this condition as self-consistent. Note that the true model m∗ and value function vπ are self-consistent by definition.
If an approximate model m̂ and value v̂ are learned independently, they may only be asymptotically self-consistent, in that the model is trained to converge to the true model, and the estimated value to the true value. Model-based RL algorithms, such as Dyna, introduce an explicit drive towards consistency throughout learning: the value is updated to be consistent with the model, in addition to the environment. However, the model is only updated to be consistent with transitions in the real environment D∗π , and not to be consistent with the approximate values. Instead, we propose to update both the model and values so that they are self-consistent with respect to trajectories D̂µ sampled by rolling out a model m̂, under action sequences sampled from some policy µ (with µ 6= π in general). We conjecture that allowing information to flow more freely between the learned rewards, transitions, and values, may make learning more efficient. Self-consistency may also implement a useful form of regularisation – in the absence of sufficient data, we might prefer a model and value to be selfconsistent. Finally, in the context of function approximation, it may help representation learning by constraining it to support self-consistent model-value pairs. In this way, self-consistency may also be valuable as an auxiliary task even if the model is not used for planning.

heading: 3.1 Self-consistency updates

text:
For the model and value to be eventually useful, they must still be in some way grounded to the true environment: self-consistency is not all you need. We are therefore interested in algorithms that include both grounded model-value updates, and self-consistency updates. This may be implemented by alternating the two kind of updates, or in single optimisation step (c.f. Algorithm 1) by gathering both real D∗π and imagined D̂µ experience and then jointly minimize the loss
L(m̂, v̂, π|D∗π + D̂µ) = Lbase(m̂, v̂, π|D∗π) + Lsc(m̂, v̂|D̂µ). (3)
There are many possible ways to leverage the principle of self-consistency to perform the additional update to the model and value (Lsc). We will first describe several possible updates based on 1-step temporal difference errors computed from K-step model rollouts following µ = π, but we will later
consider more general forms. The most obvious update enforces self-consistency directly:
Lsc-residual(m̂, v̂) = Em̂,π [ K∑ k=0 ( r̂(sk, ak) + γv̂(sk+1)− v̂(sk) )2] . (4)
When used only to learn values, updates that minimize this loss are known as “residual” updates. These have been used successfully in model-free and model-based deep RL [57], even without addressing the double sampling issue [4]. However, it has been observed in the model-free setting that residual algorithms can have slower convergence rates than the TD(0) update [4], and may fail to find the true value [35]. In our setting, the additional degrees of freedom, due to the loss depending on a learned model, could allow the system to more easily fall into degenerate self-consistent solutions (e.g. zero reward and values everywhere). To alleviate this concern, we can also design our self-consistent updates to mirror a standard TD update, by treating the entire target r̂(sk, ak) + γv̂(sk+1) as fixed:
Lsc-direct(m̂, v̂) = Em̂,π [ K∑ k=0 ( ⊥ ( r̂(sk, ak) + γv̂(sk+1) ) − v̂(sk) )2] , (5)
where ⊥ indicates a suppressed dependence (a “stop gradient” in the terminology of automatic differentiation). This is sometimes referred to as using the “direct” method, or as using a semigradient objective [52]. Dyna minimises this objective, but only updates the parameters of the value function (Fig. 1a). We propose to optimise the parameters of both the model and the value (Fig. 1b).
For k = 0, this will not update the transition model, as v(s0) does not depend on it. This necessitates the use of multi-step model rollouts (k ≥ 1). Using multi-step model rollouts is also typically better in practice [27], at least for traditional model-based value learning. This form of the self-consistency update will also never update the reward model. We hypothesise this may be a desirable property, since the grounded learning of the reward model is typically well-behaved. Further, if the grounded model update enforces value equivalence with respect to the current policy’s value, the transition dynamics are policy-dependent, but the reward model may be stationary. Consequently, we expect this form of update may have a particular synergy with the value-equivalent setting.
Alternatively, we consider applying a semi-gradient in the other direction:
Lsc-reverse(m̂, v̂) = Em̂,π [ K∑ k=0 ( r̂(sk, ak) + γv̂(sk+1)−⊥ ( v̂(sk) ))2] (6)
This passes information forward in model-time from the value to the reward and transition model, even if k = 0.

heading: 3.2 Self-consistency in a tabular setting

text:
First, we instantiate these updates in a tabular setting, using the pattern from Algorithm 1, but alternating the updates to the model and value for the grounded and self-consistency losses rather than adding the losses together for a single joint update. At each iteration, we first collect a batch of transitions from the environment (s, a, r, s′) ∼ m∗. The reward model is updated as r̂(s, a)← r(s, a) + αr(r(s, a)− r̂(s, a)). We update the value with TD(0): v̂(s)← (1− αv)v̂(s) + αv(r(s, a) + γv̂(s ′)). Here, αr and αv are learning rates.
Next, we update the model, using either a maximum likelihood approach or the value equivalence principle. The transition model P̂ is parameterised as a softmax of learned logits: P̂ (s′|s, a) = softmax(p̂(s, a))[s′]. To learn a maximum likelihood model we follow the gradient of log P̂ (s′|s, a). For a value equivalent model, we descend the gradient of the squared difference of next-state values under the model and environment: ( ∑ ŝ′,a π(a|s)P̂ (s′|s, a)v̂(ŝ′)− v̂(s′))2.
Then, we update either the value only to be consistent with the model (Dyna), or update both the model and value to be self-consistent according to the objectives in the previous section. For control, we can construct an action-value estimate Q̂(s, a) = r̂(s, a) + γ ∑ s′ P̂ (s, a, s
′)v̂(s′) with the model m̂ = {r̂, P̂}, and act -greedily using Q̂.

heading: 3.3 Self-consistency in deep RL

text:
Self-consistency may be applied to deep RL as well, by modifying a base agent that learns a differentiable model m̂ and value v̂ using the generic augmented loss (3). We follow the design of Muesli [25] for our baseline, using the objective (2) to update the encoder, model, value, and policy in a grounded manner. We then add to the loss an additional self-consistency objective that generalises to the latent state representation setting:
Lscπ = Ez0∼Z0Ea1,...,aK∼µ [ K∑ k=1 lv ( Ĝπk:K , v(z k) )] . (7)
The base objective (2) already jointly learns a model and value to minimise a temporal-difference value objective `v(vtarget, v̂). Our self-consistency update is thus closely related to the original value loss, but differs in two important ways. First, instead of using rewards and next states from the environment, the value targets Ĝπk:K are (K − k)-step bootstrapped value estimates constructed with rewards and states predicted by rolling out the model m̂. The reliance of MuZero-style VE on the environment for bootstrap targets is illustrated schematically in Figure 1c, in contrast to self-consistency. Second, because the targets do not rely on real interactions with the environment, we may use any starting state and sequence of actions, like in Dyna.
Z0 denotes the distribution of starting latent states from which to roll out the model for K steps, and µ is the policy followed in the model rollout. When µ differs from π, we may use off-policy corrections to construct the target Ĝπk:K ; in our experiments, we use V-Trace [13]. For now, we default to µ = π, and to using the same distribution for Z0 as for the base objective. We revisit these choices in Section 4.2. Note that by treating different parts of the loss in 7 as fixed, we can recover the residual, direct, and reverse forms of self-consistency.
The self consistency objective is illustrated schematically in Figure 1. Subfigures (a) and (b) contrast a Dyna update to the value function with the SC-Direct update, which uses the same objective to update both value and model. Subfigure (c) shows the type of value equivalent update used in MuZero and Muesli. The objective is similar to SC-Direct, but the value targets come from real experience. Consequently, the model unroll must use the actions that were taken by the agent in the environment. Self-consistency instead allows to update the value and model based on any trajectories rolled out in the model; we verify the utility of this flexibility in Section 4.2.
Baseline Dyna SC-Residual SC-Direct SC-Reverse
0 25 k 50 k0.0
0.1
0.2
0.3
0.4
0.5
0.6
N or
m al
is ed
v al
ue e
rr or
Evaluation (MLE)
0 25 k 50 k0.0
0.1
0.2
0.3
0.4
0.5
0.6 Evaluation (VE)
0 25 k 50 k0.75
0.80
0.85
0.90
0.95
1.00
N or
m al
is ed
v al
ue
Control (MLE)
0 25 k 50 k0.75
0.80
0.85
0.90
0.95
1.00 Control (VE)
Number of iterations
(a) Self-consistency in random tabular MDPs. Each experiment was run with 30 independent replicas using different random seeds. (Left) Normalised value prediction error for policy evaluation, using MLE and VE models respectively. (Right) Normalised policy values for control, using MLE and VE models respectively.
0 100M 200M
10
5
0
5
Av g.
e pi
so de
r et
ur n
sokoban
0 100M 200M 0
8k
16k
24k alien
0 100M 200M 0
8k
16k
24k
assault
0 100M 200M 0
10k
20k
30k asteroids
0 100M 200M 0
10k
20k
30k
wizard
0 100M 200M 0
15k
30k
zaxxon
Number of frames 0 100M 200M
10
5
0
5
Av g.
e pi
so de
r et
ur n
sokoban
1 0M 200M 0 1 0M 200M 0
8k
16k
24k
lt
0 200M 0
10k
20k
30k asteroids
0 2 0M 0
10k
20k
30k wizard
0 10 M 200 0
15k
30k
zaxxon
u ber of fra es
(b) Self-consistency on Sokoban and 5 Atari games. Each experiment was run with 5 independent replicas using different random seeds. The self-consistency updates were applied to variants of Muesli, a model-based RL agent using deep neural networks for function approximation and value equivalent models for planning.
Figure 2: An evaluation of Dyna and the self-consistency updates described in section 3.1. Each variant was evaluated in both a tabular setting (a) and with function approximation (b). The experiments included both MLE and value equivalent models. Shaded regions denote 90% CI.

heading: 4 Experiments

heading: 4.1 Sample efficiency through self-consistency

text:
Tabular. In our first set of experiments, we used random “Garnet” MDPs [2] to study different combinations of grounded and self-consistent updates for approximate models and values. We followed an alternating minimization approach as described in Section 3.2. In each case, we applied the Dyna, SC-Direct, SC-Reverse, SC-Residual updates at each iteration starting at every state in the MDP. See Appendix A for further details of our experimental setup.
Figure 2a (on the left) shows the relative value error, calculated as |(v(s)− v̂(s))/v(s)|, averaged across all states in the MDP, when evaluating a random policy. As we expected, Dyna improved the speed of value learning over a model-free baseline. We saw considerable difference among the self-consistent updates. The “residual” and “reverse” updates were ineffective and harmful, respectively. In contrast, the “direct” self-consistency update was able to accelerate learning, even compared to Dyna, for both MLE and value-equivalent models.
Figure 2a (on the right) shows the value of the policy (normalised by the optimal value), averaged across all states, for a control experiment with the same class of MDPs. At each iteration, we construct Q̂(s, a) estimates using the model, and follow an -greedy policy with = 0.1. In this case, the variation between methods is smaller. The “direct” method still achieves the best final performance. However, the “residual” and “reverse” self-consistency updates are competitive for tabular control.
Function approximation. We evaluated the same selection of imagination-based updates in a deep RL setting, using a variant of the Muesli [25] agent. The only difference is in the hyperparameters for batch size and replay buffer, as documented in the Appendix (this seemed to perform more stably as a baseline in some early experiments). All experiments were run using the Sebulba distributed
architecture [26] and the joint optimization sketched in Algorithm 1. We used 5 independent replicas, with different random seeds, to assess the performance of each update.
We evaluated the performance of the updates on a selection of Atari 2600 games [6], plus the planning environment Sokoban. In Sokoban, we omit the policy-gradient objective from the Muesli baseline because with the policy gradient, the performance ceiling for this problem is quickly reached. Using only the MPO component of the policy update relies more heavily on the model, giving us a useful testbed for our methods in this environment, but should not be regarded as state-of-the-art.
We first evaluated a Dyna baseline, using experience sampled from the latent-space value-equivalent model. As in the tabular case, we did so by taking the gradient of the objective (7) with respect to only the parameters of the value predictor v̂. Unlike the simple tabular case, where Dyna performed well, we found this approach to consistently degrade performance at scale (Figure 2b, “Dyna”). To the best of our knowledge, this was the first attempt to perform Dyna-style planning with value-equivalent models, and it was a perhaps-surprising negative result. Note the contrast with the success reported in previous work [23] for using Dyna-style planning with MLE models on Atari.
Next we evaluated the self-consistency updates. We found that the “residual” self-consistency update performed poorly. We conjecture that this was due to the model converging on degenerate trivially self-consistent solutions, at the cost of optimising the grounded objectives required for effective learning. The degradation in performance was even more dramatic than in the tabular setting, perhaps due to the greater flexibility provided by the use of deep function approximation.
This hypothesis is consistent with our next finding: both semi-gradient objectives performed much better. In particular, the “direct” self-consistency objective increased sample efficiency over the baseline in all but one environment. We found the “reverse” objective to also help in multiple domains, although less consistently across the set of environments we tested. As a result, in subsequent experiments we focused on the “direct” update, and investigated whether the effectiveness of this kind of self-consistency update can be improved further.

heading: 4.2 Search control for self-consistency

text:
Search control is the process of choosing the states and actions with which to query the model when planning [52]; we know from the literature that this can substantially influence the effectiveness of planning algorithms such as Dyna [40]. Since our self-consistency objective also allows flexibly updating the value of states that do not correspond to the states observed from real actions taken in the environment, we investigated the effect of different choices for imagination policy µ and starting states Z0. Which policy should be followed in imagination? We now explore four choices for the policy followed in imagination to generate trajectories used by the self-consistency updates. The simplest choice is to sample actions according to the same policy we use to interact with the real environment (µ = π). A second option is to pick a random action with probability , and otherwise follows π (we try =0.5). The third option is to avoid the first action that was actually taken in the environment to ensure diversity from the real behaviour. After sampling a different first action, the policy π is used to sample the rest (Avoid a0). The last option is to replay the original sequence of actions that were taken in the environment (Original actions). This corresponds to falling back to the base agent’s value loss, but with rewards and value targets computed using the model in place of the real experience.
In Figure 3a, we compare these options. A notable finding was that just replaying exactly the actions taken in the environment (in brown) performed considerably worst than re-sampling the actions from π (in blue). This confirms our intuition that the ability to make use of arbitrary trajectories is a critical aspect of the self-consistency update. Introducing more diversity by adding noise to the policy π had little effect on most environments, although the use of = 0.5 did provide a benefit in assault and asteroids. Avoiding the first behaviour action, to ensure that only novel trajectories are generated in imagination, also performed reasonably well. Overall, as long as we generated a new trajectory by sampling from a policy, the self-consistency update was fairly robust to all the alternatives considered.
From which states should we start imagined rollouts? A search control strategy must also prescribe how the starting states are selected. Since we use a latent-space model, we investigated using latent-space perturbations to select starting states close (but different) from those encountered in the environment. We used a simple perturbation, multiplying the latent representation z0t element-wise
so ko
ba n
ali en
as sa
ult
as ter
oid s
wi za
rd
za xx
on 0.00
0.25
0.50
0.75
1.00
1.25
N or
m al
is ed
A U
C
Original actions =
=0.5 Avoid a0
Latent noise, = Latent noise, =0.5
(a) Different search control strategies for self-consistency. Each experiment used 5 random seeds, error bars denote 90% CI.
0 25 50 75 100 125 150 175 200 Millions of frames
0%
100%
200%
300%
400%
500%
M ed
ia n
hu m
an -n
or m
al iz
ed s
co re
LASER
Rainbow
atari57 median Muesli Muesli + SC
(b) Median human-normalised episode return on across 57 Atari games. Each experiment used 4 random seeds, shaded region denote 90% CI.
0 500 M 1 G 1.5 G 2 G 2.5 G Environment steps
0.0
0.1
0.2
0.3
0.4
0.5
0.6
W in
ra te
v s
Pa ch i Method
Baseline + SC
MCTS True False
(c) Performance on 9x9 self-play Go, with evaluation against Pachi [5] Each experiment used 5 random seeds, shaded region denote 90% CI.
Figure 3: Search control for self consistency, and evaluations on Atari and Go.
by noise drawn from U(0.8, 1.2) (Latent noise, µ = π). The results for these experiments are also shown in Figure 3a. The effect was not large, but the variant combining latent state perturbations with − exploration in the imagination policy (Latent noise, = 0.5) performed best amongst all variants we considered. The fact that some benefit was observed with such rudimentary data augmentation suggests that diverse starting states may be of value. Self-consistency could therefore be effective in semi-supervised settings where we have access to many valid states (e.g. reachable configurations of a robotic system), but few ground-truth rewards or transitions.
A full evaluation. In the next experiments we evaluated our self-consistency update on the full Atari57 benchmark, with the search-control strategy found to be most effective in the previous section. As shown in Figure 3b, we observed a modest improvement over the Muesli baseline in the median human-normalised score across the 57 environments. We also evaluated the same self-consistency update in a small experiment on self-play 9x9 Go in combination with a Muesli-MPO baseline, with results shown in Figure 3c. Again, the agent using self-consistency updates outperformed the baseline, both with (solid lines) and without (dashed lines) the use of MCTS at evaluation time. This confirms that the self-consistent model is sufficiently robust to be used effectively for planning online.

heading: 4.3 How does self-consistency work?

text:
Representation learning. In our final set of experiments we investigated how self-consistency affects the learning dynamics of our agents. First, we analysed the impact of self-consistency on the learned representation. In Figure 4a we compare (1) a policy-gradient baseline “PG” that does not use a model (2) learning a value-equivalent model purely as an auxiliary task “+VE Aux”, and (3)
augmenting this auxiliary task with our self-consistency objective “+SC Aux”. When we configure Muesli to use the model purely as an auxiliary there is no MPO-like loss term, and the value function bootstrap uses state values rather than model-based Q values. We found that self-consistency updates were helpful, even though the model was neither used for value estimation nor policy improvement; this suggests that self-consistency may aid representation learning.
Information flow. Self-consistency updates, unlike Dyna, move information between the reward model, transition model, and value function. To study this effect, we returned to tabular policy evaluation. We analysed the robustness of self consistency to initial additive Gaussian noise in the reward model. Figure 5 (left) shows the area under the value error curve as a function of the noise. With worse initialisation the effectiveness of the self-consistency update deteriorates more rapidly than the Dyna update, especially when learning with MLE. The reward model is learned in a grounded manner, so will receive the same updates with Dyna or SC. In the Dyna case, the poor reward initialisation can only pollute the value function. With SC, information can also flow from the impaired reward model into the dynamics model, and from there damage the value function further in turn. A different effect can be seen when varying the value initialisation by adding noise to the true value, as shown in Figure 5 (right). Since the value is updated by SC, it is possible for a poorly initialised value to be fixed more rapidly if SC is effective. Indeed, we see a slight trend towards a greater advantage for SC when value initialisation is worse. However, this is not a guarantee; certain poor initialisations could lead to self-reinforcing errors with SC, which is reflected in the overlapping confidence intervals.
Scaling with computation. Self-consistency updates, like Dyna, allow to use additional computation without requiring further interactions in the environment. In Figure 4b, we show the effect of rolling out different numbers of imagined trajectories to estimate the self-consistency loss, by randomly choosing a fraction of the states in each batch to serve as starting points (for ratios greater than one we randomly duplicated starting states). In the domains where self-consistency was beneficial, the potential gains often increased with the sampling of additional starting-states. However, in most cases saturation was reached when one imagined trajectory was sampled for each state in the batch (ratio=1 was the default in the experiments described above).

heading: 5 Related work

text:
Model learning can be useful to RL agents in various ways, such as: (i) aiding representation learning [46, 28, 33, 20, 25] (ii) planning for policy optimisation and/or value learning [55, 51, 21, 10]; (iii) action selection via local planning [45, 49]. See Moerland et al. [36] or Hamrick et al. [24] for a survey of how models are used in RL. Many approaches to learning the model have been investigated in the literature, usually based on the principle of maximum likelihood estimation (MLE) [32, 51, 38, 21, 31, 10, 22]. Hafner et al. [23] notably showed that a deep generative model of Atari games can be used for Dyna-style learning of a policy and value only in the imagination of the model, performing competitively with model-free methods.
RL-informed objectives have been used in some recent approaches to learn implicit or explicit models that focus on the aspects of the environment that are relevant for control [53, 50, 39, 16, 30, 47, 18, 44, 25]. This type of model learning is connected to theoretical investigations of “value-aware” model learning [15], which discusses an equivalence of Bellman operators applied to worst-case functions from a hypothesis space; “iterative value-aware model learning” [14] which, closer to our setting, uses an equivalence of Bellman operators applied to a series of agent value-functions; and the “value equivalence” (VE) principle [19], which describes equivalences between Bellman operators applied to sets of arbitrary functions.
All the approaches to model learning discussed above leverage observed real experience for learning a model of the environment. In contrast, our self-consistency principle provides a mechanism for model learning using imagined experience. This is related to the work of Silver et al. [50], where self-consistency updates were used to aid learning of temporally abstract models of a Markov reward process. Filos et al. [17] also used self-consistency losses, for offline multi-task inverse RL. Selfconsistency has been applied in other areas of machine learning as well. Consider, for instance, back-translation [8, 12] in natural language processing, or CycleGANs [58] in generative modelling.
In this paper we considered models that can be rolled forward. It is also possible to consider backward models, that can be used to assign credit back in time [54, 11]. The principle of self-consistency we introduce in this work can be extended to these kinds of models fairly straightforwardly. A different type of self-consistency was studied concurrently by Yu et al. [56], who develop a model-based system where a forward and backward model are trained to be cyclically consistent with each other.

heading: 6 Conclusion

text:
We introduced the idea of self-consistent models and values. Our approach departs from classic planning paradigms where a model is learned to be consistent with the environment, and values are learned to be consistent with the model. Amongst possible variants, we identified as particularly promising an update modelled on semi-gradient temporal difference objectives. We found this update to be effective both in tabular settings and with function approximation. Self-consistency updates proved particularly well suited to deep value-equivalent model-based agents, where traditional algorithms such as Dyna were found to perform poorly. The self-consistency objectives discusses in this paper are based on a form of policy-evaluation; future work may investigate extensions to enable value-equivalent model-based agents to perform policy improvement for states in arbitrary trajectories drawn from the model.
Acknowledgments and Disclosure of Funding
We would like to thank Ivo Danihelka, Junhyuk Oh, Iurii Kemaev, and Thomas Hubert for valuable discussions and comments on the manuscript. Thanks also to the developers of Jax [9] and the DeepMind Jax ecosystem [3] which were invaluable to this project. The authors received no specific funding for this work.

title: Self-Consistent Models and Values


summaries: This paper investigates the effect of episode difficulty and sampling schemes on the performance of few-shot classification algorithms. It defines difficulty based on the negative log predictive likelihood of the few-shot learner and proposes to use importance sampling to mimic various sampling strategies. Experiments demonstrate that episodic difficulty is (a) approximately normally distributed, (b) tends to transfer across architectures and learning algorithms, (c) tends to be consistent over the course of training, and (d) that a uniform sampling scheme over difficulty tends to produce the best results. An online scheme for computing importance weights is also proposed.   === Post-rebuttal ===  I have read the authors' response and other reviews and will maintain my original rating.

summaries: This paper explores episode/task difficulty in few-shot learning. The task difficulty is defined as the negative log-likelihood over the query set given the support set and model parameters. The authors find evidence that episode difficulty likely follows a normal distribution and proposes a new method, called UNIFORM, based on importance sampling. The new method achieves similar or better performance than other sampling procedures (EASY, HARD, and CURRICULUM). The paper also presents two variants of UNIFORM: offline and online.

summaries: This submission studies episodic sampling, which is used to sample datasets (consisting of support and query sets) for training meta-learning models. The authors define a notion of episodic difficulty (the negative log-likelihood of the query set given the support set and model parameters) and propose using importance sampling to train using several different types of episode-sampling strategies for meta-training. By using importance sampling, there is no sampling required for the target distribution but just weighting the episode loss based on the proposal and target distribution. They conduct experiments both to confirm some of their assumptions that are critical for their proposal (that episodic difficulty is normally distributed) and to show how their episodic sampling strategies compare against random episodic sampling, which is the dominant strategy used by work in this area.

summaries: Meta-learning (few-shot learning) algorithms, in general, equally treat all tasks for the training process. Tasks are sampled uniformly at random. This paper proposes a novel meta-learning framework that controls the task sampling probabilities according to difficulties. This paper measures the difficulties of tasks using the negative log-likelihood and leverages the difficulties to introduce different weights to tasks. The main idea of the learning framework is that the prior of the difficulties follows a normal distribution while the target sampling distribution could be uniform or some other interesting forms. Then, important sampling provides different weights to different tasks. From experiments, the uniform target distribution shows good results.

abstractText: Episodic training is a core ingredient of few-shot learning to train models on tasks with limited labelled data. Despite its success, episodic training remains largely understudied, prompting us to ask the question: what is the best way to sample episodes? In this paper, we first propose a method to approximate episode sampling distributions based on their difficulty. Building on this method, we perform an extensive analysis and find that sampling uniformly over episode difficulty outperforms other sampling schemes, including curriculum and easy-/hard-mining. As the proposed sampling method is algorithm agnostic, we can leverage these insights to improve few-shot learning accuracies across many episodic training algorithms. We demonstrate the efficacy of our method across popular few-shot learning datasets, algorithms, network architectures, and protocols.

authors:
S.M.R. Arnold
P. Mahajan
D. Datta
I. Bunner
K.S. Zarkias

title: learn2learn: A library for Meta- Learning research

venue: Aug.

authors:
Y. Bengio
J. Louradour
R. Collobert
J. Weston

title: Curriculum learning

venue: Proceedings of the 26th Annual International Conference on Machine Learning, ICML ’09, pages 41–48, New York, NY, USA, June

authors:
L. Bertinetto
J.F. Henriques
P. Torr
A. Vedaldi

title: Meta-learning with differentiable closed-form solvers

venue: Sept.

authors:
D.P. Bertsekas
J.N. Tsitsiklis

title: Neuro-Dynamic programming

authors:
L. Bottou
F.E. Curtis
J. Nocedal

title: Optimization methods for large-scale machine learning

venue: June

authors:
M. Botvinick
S. Ritter
J.X. Wang
Z. Kurth-Nelson
C. Blundell
D. Hassabis

title: Reinforcement learning, fast and slow

venue: Trends in cognitive sciences, 23(5):408–422,

authors:
W.-L. Chao
H.-J. Ye
D.-C. Zhan
M. Campbell
K.Q. Weinberger

title: Revisiting meta-learning as supervised learning

venue: arXiv preprint arXiv:2002.00573,

authors:
M. Cimpoi
S. Maji
I. Kokkinos
S. Mohamed
A. Vedaldi

title: Describing textures in the wild

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

authors:
D. Csiba
P. Richtárik

title: Importance sampling for minibatches

authors:
J. Deng
W. Dong
R. Socher
L.-J. Li
K. Li
L. Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. Ieee,

authors:
G.S. Dhillon
P. Chaudhari
A. Ravichandran
S. Soatto

title: A baseline for few-shot image classification

venue: International Conference on Learning Representations,

authors:
A. Doucet
N. d. Freitas
N. Gordon
editors

title: Sequential Monte Carlo Methods in Practice

authors:
A. Fallah
A. Mokhtari
A. Ozdaglar

title: Generalization of model-agnostic meta-learning algorithms: Recurring and unseen tasks

venue: arXiv preprint arXiv:2102.03832,

authors:
L. Fei-Fei
R. Fergus
P. Perona

title: One-shot learning of object categories

venue: IEEE Trans. Pattern Anal. Mach. Intell., 28(4):594–611, Apr.

authors:
C. Finn
P. Abbeel
S. Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: International Conference on Machine Learning, pages 1126–1135. PMLR,

authors:
S. Flennerhag
A.A. Rusu
R. Pascanu
H. Yin
R. Hadsell

title: Meta-Learning with warped gradient descent

authors:
M.P. Friedlander
M. Schmidt

title: Hybrid Deterministic-Stochastic methods for data fitting

venue: SIAM J. Sci. Comput., 34(3):A1380–A1405, Jan.

authors:
M. Garnelo
D. Rosenbaum
C. Maddison
T. Ramalho
D. Saxton
M. Shanahan
Y.W. Teh
D. Rezende
S.M.A. Eslami

title: Conditional neural processes

venue: J. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1704–1713, Stockholmsmässan, Stockholm Sweden,

authors:
G. Ghiasi
T.-Y. Lin
Q.V. Le

title: Dropblock: a regularization method for convolutional networks

venue: Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 10750–10760,

authors:
S. Gidaris
N. Komodakis

title: Dynamic few-shot visual learning without forgetting

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

title: Monte Carlo methods in financial engineering

venue: Springer, New York,

authors:
E. Grant
C. Finn
S. Levine
T. Darrell
T. Griffiths

title: Recasting Gradient-Based Meta-Learning as hierarchical bayes

venue: Jan.

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

title: Deep residual learning for image recognition

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778,

authors:
T. Hospedales
A. Antoniou
P. Micaelli
others

title: Meta-learning in neural networks: A survey

venue: arXiv preprint arXiv,

authors:
S. Ioffe
C. Szegedy

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

venue: International conference on machine learning, pages 448–456. PMLR,

authors:
T.B. Johnson
C. Guestrin

title: Training deep models faster with robust, approximate importance sampling

venue: S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 7265–7275. Curran Associates, Inc.,

authors:
A. Katharopoulos
F. Fleuret

title: Not all samples are created equal: Deep learning with importance sampling

venue: Mar.

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization

venue: ICLR (Poster),

title: A note on importance sampling using standardized weights

venue: University of Chicago, Dept. of Statistics, Tech. Rep, 348,

authors:
S. Laenen
L. Bertinetto

title: On episodes, prototypical networks, and few-shot learning

venue: Dec.

authors:
B.M. Lake
T.D. Ullman
J.B. Tenenbaum
S.J. Gershman

title: Building machines that learn and think like people

venue: Behavioral and brain sciences, 40,

authors:
K. Lee
S. Maji
A. Ravichandran
S. Soatto

title: Meta-learning with differentiable convex optimization

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10657– 10665,

authors:
C. Liu
Z. Wang
D. Sahoo
Y. Fang
K. Zhang
S.C.H. Hoi

title: Adaptive task sampling for Meta-Learning

venue: July

title: Metropolized independent sampling with comparisons to rejection sampling and importance sampling

venue: Statistics and computing, 6(2):113–119,

authors:
S. Maji
E. Rahtu
J. Kannala
M. Blaschko
A. Vedaldi

title: Fine-grained visual classification of aircraft

venue: arXiv preprint arXiv:1306.5151,

authors:
E.G. Miller
N.E. Matsakis
P.A. Viola

title: Learning from one example through shared densities on transforms

venue: Proceedings IEEE Conference on Computer Vision and Pattern Recognition. CVPR 2000 (Cat. No.PR00662), volume 1, pages 464–471 vol.1, June

authors:
A. Nichol
J. Achiam
J. Schulman

title: On First-Order Meta-Learning algorithms

authors:
M.-E. Nilsback
A. Zisserman

title: A visual vocabulary for flower classification

venue: 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’06), volume 2, pages 1447–1454. IEEE,

authors:
B.N. Oreshkin
P. Rodriguez
A. Lacoste

title: Tadam: task dependent adaptive metric for improved few-shot learning

venue: Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 719–729,

authors:
E. Park
J.B. Oliva

title: Meta-Curvature

venue: Feb.

authors:
A. Raghu
M. Raghu
S. Bengio
O. Vinyals

title: Rapid learning or feature reuse? towards understanding the effectiveness of maml

authors:
A. Rajeswaran
C. Finn
S. Kakade
S. Levine

title: Meta-Learning with implicit gradients

authors:
S. Ravi
H. Larochelle

title: Optimization as a model for few-shot learning

authors:
A. Ravichandran
R. Bhotika
S. Soatto

title: Few-Shot learning with embedded class models and Shot-Free meta training

venue: May

authors:
M. Ren
E. Triantafillou
S. Ravi
J. Snell
K. Swersky
J.B. Tenenbaum
H. Larochelle
R.S. Zemel

title: Meta-learning for semi-supervised few-shot classification

authors:
H. Robbins
S. Monro

title: A stochastic approximation method

authors:
A. Santoro
S. Bartunov
M. Botvinick
D. Wierstra
T. Lillicrap

title: Meta-Learning with Memory- Augmented neural networks

venue: M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1842–1850, New York, New York, USA,

authors:
T. Schaul
J. Quan
I. Antonoglou
D. Silver

title: Prioritized experience replay

venue: Nov.

title: Evolutionary Principles in Self-Referential Learning

venue: PhD thesis,

authors:
S.S. Shapiro
M.B. Wilk

title: An analysis of variance test for normality (complete samples)

authors:
A. Shrivastava
A. Gupta
R. Girshick

title: Training region-based object detectors with online hard example mining

venue: Proceedings of the IEEE conference on computer vision and pattern recognition, pages 761–769,

authors:
P.J. Smith
M. Shafi
Hongsheng Gao

title: Quick simulation: a review of importance sampling techniques in communications systems

venue: IEEE J. Sel. Areas Commun.,

authors:
J. Snell
K. Swersky
R. Zemel

title: Prototypical networks for few-shot learning

venue: Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 4080–4090,

authors:
N. Srivastava
G. Hinton
A. Krizhevsky
I. Sutskever
R. Salakhutdinov

title: Dropout: a simple way to prevent neural networks from overfitting

venue: The journal of machine learning research, 15(1):1929–1958,

authors:
Q. Sun
Y. Liu
Z. Chen
T.-S. Chua
B. Schiele

title: Meta-Transfer learning through hard tasks

venue: Oct.

authors:
S. Thrun
L. Pratt
editors

title: Learning to Learn

authors:
E. Triantafillou
T. Zhu
V. Dumoulin
P. Lamblin
U. Evci
K. Xu
R. Goroshin
C. Gelada
K. Swersky
P.-A. Manzagol
H. Larochelle

title: Meta-dataset: A dataset of datasets for learning to learn from few examples

authors:
O. Vinyals
C. Blundell
T. Lillicrap
K. Kavukcuoglu
D. Wierstra

title: Matching networks for one shot learning

venue: Proceedings of the 30th International Conference on Neural Information Processing Systems, pages 3637–3645,

authors:
Y. Wang
Q. Yao
J.T. Kwok
L.M. Ni

title: Generalizing from a few examples: A survey on few-shot learning

venue: ACM Comput. Surv., 53(3):1–34, June

authors:
P. Welinder
S. Branson
T. Mita
C. Wah
F. Schroff
S. Belongie
P. Perona

title: Caltech-ucsd birds

title: Updating mean and variance estimates: An improved method

authors:
C. Wu
R. Manmatha
A.J. Smola
P. Krähenbühl

title: Sampling matters in deep embedding learning

venue: 2017 IEEE International Conference on Computer Vision (ICCV), pages 2859–2867, Oct.

authors:
X. Wu
E. Dyer
B. Neyshabur

title: When do curricula work

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

authors:
H.-J. Ye
H. Hu
D.-C. Zhan
F. Sha

title: Few-shot learning via embedding adaptation with set-toset functions

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

authors:
J. Yoon
T. Kim
O. Dia
S. Kim
Y. Bengio
S. Ahn

title: Bayesian Model-Agnostic Meta- Learning

venue: S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 7332–7342. Curran Associates, Inc.,

authors:
C. Zhang
C. Öztireli
S. Mandt
G. Salvi

title: Active mini-batch sampling using repulsive point processes

venue: Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 5741–5748,

authors:
C. Zhang
Y. Cai
G. Lin
C. Shen

title: Deepemd: Few-shot image classification with differentiable earth mover’s distance and structured classifiers

venue: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12203–12213,

authors:
P. Zhao
T. Zhang

title: Accelerating minibatch stochastic gradient descent using stratified sampling

authors:
Y. Zhou
Y. Wang
J. Cai
Y. Zhou
Q. Hu
W. Wang

title: Expert training: Task hardness aware Meta- Learning for Few-Shot classification

heading: 1 Introduction

text:
Large amounts of high-quality data have been the key for the success of deep learning algorithms. Furthermore, factors such as data augmentation and sampling affect model performance significantly. Continuously collecting and curating data is a resource (cost, time, storage, etc.) intensive process. Hence, recently, the machine learning community has been exploring methods for performing transfer-learning from large datasets to unseen tasks with limited data.
A popular genre of these approaches is called meta-learning few-shot approaches, where, in addition to the limited data from the task of interest, a large dataset of disjoint tasks is available for (pre-)training. These approaches are prevalent in the area of computer vision [31] and reinforcement learning [6]. A key component of these methods is the notion of episodic training, which refers to sampling tasks from the larger dataset for training. By learning to solve these tasks correctly, the model can generalize to new tasks.
However, sampling for episodic training remains surprisingly understudied despite numerous methods and applications that build on it. To the best of our knowledge, only a handful of works [69, 55, 33] explicitly considered the consequences of sampling episodes. In comparison, stochastic [46] and mini-batch [4] sampling alternatives have been thoroughly analyzed from the perspectives of optimization [17, 5], information theory [27, 9], and stochastic processes [68, 66], among many others. Building a similar understanding of sampling for episodic training will help theoreticians and practitioners develop improved sampling schemes, and is thus of crucial importance to both.
In this paper, we explore many sampling schemes to understand their impact on few-shot methods. Our work revolves around the following fundamental question: what is the best way to sample episodes? Our focus will be restricted to image classification in the few-shot learning setting – where “best” is taken to mean “higher transfer accuracy of unseen episodes” – and leave analyses and applications to other areas for future work. ∗Equal contributions; work done while at Amazon Web Services.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Contrary to prior work, our experiments indicate that sampling uniformly with respect to episode difficulty yields higher classification accuracy – a scheme originally proposed to regularize metric learning [62]. To better understand these results, we take a closer look at the properties of episodes and what makes them difficult. Building on this understanding, we propose a method to approximate different sampling schemes, and demonstrate its efficacy on several standard few-shot learning algorithms and datasets.
Concretely, we make the following contributions:
• We provide a detailed empirical analysis of episodes and their difficulty. When sampled randomly, we show that episode difficulty (approximately) follows a normal distribution and that the difficulty of an episode is largely independent of several modeling choices including the training algorithm, the network architecture, and the training iteration. • Leveraging our analysis, we propose simple and universally applicable modifications to the episodic sampling pipeline to approximate any sampling scheme. We then use this scheme to thoroughly compare episode sampling schemes – including easy/hard-mining, curriculum learning, and uniform sampling – and report that sampling uniformly over episode difficulty yields the best results. • Finally, we show that sampling matters for few-shot classification as it improves transfer
accuracy for a diverse set of popular [53, 20, 15, 41] and state-of-the-art [64] algorithms on standard and cross-domain benchmarks.

heading: 2 Preliminaries

heading: 2.1 Episodic sampling and training

text:
We define episodic sampling as subsampling few-shot tasks (or episodes) from a larger base dataset [7]. Assuming the base dataset admits a generative distribution, we sample an episode in two steps2. First, we sample the episode classes Cτ from class distribution p(Cτ ); second, we sample the episode’s data from data distribution p(x, y | Cτ ) conditioned on Cτ . This gives rise to the following log-likelihood for a model lθ parameterized by θ:
L(θ) = E τ∼q(·) [log lθ(τ)] , (1)
where q(τ) is the episode distribution induced by first sampling classes, then data. In practice, this expectation is approximated by sampling a batch of episodes B each with their set τQ of query samples. To enable transfer to unseen classes, it is also common to include a small set τS of support samples to provide statistics about τ . This results in the following Monte-Carlo estimator:
L(θ) ≈ 1 |B| ∑ τ∈B 1 |τQ| ∑
(x,y)∈τQ
log lθ(y | x, τS), (2)
where the data in τQ and τS are both distributed according to p(x, y | Cτ ). In few-shot classification, the n-way k-shot setting corresponds to sampling n = |Cτ | classes, each with k = |τS |n support data points. The implicit assumption in the vast majority of few-shot methods is that both classes and data are sampled with uniform probability – but are there better alternatives? We carefully examine these underlying assumptions in the forthcoming sections.

heading: 2.2 Few-shot algorithms

text:
We briefly present a few representative episodic learning algorithms. A more comprehensive treatment of few-shot algorithms is presented in Wang et al. [59] and Hospedales et al. [24]. A core question in few-shot learning lies in evaluating (and maximizing) the model likelihood lθ. These algorithms can be divided in two major families: gradient-based methods, which adapt the model’s parameters to the episode; and metric-based methods, which compute similarities between support and query samples in a learned embedding space.
Gradient-based few-shot methods are best illustrated through Model-Agnostic Meta-Learning [15] (MAML). The intuition behind MAML is to learn a set of initial parameters which can quickly
2We use the notation for a probability distribution and its probability density function interchangeably.
specialize to the task at hand. To that end, MAML computes lθ by adapting the model parameters θ via one (or more) steps of gradient ascent and then computes the likelihood using the adapted parameters θ′. Concretely, we first compute the likelihood pθ(y | x) using the support set τS , adapt the model parameters, and then evaluate the likelihood:
lθ(y | x, τS) = pθ′(y | x) s.t. θ′ = θ + α∇θ ∑
(x,y)∈τS
log pθ(y | x),
where α > 0 is known as the adaptation learning rate. A major drawback of training with MAML lies in back-propagating through the adaptation phase, which requires higher-order gradients. To alleviate this computational burden, Almost No Inner Loop [41] (ANIL) proposes to only adapt the last classification layer of the model architecture while tying the rest of the layers across episodes. They empirically demonstrate little classification accuracy drop while accelerating training times four-fold.
Akin to ANIL, metric-based methods also share most of their parameters across tasks; however, their aim is to learn a metric space where classes naturally cluster. To that end, metric-based algorithms learn a feature extractor φθ parameterized by θ and classify according to a non-parametric rule. A representative of this family is Prototypical Network [53] (ProtoNet), which classifies query points according to their distance to class prototypes – the average embedding of a class in the support set:
lθ(y | x, τS) = exp (−d(φθ(x), φyθ))∑ y′∈Cτ exp ( −d(φθ(x), φy ′ θ ) ) s.t. φcθ = 1k ∑
(x,y)∈τS y=c
φθ(x),
where d(·, ·) is a distance function such as the Euclidean distance or the negative cosine similarity, and φcθ is the class prototype for class c. Other classification rules include support vector clustering [32], neighborhood component analysis [30], and the earth-mover distance [67].

heading: 2.3 Episode difficulty

text:
Given an episode τ and likelihood function lθ, we define episode difficulty to be the negative loglikelihood incurred on that episode:
Ωlθ (τ) = − log lθ(τ), which is a surrogate for how hard it is to classify the samples in τQ correctly, given lθ and τS . By definition, this choice of episode difficulty is tied to the choice of the likelihood function lθ.
Dhillon et al. [11] use a similar surrogate as a means to systematically report few-shot performances. We use this definition because it is equivalent to the loss associated with the likelihood function lθ on episode τ , which is readily available at training time.

heading: 3 Methodology

text:
In this section, we describe the core assumptions and methodology used in our study of sampling methods for episodic training. Our proposed method builds on importance sampling [21] (IS) which we found compelling for three reasons: (i) IS is well understood and solidly grounded from a theoretical standpoint, (ii) IS is universally applicable thus compatible with all episodic training algorithms, and (iii) IS is simple to implement with little requirement for hyper-parameter tuning.
Why should we care about episodic sampling? A back-of-the-envelope calculation3 suggests that there are on the order of 10162 different training episodes for the smallest-scale experiments in Section 5. Since iterating through each of them is infeasible, we ought to express some preference over which episodes to sample. In the following, we describe a method that allows us to specify this preference.

heading: 3.1 Importance sampling for episodic training

text:
Let us assume that the sampling scheme described in Section 2.1 induces a distribution q(τ) over episodes. We call it the proposal distribution, and assume knowledge of its density function. We wish
3For a base dataset with N classes and K input-output pairs per class, there are a total of ( N n )( K k )n possible
episodes that can be created when sampling k pairs each from n classes.
Algorithm 1: Episodic training with Importance Sampling Input: target (p) and proposal (q) distributions, likelihood function lθ, optimizer OPT. Randomly initialize model parameters θ. repeat
Sample a mini-batch B of episodes from q(τ). for each episode τ in mini-batch B do
Compute episode likelihood: lθ(τ). Compute importance weight: w(τ) = p(τ)q(τ) .
end for Aggregate: L(θ)← ∑ τ∈B w(τ) log lθ(τ). Compute effective sample size ESS(B). Update model parameters: θ ← OPT( L(θ)ESS(B) ).
until parameters θ have converged.
to estimate the expectation in Eq. (1) when sampling episodes according to a target distribution p(τ) of our choice, rather than q(τ). To that end, we can use an importance sampling estimator which simply re-weights the observed values for a given episode τ by w(τ) = p(τ)q(τ) , the ratio of the target and proposal distributions:
E τ∼p(·) [log lθ(τ)] = E τ∼q(·) [w(τ) log lθ(τ)] .
The importance sampling identity holds whenever q(τ) has non-zero density over the support of p(τ), and effectively allows us to sample from any target distribution p(τ).
A practical issue of the IS estimator arises when some values of w(τ) become much larger than others; in that case, the likelihoods lθ(τ) associated with mini-batches containing heavier weights dominate the others, leading to disparities. To account for this effect, we can replace the mini-batch average in the Monte-Carlo estimate of Eq. (2) by the effective sample size ESS(B) [29, 34]:
E τ∼p(·)
[log lθ(τ)] ≈ 1 ESS(B) ∑ τ∈B w(τ) log lθ(τ) s.t. ESS(B) = ( ∑ τ∈B w(τ)) 2∑ τ∈B w(τ) 2 , (3)
where B denotes a mini-batch of episodes sampled according to q(τ). Note that whenw(τ) is constant, we recover the standard mini-batch average setting as ESS(B) = |B|. Empirically, we observed that normalizing with the effective sample size avoided instabilities. This method is summarized in Algorithm 1.

heading: 3.2 Modeling the proposal distribution

text:
A priori, we do not have access to the proposal distribution q(τ) (nor its density) and thus need to estimate it empirically. Our main assumption is that sampling episodes from q(τ) induces a normal distribution over episode difficulty. With this assumption, we model the proposal distribution by this induced distribution, therefore replacing q(τ) with N (Ωlθ (τ) | µ, σ2) where µ, σ2 are the mean and variance parameters. As we will see in Section 5.2, this normality assumption is experimentally supported on various datasets, algorithms, and architectures.
We consider two settings for the estimation of µ and σ2: offline and online. The offline setting consists of sampling 1, 000 training episodes before training, and computing µ, σ2 using a model pre-trained on the same base dataset. Though this setting seems unrealistic, i.e. having access to a pre-trained model, several meta-learning few-shot methods start with a pre-trained model which they further build upon. Hence, for such methods there is no overhead. For the online setting, we estimate the parameters on-the-fly using the model currently being trained. This is justified by the analysis in Section 5.2 which shows that episode difficulty transfers across model parameters during training. We update our estimates of µ, σ2 with an exponential moving average:
µ← λµ+ (1− λ)Ωlθ (τ) and σ2 ← λσ2 + (1− λ)(Ωlθ (τ)− µ)2, where λ ∈ [0, 1] controls the adjustment rate of the estimates, and the initial values of µ, σ2 are computed in a warm-up phase lasting 100 iterations. Keeping λ = 0.9 worked well for all our
experiments (Section 5). We opted for this simple implementation as more sophisticated approaches like West [61] yielded little to no benefit.

heading: 3.3 Modeling the target distribution

text:
Similar to the proposal distribution, we model the target distribution by its induced distribution over episode difficulty. Our experiments compare four different approaches, all of which share parameters µ, σ2 with the normal model of the proposal distribution. For numerical stability, we truncate the support of all distributions to [µ− 2.58σ, µ+ 2.58σ], which gives approximately 99% coverage for the normal distribution centered around µ.
The first approach (HARD) takes inspiration from hard negative mining [51], where we wish to sample only more challenging episodes. The second approach (EASY) takes a similar view but instead only samples easier episodes. We can model both distributions as follows:
U(Ωlθ (τ) | µ, µ+ 2.58σ) (HARD) and
U(Ωlθ (τ) | µ− 2.58σ, µ) (EASY)
where U denotes the uniform distribution. The third (CURRICULUM) is motivated by curriculum learning [2], which slowly increases the likelihood of sampling more difficult episodes:
N (Ωlθ (τ) | µt, σ2) (CURRICULUM)
where µt is linearly interpolated from µ− 2.58σ to µ+ 2.58σ as training progresses. Finally, our fourth approach, UNIFORM, resembles distance weighted sampling [62] and consists of sampling uniformly over episode difficulty:
U(Ωlθ (τ) | µ− 2.58σ, µ+ 2.58σ). (UNIFORM)
Intuitively, UNIFORM can be understood as a uniform prior over unseen test episodes, with the intention of performing well across the entire difficulty spectrum. This acts as a regularizer, forcing the model to be equally discriminative for both easy and hard episodes.

heading: 4 Related Works

text:
This paper studies task sampling in the context of few-shot [36, 14] and meta-learning [49, 56].
Few-shot learning. This setting has received a lot of attention over recent years [58, 43, 47, 18]. Broadly speaking, state-of-the-art methods can be categorized in two major families: metric-based and gradient-based.
Metric-based methods learn a shared feature extractor which is used to compute the distance between samples in embedding space [53, 3, 44, 30]. The choice of metric mostly differentiates one method from another; for example, popular choices include Euclidean distance [53], negative cosine similarity [20], support vector machines [32], set-to-set functions [64], or the earth-mover distance [67].
Gradient-based algorithms such as MAML [15], propose an objective to learn a network initialization that can quickly adapt to new tasks. Due to its minimal assumptions, MAML has been extended to probabilistic formulations [22, 65] to incorporate learned optimizers – implicit [16] or explicit [40] – and simplified to avoid expensive second-order computations [37, 42]. In that line of work, ANIL [41] claims to match MAML’s performance when adapting only the last classification layer – thus greatly reducing the computational burden and bringing gradient and metric-based methods closer together.
Sampling strategies. Sampling strategies have been studied for different training regimes. Wu et al. [62] demonstrate that “sampling matters” in the context of metric learning. They propose to sample a triplet with probability proportional to the distance of its positive and negative samples, and observe stabilized training and improved accuracy. This observation was echoed by Katharopoulos and Fleuret [27] when sampling mini-batches: carefully choosing the constituent samples of a minibatch improves the convergence rate and asymptotic performance. Like ours, their method builds on importance sampling [52, 12, 26] but whereas they compute importance weights using the magnitude
of the model’s gradients, we use the episode’s difficulty. Similar insights were also observed in reinforcement learning, where Schaul et al. [48] suggests a scheme to sample transitions according to the temporal difference error.
Closer to our work, Sun et al. [55] present a hard-mining scheme where the most challenging classes across episodes are pooled together and used to create new episodes. Observing that the difficulty of a class is intrinsically linked to the other classes in the episode, Liu et al. [33] propose a mechanism to track the difficulty across every class pair. They use this mechanism to build a curriculum [2, 63] of increasingly difficult episodes. In contrast to these two approaches, our proposed method makes use of importance sampling to mimic the target distribution rather than sampling from it directly. This helps achieve fast and efficient sampling without any preprocessing requirements.

heading: 5 Experiments

text: We first validate the assumptions underlying our proposed IS estimator and shed light on the properties of episode difficulty. Then, we answer the question we pose in the introduction, namely: what is the best way to sample episodes? Finally, we ask if better sampling improves few-shot classification.

heading: 5.1 Experimental setup

text:
We review the standardized few-shot benchmarks and provide a detailed description in Appendix A.
Datasets. We use two standardized image classification datasets, Mini-ImageNet [58] and TieredImageNet [45], both subsets of ImageNet [10]. Mini-ImageNet consists of 64 classes for training, 16 for validation, and 20 for testing; we use the class splits introduced by Ravi and Larochelle [43]. Tiered-ImageNet contains 608 classes split into 351, 97, and 160 for training, validation, and testing, respectively.
Network architectures. We train two model architectures. A 4-layer convolution network conv(64)4 Vinyals et al. [58] with 64 channels per layer. And ResNet-12, a 12-layer deep residual network [23] introduced by Oreshkin et al. [39]. Both architectures use batch normalization [25] after every convolutional layer and ReLU as the non-linearity.
Training algorithms. For the metric-based family, we use ProtoNet with Euclidean [53] and scaled negative cosine similarity measures [20]. Additionally, we use MAML [15] and ANIL [41] as representative gradient-based algorithms.
Hyper-parameters. We tune hyper-parameters for each algorithm and dataset to work well across different few-shot settings and network architectures. Additionally, we keep the hyper-parameters the same across all different sampling methods for a fair comparison. We train for 20k iterations with a mini-batch of size 16 and 32 for Mini-ImageNet and Tiered-ImageNet respectively, and validate every 1k iterations on 1k episodes. The best performing model is finally evaluated on 1k test episodes.

heading: 5.2 Understanding episode difficulty

text: All the models in this subsection are trained using baseline sampling as described in Section 2.1, i.e., episodic training without importance sampling.

heading: 5.2.1 Episode difficulty is approximately normally distributed

text:
We begin our analysis by verifying that the distribution over episode difficulty induced by q(τ) is approximately normal. In Fig. 1, we use the difficulty of 10k test episodes sampled with q(τ). The difficulties are computed using conv(64)4 trained with ProtoNet and MAML on Mini-ImageNet for 1-shot 5-way classification. The episode difficulty density plots follow a bell curve, which are naturally modeled with a normal distribution. The Q-Q plots, typically used to assess normality, suggest the same – the closer the curve is to the identity line, the closer the distribution is to a normal.
Finally, we compute the Shapiro-Wilk test for normality [50], which tests for the null hypothesis that the data is drawn from a normal distribution. Since the p-value for this test is sensitive to the
sample size4, we subsample 50 values 100 times and average rejection rates over these subsets. With α = 0.05, the null hypothesis is rejected 14% and 17% of the time for Mini-ImageNet and Tiered-ImageNet respectively, thus suggesting that episode difficulty can be reliably approximated with a normal distribution.

heading: 5.2.2 Independence from modeling choices

text:
By definition, the notion of episode difficulty is tightly coupled to the model likelihood lθ (Section 2.3), and hence to the modeling variables such as learning algorithm, network architecture, and model parameters. We check if episode difficulty transfers across different choices for these variables.
1.0
1.5
2.0
2.5
Pr ot
ot yp
ica l N
et wo rk (E uc lid ea n)
1.0
1.5
2.0
2.5
Pr ot
ot yp
ica l N
et wo rk (c os in e)
1.0
1.5
2.0
2.5
M AM
L
1 2 Prototypical Network
(Euclidean)
1.0
1.5
2.0
2.5
AN IL
1 2 Prototypical Network
(cosine)
1 2 MAML
1 2 3 4 ANIL
Network architecture. Next we analyze the dependence on the network architecture. We trained conv(64)4 and ResNet-12’s using all training algorithms for Mini-ImageNet 1-shot 5-way classi-
4For a large sample size, the p-values are not reliable as they may detect trivial departures from normality.
fication. We compute the episode difficulties for 10k test episodes and compute their Spearman rank-order correlation coefficients across the two architectures, for a given algorithm. The correlation coefficients are 0.57 for ProtoNet (Euclidean), 0.72 for ProtoNet (cosine), 0.58 for MAML, and 0.49 for ANIL. Fig. 3 illustrates this positive correlation, suggesting that episode difficulty is transferred across network architectures with high probability.
Model parameters during training. Lastly, we study the dependence on model parameters during training. We select the 50 easiest and 50 hardest episodes, i.e., episodes with the lowest and highest difficulties respectively, from 1k test episodes. We track the episode difficulty for all episodes over the training phase and visualize the trend in Fig. 4 for conv(64)4’s trained using different training algorithms on Mini-ImageNet (1-shot 5-way). Throughout training, easy episodes remain easy and hard episodes remain hard, hence suggesting that episode difficulty transfers across different model parameters during training. Since the episode difficulty does not change drastically during the training process, we can estimate it with a running average over the training iterations. This justifies the online modeling of the proposal distribution in Section 3.2.

heading: 5.3 Comparing episode sampling methods

text:
We compare different methods for episode sampling. To ensure fair comparisons, we use the offline formulation (Section 3.2) so that all sampling methods share the same pre-trained network (the network trained using baseline sampling) when computing proposal likelihoods. We compute results over 2 datasets, 2 network architectures, 4 algorithms and 2 few-shot protocols, totaling in 24 scenarios. Table 1 presents results on ProtoNet (cosine), while the rest are in Appendix D.
We observe that, although not strictly dominant, UNIFORM tends to outperform other methods as it is within the statistical confidence of the best method in 19/24 scenarios. For the 5/24 scenarios
where UNIFORM underperforms, it closely trails behind the best methods. Compared to baseline sampling, the average degradation of UNIFORM is −0.83% and at most −1.44% (ignoring the standard deviations) in 4/24 scenarios. Conversely, UNIFORM boosts accuracy by as much as 8.74% and on average by 3.86% (ignoring the standard deviations) in 13/24 scenarios. We attribute this overall good performance to the fact that uniform sampling puts a uniform distribution prior over the (unseen) test episodes, with the intention of performing well across the entire difficulty spectrum. This acts as a regularizer, forcing the model to be equally discriminative for easy and hard episodes. If we knew the test episode distribution, upweighting episodes that are most likely under that distribution will improve transfer accuracy [13]. However, this uninformative prior is the safest choice without additional information about the test episodes.
Second best is baseline sampling as it is statistically competitive on 10/24 scenarios, while EASY, HARD, and CURRICULUM only appear among the better methods in 4, 4, and 9 scenarios, respectively.

heading: 5.4 Online approximation of the proposal distribution

text:
Although the offline formulation is better suited for analysis experiments, it is expensive as it requires a pre-training phase for the proposal network and 2 forward passes during episodic training (one for the episode loss, another for the proposal density). In this subsection, we show that the online formulation faithfully approximates offline sampling and can retain most of the performance improvements from UNIFORM. We take the same 24 scenarios as in the previous subsection, and compare baseline sampling against offline and online UNIFORM. Table 2 reports the full suite of results.
We observe that baseline is statistically competitive on 9/24 scenarios; on the other hand, offline and online UNIFORM perform similarly on aggregate, as they are within the best results in 16/24 and 15/24 scenarios respectively. Similar to its offline counterpart, online UNIFORM does better than or comparable to baseline sampling in 21 out of 24 scenarios. On the 3/24 scenarios where online UNIFORM underperforms compared to baseline, the average degradation is −0.92%, and at most −1.26% (ignoring the standard deviations). Conversely, in the remaining scenarios, it boosts accuracy by as much as 6.67% and on average by 2.24% (ignoring the standard deviations). Therefore, using online UNIFORM, while computationally comparable to baseline sampling, results in a boost in few-shot performance; when it underperforms, it trails closely. We also compute the mean accuracy difference between the offline and online formulation, which is 0.07%± 0.35 accuracy points. This confirms that both the offline and online methods produce quantitatively similar outcomes.

heading: 5.5 Better sampling improves cross-domain transfer

text: To further validate the role of episode sampling as a way to improve generalization, we evaluate the models trained in the previous subsection on episodes from completely different domains. Specifically, we use the models trained on Mini-ImageNet with baseline and online UNIFORM sampling to evaluate on the test episodes of CUB-200 [60], Describable Textures [8], FGVC Aircrafts [35], and VGG Flowers [38], following the splits of Triantafillou et al. [57]. Table 3 displays results for ProtoNet (cosine) and MAML on CUB-200 and Describable Textures, with the complete set of experiments available in Appendix F. Out of the 64 total cross-domain scenarios, online UNIFORM does statistically better in 49/64 scenarios, comparable in 12/64 scenarios and worse in only 3/64 scenarios. These results further go to show that sampling matters in episodic training.

heading: 5.6 Better sampling improves few-shot classification

text: The results in the previous subsections suggest that online UNIFORM yields a simple and universally applicable method to improve episode sampling. To validate that state-of-the-art methods can also benefit from better sampling, we take the recently proposed FEAT [64] algorithm and augment it with our IS-based implementation of online UNIFORM. Concretely, we use their open-source implementation5 to train both baseline and online UNIFORM sampling. We use the prescribed hyper-parameters without any modifications. Results for ResNet-12 on Mini-ImageNet and TieredImageNet are reported in Table 4, where online UNIFORM outperforms baseline sampling on 3/4 scenarios and is matched on the remaining one. Thus, better episodic sampling can improve few-shot classification even for the very best methods.

heading: 6 Conclusion

text:
This manuscript presents a careful study of sampling in the context of few-shot learning, with an eye on episodes and their difficulty. Following an empirical study of episode difficulty, we propose an importance sampling-based method to compare different episode sampling schemes. Our experiments suggest that sampling uniformly over episode difficulty performs best across datasets, training algorithms, network architectures and few-shot protocols. Avenues for future work include devising better sampling strategies, analysis beyond few-shot classification (e.g., regression, reinforcement learning), and a theoretical grounding explaining our observations.
5Available at: https://github.com/Sha-Lab/FEAT

title: Uniform Sampling over Episode Difficulty