Deep Reference Priors: What is the best way to pretrain a model?
Abstract
What is the best way to exploit extra data—be it unlabeled data from the same task, or labeled data from a related task—to learn a given task?
This paper formalizes the question using the theory of reference priors.
Reference priors are objective, uninformative Bayesian priors that maximize the mutual information between the task and the weights of the model.
Such priors enable the task to maximally affect the Bayesian posterior, e.g., reference priors depend upon the number of samples available for learning the task and for very small sample sizes, the prior puts more probability mass on low-complexity models in the hypothesis space.
This paper presents the first demonstration of reference priors for medium-scale deep networks and image-based data.
We develop generalizations of reference priors and demonstrate applications to two problems.
First, by using unlabeled data to compute the reference prior, we develop new Bayesian semi-supervised learning methods that remain effective even with very few samples per class.
Second, by using labeled data from the source task to compute the reference prior, we develop a new pretraining method for transfer learning that allows data from the target task to maximally affect the Bayesian posterior.
Empirical validation of these methods is conducted on image classification datasets.
Introduction
Exploiting extra data, e.g., labeled data from a related task, or unlabeled data from the same task, is a powerful way of reducing the number of training data required to learn a given task.
This idea lies at the heart of burgeoning fields like transfer, meta-, semi- and self-supervised learning, and these fields have developed a wide variety of methods to incorporate such extra information.
To give a few examples, methods for transfer learning fine-tune a representation that was pretrained on labeled data from another—ideally related—task.
Methods for semi-/self-supervised learning pretrain the representation using unlabeled data, which may come from the same task or from other related tasks, before using the labeled data.
In this paper, we ask the question: what is the best way to exploit extra data for learning a task?
In other words, if we have some pool of data—be it labeled or unlabeled, from the same task, or from another task—what is the optimal way to pretrain a representation?
Asposed,the answer to the above question depends upon the downstream task that we seek to solve and therefore it seems difficult to obtain a general solution.
We can formulate a more reasonable question by recognizing that a pretrained representation can be thought of as a Bayesian prior (or a sample from it).
Fundamentally, a prior restricts the set of models that can be fitted upon the task.
So we could instead ask: how to best use the extra data to restrict the set of models that we could fit on the desired task.
We formalize this question in this paper using the theory of reference priors in Bayesian statistics.
Our paper has the following contributions.
(1) We formalize the problem of “how to best pretrain a model” using the theory of reference priors, which are objective, uninformative Bayesian priors computed by maximizing the mutual information between the task and the weights.
We show how these priors maximize the KL-divergence between the posterior computed from the task and the prior, on average over the distribution of the unknown future data.
This allows the samples from the task to maximally influence the posterior.
We discuss how reference priors are supported on a discrete set of atoms in the weight space.
We develop a scalable, particle-based algorithm to compute reference priors for deep networks.
To our knowledge, this is the first demonstration of reference priors for deep networks that maintains their characteristic nature, namely that they are supported in a discrete set of atoms.
) We formalize semi-supervised and self-supervised learning as computing a reference prior where the learner is given access to a pool of unlabeled data and seeks to compute a prior using this data.
This formulation sheds light upon the theoretical underpinnings of existing state of the art methods such as FixMatch.
In particular, we show that techniques such as consistency regularization and entropy minimization which are commonly used in practice can be directly understood using the reference prior formulation.
This work therefore provides an information-theoretic formulation of semi-supervised and self-supervised learning.
(3) We develop a novel two-stage reference prior where the learner gets access to data in two stages and seeks to compute a prior that is optimal for data from the second stage.
By solving for the optimal prior in this case, we show how it has the flavor of ignoring certain parts of the weight space depending upon whether data from the first stage was similar to that from the second stage, or not.
This objective is closely related to the predictive Information Bottleneck principle.
(4) We formalize transfer learning using the two-stage reference prior where the goal is to build a prior using data from a source task that maximally leverages data from the target task.
This formulation is useful because it is an information-theoretically optimal way to pretrain using a source task for the goal of transferring to the target task.
We show an empirical study of our formulations on the CIFAR-10 and CIFAR-100 datasets.
We show that our methods to compute reference priors are scalable, provide results that are competitive with state of the art methods for semi-supervised learning, and obtain significantly better accuracy than fine-tuning for transfer learning, even for very small sample sizes.
Background
This section introduces notation, develops the formulation for objective Bayesian priors and gives a few examples that explain how these priors work.
2. Setup
Consider a dataset ˆ P n = { ( x i , y i ) } ni =1 with n samples that consists of inputs x i ∈ R d and labels y i ∈ { 1 , . .
, C } .
Each sample of this dataset is drawn from a joint distribution P ( x, y ) which we define to be the “task”.
It will be useful to use the shorthand x n = ( x 1 , . . .
, x n ) and y n = ( y 1 , . . .
, y n ) to denote all inputs and labels.
Let w ∈ R p be the weights of a probabilistic model which evaluates p w ( y | x ) .
We will use a random variable z with a probabilistic model denoted by p w ( z ) when we do not wish to distinguish between inputs and labels.
Given a prior on weights π ( w ) , Bayes law gives the posterior
[Equation]
The Fisher Information Matrix (FIM) g ∈ R p × p has entries,
[Equation]
and it can be used to understand the differences between two probabilistic models.
For a model with weights w , the Kullback-Leibler (KL) divergence KL ( p w , p w +d w ) =
(cid:82) d P ( x ) d y p w ( y | x ) log( p w ( y | x ) /p
w +d w ( y | x )) can be written as
[Equation]
The FIM is preserved under diffeomorphisms of the weights w (Amari, 2016).
In Bayesian statistics, it is often used to build the Jeffreys prior π J ( w ) ∝ (cid:112) det g ( w ) .
Jeffreys prior is reparameterization invariant, i.e., it assigns the same probability to a set of models irrespective of our choice of parameterization of those models.
It is an example of an “uninformative prior” which means that it allows us to impose some generic structure upon the problem, e.g., that the prior be reparameterization invariant.
Informative priors are different from such uninformative priors. For instance, if we were to choose π to be a Gaussian distribution with mean at a pretrained model then such a prior would express a very definite information and bias the ensuing posterior towards the pretrained model.
2. Reference Priors
The choice of a prior is typically subjective and left to the user.
In order to make this choice more objective,
Bernardo (1979) suggested that uninformative priors should maximize some divergence, say the KullbackLeibler (KL) divergence KL ( p ( w | z ) , π ( w ))
= (cid:82) d w p ( w | z ) log ( p ( w | z ) /π ( w )) , between the prior and the posterior for data z .
The rationale for doing so is to allow the data to dominate the posterior rather than our choice of the prior.
Since we do not know the data a priori while picking the prior, we should maximize the average KL-divergence over the data distribution p ( z ) .
This amounts to maximizing the mutual information
[Equation]
where p ( z ) =
(cid:82) d w π ( w ) p ( z | w ) and H ( w ) =
(cid:82) d w π ( w ) log π ( w ) is the Shannon entropy; the conditional entropy H ( w | z ) is defined analogously.
Mutual information is a natural quantity for measuring the amount missing information about w provided by data z when the initial belief was π .
The prior π ∗ ( w ) is known as a reference prior. It is invariant to a reparameterization of the weight space because mutual information is invariant to reparameterization.
The reference prior does not depend upon the samples ˆ P n but only depends on their distribution P .
The objective to calculate reference prior π ∗ above may not be analytically tractable.
Bernardo therefore also suggested that we compute n -reference priors.
In this paper, we call n the “order” and have deliberately overloaded the notation for the number of samples n ; the reason for doing so will be clear soon.
[Equation]
using n samples and then setting π ∗ : = lim n →∞ π ∗ n under appropriate technical conditions (Berger et al., 1988).
Reference priors are asymptotically equivalent to Jeffreys prior for one-dimensional problems.
They differ in general for multi-dimensional problems but it can be shown that Jeffreys prior is the continuous prior that maximizes the mutual information (Clarke and Barron, 1994).
2. Blahut-Arimoto algorithm
The Blahut-Arimoto algorithm (Arimoto, 1972; Blahut, 1972) is a general method for maximizng functionals like (1) and leads to iterations of the form π t +1 ( w ) ∝ exp ( KL ( p ( z | w ) , p ( z )))
π t ( w ) .
It exploits the identity
[Equation]
where the optimum on the right-hand-side is achieved when q ( w | z ) ≡ p ( w | z ) .
The BA algorithm alternates the maximization over π and q in this identity to result in Maximizing mutual information is a convex problem and therefore the BA algorithm is guaranteed to converge.
It is typically implemented for discrete variables, e.g., Information Bottleneck (Tishby et al., 1999).
Such a discretization of the weight space is prohibitive for high-dimensional models, we therefore implement this algorithm using particles in Remark 4.
Example 1 (Estimating the bias of a coin).
Consider the estimation of the bias of a coin w ∈ [0 , 1] using n trials.
If z denotes the numberof heads (which isa sufficient statistic), we have p ( z | w ) =
w z (1 − w ) n − z n ! /
( z !( n − z )!) .
For n = 1 , since we know that I ( w ; z 1 ) ≤ log 2 with this one bit of information, we can see that π ∗ 1 ( z ) =
( δ ( w ) + δ (1 − w )) / 2 is the reference prior that achieves the upper bound (Berger et al., 1988).
This result is intuitive: if we know that we have only one observation and the outcome is either a heads or a tails, the optimal uninformative prior puts equal probability mass on w = 0 and w = 1 .
We can numerically find π ∗ n for different values of n using the BA algorithm (Fig. 2).
Example 2 (Visualizing the reference priors in data space).
We next construct a way to visualize the reference prior for deep learning classification problem.
To compute the reference prior, we implement the algorithm in Remark 4 using particles(atoms).
Note that w are weights of a deep network, it is not reasonable to visualize π ∗ n ( w ) in such a high dimension weight space.
We use the following strategy to visualize the prior. Given a classification problem with C classes and n samples in the training set, we think of the predictions of the network as a vector in n × C dimensions given by [Equation] This vector f ( w ) is parameterized by the weights w of the network.
The rationale behind this choice is to observe that for two weights w, w ′ , the Euclidean distance between the prediction vectors is the Hellinger
[Figure]
divergence between the prediction probability distributions,
[Equation]
where
[Equation]
is the Hellinger distance.
In other words, the prediction vector f ( w ) maps the weights w into a ( n ×
C ) − dimensional Euclidean space, which makes it possible to visualize the reference prior, as shown in Fig.
At the beginning stage of the optimization, all atoms w k make similar predictions (blue points in Fig.concentrate near each other).
This is not surprising because the atoms are initialized randomly and w k has a uniform distribution at its output.
As the prior is optimized to maximize the mutual information, its atoms increasingly make more diverse predictions (orange) and towards the end stage (green), the prediction of the atoms forms a low-dimension manifold.
We hypothesize that these green particles form a “boundary” of the possible prediction vectors made by the model and consist of simple, low-dimensional models (Mattingly et al., 2018).
Methods
This section first discusses a key property of these priors that enables us to calculate them numerically, namely the fact that they are supported on a discrete set in the weight space for finite order (§3.1).
It then constructs priors that are designed to make maximal use of pre-training data for semi-supervised (§3.2) and transfer learning (§3.4).
3. Existence and discreteness of reference priors
Existence and boundedness. As (Berger et al., 1988) say, reference prior does not exist if I π ( w ; z n ) is infinite.
To address the existence of the reference prior, we need to assume that (i) π is supported on a compact set Ω ⊂ R p , and (ii) if p π ( z n ) = (cid:82) Ω d w π ( w ) p ( z n | w ) is the marginal, KL ( p w , p π ) is a continuous function of w for any π . Then the n -reference prior π ∗ n exists and I π n ( w ; z n ) is finite; see Lemma 2.14 in (Zhang, 1994).
Discreteness.
Suppose that z n is a discrete random variable with number of atoms denoted by C . Suppose p−   p−2. p Beginning
Middle End that the n − reference prior π ∗ n exists and is unique(up to measure zero), Let Ω n = { w ∈ Ω : π ∗ n ( w ) > 0 } be the support of π ∗ n . Suppose { p ( z n | w ) : w ∈ Ω n } is compact. Then π ∗ n is discrete with number of atoms no greater than C (see Lemma 2.18 in (Zhang, 1994)).
Rigorous development of reference priors has been done by these earlier works.
Our paper focuses on some innovative applications to problems in machine learning.
Remark
(The reference prior depends upon the number of samples in the task and finds a small set of diverse models).
The likelihood of each atom in the optimal prior p ( z n | w ) is maximally different from that of other atoms, and thereby the marginal likelihood p ( z n ) .
It is intuitive why the prior should have finite atoms, this is because it is not possible to distinguish between a large set of models using few samples n .
For models like deep networks, this is extremely promising in the low-sample regime because it automatically selects a small set of models from a large model space to fit to the few data.
This is analogous to covering numbers in learning theory (Bousquet et al., 2003) where we endow the hypothesis space with a metric that measures disagreement between two hypotheses over n samples.
Smaller the n , smaller the covering number, and smaller the set of models considered by the learner.
Remark 4 (Using multiple particles in place of the Blahut-Arimoto algorithm).
Since the optimal prior is discrete under certain assumptions, we can maximize the mutual information directly by identifying the best set of atoms.
We set π ∗ n = (cid:80) Ki
=1 λ i δ ( w − w i ) where { w 1 , . . .
, w K } are the K atoms of our prior π ∗ n
and λ i are coefficients with (cid:80) i λ i = 1 .
We call these atoms “particles”.
Abbott and Machta (2019) also suggest a scaling law for K in terms of n .
For example, they suggest that K ∼ n 4 / 3 for a problem with two biased coins.
Each particle for us will be a deep network, so we will instead treat K as a hyper-parameter.
Remark 5 (Variational approximation of reference priors).
Nalisnick and Smyth (2017) maximize a lower bound on I π ( w ; z ) and replace the term p ( z ) =
(cid:82) d w π ( w ) p ( z | w ) in (1) by the so-called VR-max estimator max w log p ( z | w ) where the maximum is evaluated across a set of samples from π ( w ) (Li and Turner, 2016).
They use a continuous variational family parameterized by neural networks.
Our arguments show that the prior is supported on a discrete set and we need not use a continuous-valued distribution to model it.
We therefore maintain π ( w ) as a set of particles.
3. Reference priors for semi-supervised learning
Consider the situation where we are given inputs x n and labels y n , and an additional pool of unlabeled inputs x u from the same task which we can use to compute π ∗ ( w ) .
The key idea behind using reference priors for semi-supervised learning is that since the KL-divergence between the posterior on ( x n , y n ) and the prior π ( w ) is maximized on average over the inputs, we can use the unlabeled inputs x u to perform this averaging.
More precisely, we set
[Equation]
where p π ( y n | x n )
= (cid:82) d w π ( w ) (cid:81)
ni =1 p ( y i | x i , w ) and likewise for p π ( y u | x u ) .
We have switched the expectation in the final step to over x u instead of labeled inputs x n because they have the same distribution.
For the same reason, we also switched the expectation on all possible labels y n of the labeled data, to y u , the unknown labels of unlabeled data.
Assume for now that we know the number of classes C and can compute expectations over p ( y n | x n , w ) .
If our prior has K particles, then the second term is the average of the per-particle entropy of the predictions.
The objective encourages each particle w i to predict confidently, i.e., to have a small entropy in its output distribution p w i ( y | x ) .
The first term is the entropy of the average predictions: p π ( y n | x n ) , and it is large if particles predict different outputs y n for the same inputs x n , i.e., they disagree with each other.
Effectively, the objective encourages particles to be very dissimilar but confident (not necessarily correct) models.
We use the empirical average over unlabeled inputs to maximize (3).
Predictions on new samples x are made using the Bayesian posterior predictive distribution
[Equation]
Remark 6 (The number of classes need not be known a priori).
Consider a model that has a feature generator with weights w in conjunction with a fixed (or a non-parametric) classifier that makes predictions over C ′ classes.
Since (3) involves an expectation over the predictions p ( y u | x u , w ) , C ′ need not be the same as the number of classes C for the task.
We can compute the prior over weights w of the feature generator and use it to learn a model for labeled data.
This property mirrors the prevalent practice in self-supervised learning, e.g., (Chen et al., 2020b), where the feature generator is trained on pretext tasks that are different from the eventual goal of classification.
Remark 7 (Reference prior spreads out on the statistical manifold).
Given the weight space, the prior spreads its probability mass such that while each particle that is likely under the prior can confidently (not necessarily correctly) make predictions on inputs (first term in (3)), predictions of these particles are maximally entropic when combined together (second term).
If we consider the statistical manifold from information geometry (Amari, 2016) which consists of all models p w ( y | x ) , then the reference prior has a finite number of atoms and is maximally spread out on this manifold to ensure that the particles are distinct.
Asymptotically, as the number of atoms grows to infinity, it approaches the flat prior on the manifold, which corresponds to
Jeffreys prior on the weight space (Clarke and Barron, 1994).
3. Reference priors for a two-stage experiment
In this section, we do not distinguish between inputs and labels and first develop the idea in terms of a generic random variable z n .
Consider a situation when we obtain data in two stages, first z m , and then z n .
How should we select a prior, and thereby the posterior of the first stage, such that the posterior of the second stage makes maximal use of the new n samples?
We can extend the idea in §2.2 in a natural way to address this question.
We can maximize the divergence between the posterior of the second stage and the posterior after the first stage, on average, over the new samples z n .
Since we have access to samples z m , we need not average over them, we can compute the posterior p ( w | z m ) from these samples given the prior π ( w ) .
First, notice that by chain rule p ( w, z n | z m ) =
p ( w | z m + n ) p ( z n | z m ) = p ( z n | w ) p ( w | z m ) .
We can now write
[Equation]
[Equation]
where p ( w | z m ) ∝ p ( z m | w ) π ( w ) and p ( z n | z m )
= (cid:82) d w p ( z n | w ) p ( w | z m ) .
If (2) has a unique solution, we should have the optimal p ( w | z m ) ≡ π ∗ n ( w ) .
This leads to
[Equation]
Notice that the reference prior puts less probability mass on regions which have high likelihood on old data z m .
This is consistent with intuition because the prior is designed to be such that the posterior is maximally informed by the new samples z n .
Given knowledge of old data, the prior downweighs regions in the weight space that could bias the posterior of the new data.
We will exploit this peculiar property of reference priors in this section.
We also have π ∗ n | m = π ∗ n for m = 0 which is consistent with (2).
As m → ∞ , this prior ignores the part of the weight space that was ideal for z m .
Appendix C shows this two-stage prior for the biased coin experiment.
Remark 8 (Averaging over z m in the two-stage experiment).
If we do not know the outcomes z m yet, the prior should be calculated by averaging over both z m , z n
[Equation]
The reference prior encourages multiple explanations of initial data z m , i.e., high H ( w | z m ) , so as to let the future samples x n select the best one among these explanations, i.e., reduce the entropy H ( w | z m + n ) .
It is interesting to note that neither is this
two-stage prior equivalent to maximizing I π ( w ; z m + n ) , nor is it simply the optimal prior corresponding to objectives
I π ( w ; z m ) or I π ( w ; z n ) .
Both (6) and (7) suggest that such priors may be useful when we have some data a priori , e.g., either unlabeled samples from the same task, or labeled samples from some other task.
Remark 9 (Connection to the Predictive Information Bottleneck).
The objective in (7) resembles that of Bialek et al.
(2001), or its variational version in Alemi (2020), which seek to learn a representation, say w , which maximally forgets past data while remaining predictive of future data
[Equation]
The parameter β in (8) gives this objective control over how much information from the past is retained in the representation w .
This objective is therefore called the predictive information bottleneck (IB).
A softer version of reference prior for the two-stage experiment.
We take inspiration from the predictive IB and construct a variant of the prior in (6)
[Equation]
for β ∈ (0 , 1) .
[Equation]
In this prior, we should use a value of β = 0 when we expect that data from the first stage z m is very similar to data z n from the second stage; this allows the posterior to benefit from the past samples .
If we expect that the data are different, then we should set β = 1 which ignores regions in the weight space that predict well for z m while selecting the prior for z n .
This is similar to the predictive IB where a small β encourages remembering the past data and β = 1 encourages forgetting.
The difference between our reference prior approach and the predictive IB is that the prior in the latter need not be a reference prior.
3. Reference priors for transfer learning
Now consider the two-stage experiment where at each stage we obtain data from a different task.
The first stage consists of m samples from a “source” task P s and the second stage consists of n samples from the “target” task
P t .
Our goal is to calculate a prior π ( w ) that best utilizes the target task data.
Bayesian inference for this problem involves first computing the posterior p ( w | x ms , y ms ) ∝ p ( y ms | w, x ms ) π ( w ) from the source task and then using it as a prior to compute the posterior for the target task p ( w | x nt , y nt , x ms , y ms ) .
Just like §2.2, we would like to maximize the average KL-divergence between the two posteriors KL [ p ( w | x nt ,
y nt , x ms , y ms ) , p ( w | x ms , y ms )] which measures the amount of the information learnt from samples from the target task.
However, note that this KL-divergence is a random variable: its randomness comes from samples x ms and x nt .
Therefore, instead of maximizing the divergence given a specific x ms and x nt , we maximize the average divergence.
We consider the following cases.
Case 1: Access to a pool of unlabeled data from the source x ms and the target task x nt .
In this situation, we should average the KL-divergence over the predictions of the model y m s on the source data x m s ; this is in addition to the unknown labels y nt of the target task.
This involves maximizing
[Equation]
over the prior π .
Here p π ( y ms | x ms ) =
E w ∼ π p ( y ms | x ms , w ) and p π ( y nt | x nt ) =
E w ∼ π p ( y nt | x nt , w ) , respectively.
Case 2: x ms , y ms are fixed and known, and we have a pool of unlabeled data for x nt .
Since we already know the data and labels for the source task, we will only average over x nt and y nt and maximize the objective
[Equation]
over π ; here p π ( y nt | x nt ) = (cid:82) d w π ( w ) p ( y nt | x nt , w ) .
Note that we need not have a separate pool of unlabeled data from the target, we can simply use the samples x nt and average over the labels; this also applies to (3).
Remark 10 (Connecting (10) and (11) to practice).
Both objectives can be written down as
[Equation]
with the distinction that while in Case 1, we average over all quantities, namely p ( x ms ) , p ( y ms ) , p ( x nt ) , p ( y nt ) while in Case 2, we fix the distributions of p ( x ms ) and p ( y ms ) to the empirical distribution of the source task data.
Case 2 is commonly considered transfer learning.
Case 1, where one has access to only unlabeled data from a source task that is different from the target task is not typically studied in practice.
If the source and target tasks are the same, then the objective in (12) is the same as that of (3).
If they are not, our theory indicates that the prior in (12) should minimize the mutual information with respect to the source task; this is similar to the argument in §3.3.
Like (9), we can again introduce a coefficient β on the second term in (12) to control handle the relatedness between source and target tasks.
Remark 11 (Reference priors for semi-supervised and transfer learning are discrete). ?? extends directly to the semi-supervised learning case in (3) and the two-stage experiment in (6) and (9).
Therefore we can show that the prior should be discrete if there are a finite amount of labeled samples in semi-supervised learning, or if the target task in transfer learning provides finitely many samples.
3. Reference priors for self-supervised learning
Reference priors shed light on understanding self-supervised learning techniques.
Contrastive approaches of self-supervised learning (SelfSL) learn representations by minimizing the distance between two augmented views of the same data point (positive pairs) and maximizing views from different data points (negative pairs).
In this section, w are weights of a stochastic feature extractor.
Once an input data x was feed into the model, a stochastic representation h ∈ Γ ⊂ R l is generated, where Γ is a compact subset in l − dimension real space.
Let G be a set of data augmentations, we write the joint distribution, for g ∈ G
[Equation]
then h | x, w distributes as the ensemble averaging law
[Equation]
where G is a set of data augmentations.
We extend the spirit of reference prior in a natural way and write down the mutual information bwtween h n and x n given w .
[Equation]
where p ave w ( h n | x n )
= (cid:81) ni
[Equation]
1 p ave w ( h
[Equation]
| x i
[Equation]
and p ave w ( h n ) =
E x
Remark 12.
The first term in (15) matches with the negative pair penalty in contrastive learning, while the second tern in (15) coincides with the positive pairs.
Therefore contrastive learning is equivalent to maximizing mutual information between representations and data.
That explains why contrastive approaches works since they are selecting models that do their best to learn information from data:
[Equation]
Remark 13 (connecting with free energy principle for representation learning).
(Gao and Chaudhari, 2020b) pick reconstruction of the original data—as a way to measure the discarded information in the representation when it is fitted on a specific task, this idea leads to the study of the following Lagrangian which is similar to the Information Bottlenck of (Tishby et al., 2000)
[Equation]
where the rate R is an variational upper bound on the mutual information
I w ( x | h ) , distortion D measures the quality of reconstruction of the decoder, and C measures the classification loss.
The canonical reconstruction task forces the representation learning redundant information about the data, (Gao and Chaudhari, 2020b) empirically show that these extra information are potentially helpful for transfer learning.
Instead of attaching a canonical task, we directly setting the constrains min w C such that I w ( x n | h n ) ≥ MI, which forces the representation learning redundant information.
3. Practical tricks for implementing reference priors
The reference prior objective is conceptually simple but it is difficult to implement it directly using deep networks and modern datasets.
We next discuss some practical tricks to compute reference priors.
Order of the prior n versus the number of samples. Theory sets the order of the prior n to be the same as the number of samples (Bernardo, 1979).
We make a distinction between the two and restrict our experiments to order n = 2 , 3 .
Mathematically, this amounts to computing averages in (2) or (3) over only sets of n samples at a time rather than all of them.
This significantly reduces the class of models considered in the reference prior by pretending that there is an extremely small number of samples available for training the task; this is useful for large models like deep networks.
Note that we are not restricting to small order n for computational reasons, i.e., computing the expectation over all classes y n in (3) can be done in a single forward pass.
Using cross-entropy loss to bias particles towards good parts of the weight space. Consider Case 2 in §3.4.
The posterior (4) suggests that we should first compute the prior, and then weight each sample by the likelihood of the labeled data.
In practice, we combine these two steps into a single objective
[Equation]
where γ is a hyper parameter, x n , y n are labeled samples.
(17) makes sense since it allows us to directly obtain particles that both have high probability under the prior and that have a high likelihood.
Using such particles for inference is different from using the correct Bayesian posterior, in particular due to the coefficient γ , which we set to be γ = 1 / 2 .
We initialize the prior π ( w )
= (cid:80) Ki
=1 λ i δ ( w − w i ) at randomly chosen weights { w i } and perform gradient ascent steps on the above objective to directly compute the likelihood weighted particles.
This strategy further restricts the search space for the particles in π ( w ) .
A warm posterior, i.e., γ < 1 , as opposed to cold posteriors (Wenzel et al., 2020) used in supervised learning, is a reasonable thing to do here because our prior π ( w ) is highly entropic.
Data augmentations.
Data augmentation is a key component of state-of-the-art approaches for semi-supervised learning.
Methods like Fix-match employ heavy data augmentations.
We also use data augmentations for experiments on CIFAR-10 and CIFAR-100.
For implementing (17), we define the averaging prediction
[Equation] as a linear combination of the predictions on the same input augmented by different augmentations G weak and G strong , where weak augmentations are flips/shifts, and strong ones are from AutoAugment (Sohn et al., 2020).
Notice that by Jensen’s inequality, [Equation]
The reference prior for semi-supervised learning with weak and strong data augmentations is computed using
[Equation]
where x u is sampled from unlabelled data, ( x n , y n ) are labeled samples.
If G weak and G strong are identical augmentations, then p ave w ( y | x ) =
p ( y | x, w ) and the first term converges to − γH ( y u | x u , w ) .
We can therefore recover (17) from (18).
3. Relationship of the reference prior objective to existing methods for semisupervised learning
The reference prior objective in (3) is closely connected to existing methods in semi-supervised learning.
It is a difference of two entropries: I π ( [Equation] [Equation] [Equation]
=
E [Equation] [Equation] |
[Equation]
] − E
[Equation] [Equation] ∼ π [ S ( y n | x n , w )] . A number of methods, e.g., FixMatch (Sohn et al., 2020) or MixMatch (Berthelot et al., 2019), force the model to identify whether the underlying input image is the same across strong augmentations.
Let G be a set of such augmentations.
These methods use a quadratic penalty for p ( y n | x n ) with respect to deviations from p G ( y n | x n , w ) =
| G | − 1 (cid:80) g ∈ G p ( y n | g ( x n ) , w ) where g ( x n ) denotes the result of applying the augmentation g to inputs x n .
For our objective, minimizing the second term, E x n ,w ∼ π [ H ( y n | x n , w )] with this augmentationaveraged distribution p G ( y n | x n , w ) achieves the same purpose.
Disagreement-based methods (Zhou and Li, 2010) employ multiple models and use the confident models to soft-annotate unlabeled samples for others.
Our reference prior is discrete and therefore also consists of multiple models.
Disagreements among them are measured by H ( y n | x n ) .
If p ( y n | x n ) is uniform, which is encouraged by the reference prior objective, we obtain strong disagreements.
Thefirstterm
[Equation]
(cid:82)
[Equation]
[Equation]
12 log
[Equation]
G weak (
[Equation] [Equation] 12 log p ( [Equation]
u | G strong (
[Equation]
u )
[Equation]
w ) (cid:3) in(18) closely relates to FixMatch which uses a cross-entropy loss that encourages consistent predictions across weak and strong augmentations.
As we saw in (18), the key difference between FixMatch and our objective is the second term E x u [ H ( y n | x n )] that encourages disagreements between particles.
Empirical Study 4. Setup
We evaluate on the CIFAR-10 and CIFAR-100 (Krizhevsky, 2009) datasets.
For semi-supervised learning, we show results using 50–1000 labeled samples, i.e., 5–100 samples/class and use the rest of the samples in the training set as unlabeled samples.
For transfer learning, we construct 20 5-way classification tasks using the super classes of the CIFAR-100 dataset.
For all experiments on CIFAR-10/CIFAR-100, we use the WRN 28architecture (Zagoruyko and Komodakis, 2016) which is identical to that used in Berthelot et al.
We use the Adam (Kingma and Ba, 2015) to train with a constant learning rate of 0.002.
For all our experiments, the reference prior is of order n = 2 and has K = 4 particles.
We run our methods for 512 epochs, with the coefficient γ scaled linearly from 0 to 0.25 over the entire training period and α fixed at 0.1.
During evaluation, we use the Exponential moving average (EMA) of the model parameters which is commonly used in semi-supervised learning (Tarvainen and Valpola, 2017).
Appendix A provides more details.
4. Baseline methods for semi-supervised and transfer learning
We compare Deep reference priors to existing methods in semi-supervised learning such as FixMatch (Sohn et al., 2020), MixMatch (Berthelot et al., 2019), Mean Teacher (Tarvainen and Valpola, 2017), Π -Model (Sajjadi et al., 2016b), Pseudo-Label (Lee et al., 2013), Virtual Adversarial Training (Miyato et al., 2018), and
Mixup between labeled and unlabeled data (Berthelot et al., 2019).
For transfer learning, we use fine-tuning to adapt a model trained on source to the target data.
We additionally compare to using just the labeled target data (standard supervised learning), using labeled and unlabeled target data (self-supervised learning).
We consider 5 different tasks from CIFAR-100, with 1000 labeled source samples and 100 labeled target samples for all experiments.
4. Semi-supervised learning
We evaluate Deep reference priors in the semi-supervised learning setting using the formulation described in §3.6.
Table 1 compares Deep reference priors with competing semi-supervised learning methods on the
CIFAR-10 dataset.
Our methods outperform 5 other baseline methods (by 5-10%), except MixMatch and
FixMatch.
This indicates that reference priors can effectively leverage unlabeled data to pretrain representations and are competitive, if not better, than the state of the art.
Deep reference priors are applicable to semisupervised learning without being explicitly formulated to enforce properties like confident predictions on unlabeled samples or label consistency with respect to augmentations.
Fig.
4 compares Deep reference priors to MixMatch in settings where the learner is given access to different number of labeled data.
Both methods have similar trends as the number of labeled samples is changed.
The implementation of MixMatch by the original authors implements a number of techniques such as a moving averaging of the weights, an “interleaved” version of Mixup across labeled and unlabeled data, intricate schedules for various coefficients in their method.
Our theoretical argument in §3.7 shows that the objective used in reference prior is very close to that of MixMatch, so we expect to match their accuracy on this problem.
We have investigated the gap in accuracy at depth and we currently believe that the gap could be mitigated by using a larger number of particles in the reference prior.
The utility of Deep reference priors is further validated in Table 2 (left) which evaluates on 5 semi-supervised learning tasks from CIFAR-100.
Deep reference priors make use of unlabeled samples in addition to the 100 labeled samples for each task using, which results in accuracy gains as large as 18.6% when compared to using just labeled samples (vanilla supervised learning) [Table] [Figure]
4. Transfer learning
We next evaluate Deep reference priors in the transfer learning setting.
The posterior (9) looks to find a model that is influenced by the target data by varying degrees controlled by β , with β = 1 corresponding to the scenario where the posterior is maxmimally influenced by target data after being pre-trained on the source data.
We instantiate (9), by combining prior selection, pre-training on the source task and likelihood on the target task, into one objective,
[Equation]
19) where γ and β are hyper-parameters.
[Figure]
V e h
i c
l e s
V e h
i c
l e s
F i s h
P e o p l e
A q .
M a mm a l s
T a r g e t t a s k                    (1.5) (1.7) (1.6) (1.9) (2.8) (2.4) (1.5) (1.3) (1.8) (2.8) (0.8) (0.5) (2.0) (2.8) (2.5) (0.9) (2.4) (3.1) (2.2) (3.4) (a)
Vehicles
[Figure]
Fish
People Aq.
Mammals
[Figure]
V e h
i c
l e s
V e h
i c
l e s
F i s h
P e o p l e
A q .
M a mm a l s
T a r g e t t a s k                    (5.3) (2.4) (2.2) (5.2) (3.6) (4.0) (1.8) (4.4) (2.6) (6.0) (1.1) (2.6) (3.7) (3.9) (2.7) (1.8) (
[Figure]
[Figure]
[Figure]
[Table]
4. Semi-Supervised Learning Ablation Experiments
The following ablation experiments are conducted for semi-supervised learning on CIFAR-10 withlabeled samples.
Order of prior. The order of the prior is an independent hyper-parameter as discussed in §3.6.
Changing the order leads to a marginal (about 1%) improvement in accuracy.
This justifies our choice of using n = 2 as the order in all experiments.
[Table]
Number of particles in the prior. We vary the number of particles in the prior and observe an improvement in accuracy as we increase the number of particles.
The downside of this is however that a large number of particles implies increasing training time for the prior, which grows linearly with the number of particles.
[Table]
Particles Weights in the Prior. As discussed in Remark 4, the prior is represented by π ∗ n =
K (cid:88) i =
λ i δ ( w − w i ) where { w 1 , . . .
, w K } are the K atoms of the prior and λ i are coefficients with (cid:80) i λ i = 1 .
In our implementation, we set λ i to be trainable parameters.
We observe that λ i assumes a uniform distribution over the particles in our experiments (Fig. 6).
We expect this observation to change for a very large number of particles.
For a small number of particles, the individual particles are extremely confident in their predictions.
Hence, the term
H ( y u | x u ) is maximized when the predictions averaged across all particles is close to uniform distribution which is achieved when λ i is uniform over all particles. [Figure] [Figure]
4. Analysis and additional ablation studies
Diversity of the models in the prior. Fig.
7 shows how the entropy H
( y u | x u ) enforces diversity between the different particles.
We see that the accuracy of particles diverge during training.
While one particle clearly dominates the posterior in terms of predicting correctly on most of the samples, the Bayesian posterior predictive distribution has a slightly higher accuracy than the best particle.
Design Choices.
We next seek to dissect different components in our model, like EMA, Mixup and the inclusion of the cross-entropy term (17) (referred to as cross-entropy warmup).
We evaluate on a 2-way classification problem from CIFAR-10 (Cat vs. Dog) on 200 labeled samples run for 70 epochs.
This is a challenging binary classification problem (relatively speaking, it is the hardest pair of classes in CIFAR-10).
Fig.
9a indicates the benefit of including Mixup and the cross-entropy term.
This benefit is more pronounced in the 10-way classification problem.
Fig.
4 shows why using EMA in the evaluation step in beneficial.
Choice of α .
We again evaluate on the Cats vs. Dogs 2-way classification task to understand the impact of the hyper-parameter choice for α .
Fig
[Figure]
[Figure]
Ref
[Figure]
Related Work and Discussion
Reference priors in Bayesian statistics. We build upon the theory of reference priors which was developed in the objective Bayesian statistics literature Bernardo (1979); Berger et al.
The main idea used in our work is that non-asymptotic reference priors allow us to exploit the finite samples from the task in a fundamentally different way than classical Bayesian inference.
If the number of samples from the task available to the learner is finite, then the prior should also select only a finite number of models.
Reference priors are not common in the machine learning literature.
A notable exception is Nalisnick and Smyth (2017) who optimize a variational lower bound and demonstrate results on small-scale problems.
The main technical distinction of our work is that we explicitly use the discrete prior instead of a variational approximation.
To our knowledge, our work is the first instantiation of reference priors for medium-scale deep networks and image-based tasks.
Information theory. Discreteness is seen in many problems with an information-theoretic formulation, e.g., capacity of a Gaussian channel under an amplitude constraint (Smith, 1971), representations of neurons in the brain Laughlin (1981); Tkačik et al.
(2008), and biological systems (Mayer et al., 2015).
(Mattingly et al., 2018; Abbott and Machta, 2019) have developed a rich theory of scaling laws that give the number of atoms in the prior as a function of the number of data.
They use examples to study how reference priors with finite data select “simple models” which lie on the edges of the parameter space (see Fig. 2).
Although we cannot yet characterize a similar simplicity bias of reference priors for deep networks, we believe that the methods developed in our paper are effective because of this phenomenon; our choice of using a small order n for the prior is directly motivated by their examples.
On the other hand, information based methods are widely applied in miscellaneous aspects of deep learning.
One of the common motivation behind unsupervised contrastive representation learning ?? Oord et al. ; Henaff (2020); Tian et al. ; Bachman et al.
; Hjelm et al.
; He et al.
; Chen et al.
(2020a)
; Wang and Isola (2020) is the InfoMax principle (Linsker, 1988), which instantiated by maximizing the mutual information (MI) between multi-views of the data.
While in regularized information maximization for clustering, (Gomes et al., 2010; Bridle et al., 1992; Hu et al., 2017) directly maximize the mutual information between inputs and their representations.
Beyond all these works on various information criterion, our work heavily explores the statistical dependency between the model weights and the data, which shed lights on understanding what the optimal way to pre-train a representation is.
Semi-Supervised Learning.
Semi-supervised learning is a mature field with a wide diversity of approaches (Chapelle et al., 2009; Xiaojin, 2008; Zhu and Goldberg, 2009; Ouali et al., 2020).
Ideas such as consistency regularization, which enforce consistent predictions across different augmentations of the inputs, form an important component of state of the art methods.
Prior work (Bachman et al., 2014; Laine and Aila, 2016; Sajjadi et al., 2016b; Tarvainen and Valpola, 2017; Berthelot et al., 2019) explicitly enforces such losses while we show in §3.7 how the reference prior can automatically enforce consistency regularization.
Minimizing the entropy of predictions on unlabeled data, either explicitly (Grandvalet et al., 2005; Miyato et al., 2018) or using pseudo-labeling methods (Lee et al., 2013; Sajjadi et al., 2016a), is another popular technique.
We show in §3.7 how entropy minimization is automatically achieved by the reference prior objective.
Transfer learning. is a key component of a large number of applications today, e.g, (Devlin et al., 2019;
Kolesnikov et al., 2020; Dhillon et al., 2020).
A central question that remains unanswered today is how one should pretrain a model on the source task given that the eventual purpose is to transfer to a target task, although there have been some partial attempts at addressing it via the Information Bottleneck, e.g.,Gao and Chaudhari (2020a).
This question becomes particularly challenging when transferring across domains, or for small sample sizes (Davatzikos, 2019).
Reference priors are uniquely suited to tackle this question: the two-stage experiment that we have defined is the “optimal” way to leverage data from the source task for the target task, and learning such a prior is particularly useful in low-sample regimes (see §4.4).
A Details of the experimental setup
Architecture. For experiments on CIFAR-10 and CIFAR-100 (§4), we consider a modified version of the Wide-Resnet 28-2 architecture (Zagoruyko and Komodakis, 2016), which is identical to the one used in Berthelot et al.
This architecture differs from the standard Wide-Resnet architecture in a few important aspects.
The modified architecture has Leaky-ReLU with slope 0.1 (as opposed to ReLU), no activations or batch normalization before any layer with a residual connection, and a momentum of 0.001 for batch-normalization running mean and standard-deviation (as opposed to 0.1, in other words these statistics are made to change very slowly).
We observed that the change to batch-normalization momentum has a very large effect on the accuracy of semi-supervised learning.
For experiments on MNIST (Appendix B.1), we use a fully-connected network with 1 hidden layer of size 32.
We use the hardtanh activation in place of ReLU for this experiment; this is because maximizing the mutual information has the effect of increasing the magnitude of the activations for ReLU networks.
One may use weight decay to control the scale of the weights and thereby that of the activations but in an effort to implement the reference prior exactly, we did not use weight decay in this model.
Note that the nonlinearities for the CIFAR models are ReLUs.
Datasets. For semi-supervised learning, we consider the CIFAR-10 dataset with the number of labeled samples varying from 50–1000 (i.e., 5–100 labeled samples per class).
We also consider a 2-class classification problem in some ablation experiments (§4.6) with the two classes being Cat/Dog, which is among the harder 2-class classification problems in CIFAR-10.
Semi-supervised learning experiments use all samples that are not a part in the labeled set, as unlabeled samples.
For transfer learning, we construct two tasks from MNIST (task one is a 5-way classification task for digits 0–4, and task two is another 5-way classification task for digits 5–9).
For this experiment, we use labeled source data but do not use any labeled target data.
This makes our approach using a reference prior similar to a purely unsupervised method.
The CIFAR-100 dataset is also utilized in the transfer learning setup (§4.4).
We consider five 5-way classification tasks from CIFAR-100 constructed using the super-classes.
The five tasks considered are
Vehicles-1, Vehicles-2, Fish, People and Aquatic Mammals.
The selection of these tasks were motivated from the fact that some pairs of tasks are known to positively impact each other (Vehicles-1, Vehicles-2), while other pairs are known to be detrimental to each other (Vehicles-2, People); see the experiments in Ramesh and
Chaudhari (2021).
Optimization. Adam with a constant learning rate of 0.002 was used in our experiments on CIFAR-10 and
CIFAR-100.
Mixed-precision (32-bit weights, 16-bit gradients) was used to expedite training.
Training was performed for 512 epochs unless specified otherwise.
As is common practice in semi-supervised learning literature Berthelot et al.
(2019), the coefficient of the objective of the unsupervised data (mutual information term in our case) was increased linearly during the course of training.
Doing so uses the cross-entropy loss on the labeled data to effectively reduce the space of weight configurations that could be the atoms of the reference prior; this is beneficial while implementing the reference prior over a high-dimensional weight space.
We also conducted experiments with SGD for computing the reference prior on MNIST.
SGD was used with a constant learning rate of 0.001 with Nesterov’s acceleration, momentum coefficient of 0.9 and weight decay of 10 − 5 .
Definition of a single Epoch. Note that since we iterate over the unlabeled and labeled data (each with different number of samples), the notion of what is an epoch needs to be defined differently.
In our work, one epoch refers to 1024 weight updates, where each weight update is calculated using a batch-size of 64 for the labeled data of batch size 64, and a batch-size of 60 for the unlabeled data.
Exponential Moving Average (EMA).
In all CIFAR-10 and CIFAR-100 experiments, we also implement the
Exponential Moving Average (EMA) (Tarvainen and Valpola, 2017).
In each step, the EMA model is updated such that the new weights are the weighted average of the old EMA model weights, and the latest trained model weights.
The weights for averaging used in our work (and most other methods) are 0.999 and 0. respectively.
Note that EMA only affects the model that used for testing, it does not affect how weight updates are calculated during training.
Data Augmentations. We use random-horizontal flips and random-pad-crop (padding of 4 pixels on each side) as augmentations for the CIFAR-10 and CIFAR-100 datasets.
Additionally, Mixup (Zhang et al., 2018) across both labeled and unlabeled samples was used.
Pseudo-labels of the unlabeled samples where computed using the the average softmax predictions across multiple augmentations of the same image, as explained in §3.6.
Note that this step is necessary in order to apply Mixup between labeled and unlabeled samples.
No data augmentations were used for MNIST.
B Additional Experiments B. Unsupervised transfer learning on MNIST
For all our MNIST experiments, the reference prior is of order n = 2 and has K = 50 particles.
We run our methods for 1024 epochs.
We first compare Deep reference priors with fine-tuning for transfer learning.
The posterior (9) seeks to find a model that is influenced by the target data to varying degrees controlled by β , with β = 1 corresponding to the scenario where the posterior is maximally influenced by target data after being pre-trained on the source data.
We instantiate (9), by combining prior selection, pre-training on the source task into one objective,
[Equation]
where γ and β are hyper-parameters.
Solving (S-20) requires no knowledge from target data labels, therefore the setting here is pure unsupervised clustering for target task dataset.
We compare the deep reference unsupervised transfer to fine-tuning method which adapts a model trained on labelled source to the labelled target data.
All samples from the source task (about 30,000 images across 5 classes) were used for both the reference prior and fine-tuning.
C Two-stage experiment for coin tossing
In §3.3, we consider a situation when we obtain data in two stages, first z m , and then z n .
We propose a prior π ∗ ( (7)) such that the posterior of the second stage makes maximal use of the new n samples.
In this section, we visualize π ∗ in the parameter space using a two-stage coin tossing experiment.
[Table] [Table]
Consider the estimation of the bias of a coin w ∈ [0 , 1] using two-stage m + n trials.
If z denotes the number of heads in total [Figure] we have p ( z | w ) =
w z (1 − w ) m + n − z ( m + n )! /
( z !
( m + n − z )!) .
There are m trials in first stage and n trails in second stage.
We numerically find π ∗ for different values of m and n using the BA algorithm (Fig.
S-10 and Fig. S-11)
[Figure]
w 0.
π * ( w ) 0.
w 0.
π * ( w )