Deep Reference Priors:
What is the best way to pretrain a model? ¹¹1Proceedings of the 39${}^{\text{th}}$ International Conference on Machine Learning, Baltimore, Maryland, USA. Copyright 2022 by the authors.

Consider a dataset $\hat{P}_{n}=\left\{(x_{i},y_{i})\right\}_{i=1}^{n}$ with $n$ samples that consists of inputs $x_{i}\in\mathbb{R}^{d}$ and labels $y_{i}\in\left\{1,\ldots,C\right\}$. Each sample of this dataset is drawn from a joint distribution $P(x,y)$ which we define to be the “task”. We will use the shorthand $x^{n}=(x_{1},\ldots,x_{n})$ and $y^{n}=(y_{1},\ldots,y_{n})$ to denote all inputs and labels. Let $w\in\mathbb{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 $p_{w}(z)$ when we do not wish to distinguish between inputs and labels.

Given a prior on weights $\pi(w)$, Bayes law gives the posterior $p(w\,|\,x^{n},y^{n})\propto p(y^{n}\,|\,x^{n},w)\pi(w).$ The Fisher Information Matrix (FIM) $g\in\mathbb{R}^{p\times p}$ has entries $g(w)_{kl}=$

It can be used to define the Jeffreys prior $\pi_{J}(w)\propto\sqrt{\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 uninformative prior, e.g., it imposes some generic structure on the problem (reparameterization invariance).

To make the choice of a prior more objective, Bernardo (1979) suggested that uninformative priors should maximize some divergence, say the Kullback-Leibler (KL) divergence $\text{KL}(p(w\,|\,z),\pi(w))=\int\differential{w}p(w\,|\,z)\log\left(p(w\,|\,z)/\pi(w)\right)$, 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

where $p(z)=\int\differential{w}\pi(w)p(z\,|\,w)$ and $H(w)=-\int\differential{w}\pi(w)\log\pi(w)$ is the Shannon entropy; the conditional entropy $H(w\,|\,z)$ is defined analogously. Mutual information is a natural quantity for measuring the amount of missing information about $w$ provided by data $z$ if the initial belief was $\pi$. The prior $\pi^{*}(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 $\hat{P}_{n}$ but only depends on their distribution $P$.

The objective to calculate reference prior $\pi^{*}$ above may not be analytically tractable and therefore Bernardo also suggested computing $n$-reference priors. We call $n$ the “order” and deliberately overload the notation for the number of samples $n$; the reason will be clear soon.

using $n$ samples and then setting $\pi^{*}:=\lim_{n\to\infty}\pi_{n}^{*}$ under appropriate technical conditions (Berger et al., 1988). Reference priors are asymptotically equivalent to Jeffreys prior for one-dimensional problems. In general, they differ 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).

The Blahut-Arimoto algorithm (Arimoto, 1972; Blahut, 1972) is a method for maximizing functionals like Eq. 1 and leads to iterations of the form $\pi^{t+1}(w)\propto\exp\left(\text{KL}(p(z\,|\,w),p(z))\right)\pi^{t}(w)$. It is typically implemented for discrete variables, e.g., in the Information Bottleneck (Tishby et al., 1999). In this case, maximizing mutual information is a convex problem and therefore the BA algorithm is guaranteed to converge. Such discretization is difficult for high-dimensional deep networks. We therefore implement the BA algorithm using particles; see Remark 2.

Rigorous theoretical development of reference priors has been done in the statistics literature. We focus on their applications. We however mention some technical conditions under which our development remains meaningful.

A reference prior does not exist if $I_{\pi}(w;z^{n})$ is infinite (Berger et al., 1988). For the concept of a reference prior to remain meaningful, we make the following technical assumptions. (i) $\pi$ is supported on a compact set $\Omega\subset\mathbb{R}^{p}$, and (ii) if $p_{\pi}(z^{n})=\int_{\Omega}\differential{w}\pi(w)p(z^{n}\,|\,w)$ is the marginal, then $\text{KL}(p_{w},p_{\pi})$ is a continuous function of $w$ for any $\pi$. Under these conditions, the $n$-order prior $\pi_{n}^{*}$ exists and $I_{\pi_{n}}(w;z^{n})$ is finite; see (Zhang, 1994, Lemma 2.14). Now assume that $\pi_{n}^{*}$ exists and is unique up to a set of measure zero. Let $\Omega_{n}=\{w\in\Omega:\pi_{n}^{*}(w)>0\}$ be the support of $\pi_{n}^{*}$ and $z^{n}$ be a discrete random variable with $C$ atoms. If $\{p(z^{n}\,|\,w):w\in\Omega_{n}\}$ is compact, then $\pi_{n}^{*}$ is discrete with no more than $C$ atoms (Zhang, 1994, Lemma 2.18)).

One cannot directly visualize the high-dimensional particles $w$ in $\pi_{n}^{*}$. But we can think of each particle $w$ as representing a probability distribution $f(w)\in\mathbb{R}^{nC}$ given by

and use a method for visualizing such distributions developed in Quinn et al. (2019) that computes a principal component analysis (PCA) of such vectors $\{f(w^{1}),\ldots,f(w^{K})\}$ shown in Fig. 2. See Appendix C for more details.

This experiment demonstrates that we can instantiate reference priors for deep networks in a scalable fashion even for a large number of particles $K$. It provides a visual understanding of how atoms of the prior are diverse models in prediction space, just like the atoms in Fig. 1.

Consider the situation where we are given inputs $x^{n}$, their corresponding labels $y^{n}$ and unlabeled inputs $x^{u}$. Our goal is semi-supervised learning, i.e., to use $x^{u}$ to build a prior $\pi^{*}(w)$ that selects models that can be learned using the labeled data $(x^{n},y^{n})$. Recall that since $\pi^{*}$ is a prior, it should not depend on $(x^{n},y^{n})$. Just like the construction of the reference prior in Section 2.2, we can maximize

where $p_{\pi}(y^{n}\,|\,x^{n})=\int\differential{w}\pi(w)\prod_{i=1}^{n}p(y_{i}\,|\,x_{i},w)$ and likewise for $p_{\pi}(y^{u}\,|\,x^{u})$. The first step is simply the definition of $I_{\pi}$: it is the KL-divergence of the posterior after seeing $(x^{n},y^{n})$ with respect to the prior $\pi(w)$. The second step is the key idea and its rationale is as follows. If we know that inputs $x^{u}$ and $x^{n}$ come from the same task, then we can use samples $x^{u}$ to compute the expectation over $x^{n}$. For the same reason, we can average over outputs $y^{u}$ which are predicted by the network in place of the fixed labels $y^{n}$. Let us emphasize that both $x^{u}$ and $y^{u}$ are averaged out in the objective above. Predictions on new samples $x$ are made using the Bayesian posterior predictive distribution

We first develop the idea using generic random variables $z^{n}$. Consider a situation when we see 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 Section 3.3 in a natural way to address this question. We can maximize the KL-divergence between the posterior of the second stage and the posterior after the first stage, on average, over 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 $\pi(w)$. First, notice that $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

where $p(w\,|\,z^{m})\propto p(z^{m}\,|\,w)\pi(w)$ and $p(z^{n}\,|\,z^{m})=\int\differential{w}p(z^{n}\,|\,w)p(w\,|\,z^{m})$. The key observation is that if the reference prior Eq. 2 has a unique solution, we should have that the optimal $p(w\,|\,z^{m})\equiv\pi_{n}^{*}(w)$. This leads to

This prior puts less probability on regions which have high likelihood on old data $z^{m}$ whereby 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 also have $\pi_{n\,|\,m}^{*}=\pi_{n}^{*}$ for $m=0$ which is consistent with Eq. 2. As $m\to\infty$, this prior ignores the part of the weight space that was ideal for $z^{m}$. See Section D.3 for an example.

Consider the two-stage experiment where in the first stage we obtain $m$ samples $(x^{m}_{s},y^{m}_{s})$ from a “source” task $P^{s}$ and the second stage consists of $n$ samples $(x^{n}_{t},y^{n}_{t})$ from the “target” task $P^{t}$. Our goal is to calculate a prior $\pi(w)$ that best utilizes the target task data.

Bayesian inference for this problem involves first computing the posterior $p(w\,|\,x^{m}_{s},y^{m}_{s})\propto p(y^{m}_{s}\,|\,w,x^{m}_{s})\pi(w)$ from the source task and then using it as a prior to compute the posterior for the target task $p(w\,|\,x^{n}_{t},y^{n}_{t},x^{m}_{s},y^{m}_{s})$. Just like Section 2.2, the key idea again is to maximize the KL-divergence between the two posteriors $\text{KL}\left(p(w\,|\,x^{n}_{t},y^{n}_{t},x^{m}_{s},y^{m}_{s}),\ p(w\,|\,x^{m}_{s},y^{m}_{s})\right)$, but averaged over samples $x^{m}_{s}$ and $x^{n}_{t}$.

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 that we have developed.

We evaluate on CIFAR-10 and CIFAR-100 (Krizhevsky, 2009). For SSL, we use 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 five-way classification tasks from CIFAR-100 and use 1000 labeled samples from the source and 100 labeled samples from the target task. All experiments use the WRN 28-2 architecture (Zagoruyko and Komodakis, 2016), same as in Berthelot et al. (2019b).

For all our experiments, the reference prior is of order $n=2$ and has $K=4$ particles. We run all our methods for 200 epochs, with $\tau=1/3$ in Eq. 14 and $\alpha=0.1$ in Eq. 3. We set $\gamma=(1-\tau^{2})^{-1}$ as discussed in Section 3.6. For inference, each particle maintains an exponential moving average (EMA) of the weights (this is common in SSL (Tarvainen and Valpola, 2017)). Appendix A provides more details.

Just like we did in Section 3.6 for SSL, we instantiate Eq. 9 and Eq. 11, by combining prior selection, pretraining on the source task and likelihood of the target task, into one objective,

where $\gamma=1/2$ and $\beta=1/2$ are hyper-parameters, $(x^{m}_{s},y^{m}_{s})$ are labeled data from the source task ($m=1000$), $(x^{n}_{t},y^{n}_{t})$ are labeled data from the target task ($n=100$) and $x^{u}_{t}$ are unlabeled samples from the target task (all other samples).

This section presents ablation and analysis experiments for SSL on CIFAR-10 with 1000 labeled samples. We study the reference prior for different settings (i) varying the order $n$ of the prior, (ii) varying the number of particles in the BA algorithm ($K$), (iii) exponential moving averaging of the weights for each particle. We also study the two entropy terms in the reference prior objective individually.

We use a reference prior of order $n=2$ in all our experiments. We see in Table 2 that changing the order of the prior leads to marginal (about 1%) changes in the accuracy.

We next vary the number of particles in the prior in Table 3 and find that the accuracy is relatively consistent when the number of particles varies from $K=2$ to $K=16$. This seems surprising because a reference prior ideally should have an infinite number of atoms, when it approximates Jeffreys prior. We should not a priori expect $K=2$ particles to be sufficient to span the prediction space of deep networks. But our experiment in  Fig. 2 provides insight into this phenomenon. It shows that the manifold of diverse predictions is low-dimensional. Particles of the reference prior only need to span these few dimension and we can fruitfully implement our approach using very few particles.