# Principled Weight Initialization for
Hypernetworks

Oscar Chang, Lampros Flokas, Hod Lipson
Columbia University
New York, NY 10027
{oscar.chang, lf2540<EMAIL_ADDRESS>

###### Abstract

Hypernetworks are meta neural networks that generate weights for a main neural
network in an end-to-end differentiable manner. Despite extensive applications
ranging from multi-task learning to Bayesian deep learning, the problem of
optimizing hypernetworks has not been studied to date. We observe that
classical weight initialization methods like Glorot & Bengio (2010) and He et
al. (2015), when applied directly on a hypernet, fail to produce weights for
the mainnet in the correct scale. We develop principled techniques for weight
initialization in hypernets, and show that they lead to more stable mainnet
weights, lower training loss, and faster convergence.

##  1 Introduction

Meta-learning describes a broad family of techniques in machine learning that
deals with the problem of learning to learn. An emerging branch of meta-
learning involves the use of _hypernetworks_ , which are meta neural networks
that generate the weights of a main neural network to solve a given task in an
end-to-end differentiable manner. Hypernetworks were originally introduced by
Ha et al. (2016) as a way to induce weight-sharing and achieve model
compression by training the same meta network to learn the weights belonging
to different layers in the main network. Since then, hypernetworks have found
numerous applications including but not limited to: weight pruning (Liu et
al., 2019), neural architecture search (Brock et al., 2017; Zhang et al.,
2018), Bayesian neural networks (Krueger et al., 2017; Ukai et al., 2018;
Pawlowski et al., 2017; Henning et al., 2018; Deutsch et al., 2019), multi-
task learning (Pan et al., 2018; Shen et al., 2017; Klocek et al., 2019; Serrà
et al., 2019; Meyerson & Miikkulainen, 2019), continual learning (von Oswald
et al., 2019), generative models (Suarez, 2017; Ratzlaff & Fuxin, 2019),
ensemble learning (Kristiadi & Fischer, 2019), hyperparameter optimization
(Lorraine & Duvenaud, 2018), and adversarial defense (Sun et al., 2017).

Despite the intensified study of applications of hypernetworks, the problem of
optimizing them to this day remains significantly understudied. In fact, even
the problem of initializing hypernetworks has not been studied. Given the lack
of principled approaches, prior work in the area is mostly limited to ad-hoc
approaches based on trial and error (c.f. Section 3). For example, it is
common to initialize the weights of a hypernetwork by sampling a “small”
random number. Nonetheless, these ad-hoc methods do lead to successful
hypernetwork training primarily due to the use of the Adam optimizer (Kingma &
Ba, 2014), which has the desirable property of being invariant to the scale of
the gradients. However, even Adam will not work if the loss diverges (i.e.
overflow) at initialization, which will happen in sufficiently big models. The
normalization of badly scaled gradients also results in noisy training
dynamics where the loss function suffers from bigger fluctuations during
training compared to vanilla stochastic gradient descent (SGD). Wilson et al.
(2017); Reddi et al. (2018) showed that while adaptive optimizers like Adam
may exhibit lower training error, they fail to generalize as well to the test
set as non-adaptive gradient methods. Moreover, Adam incurs a computational
overhead and requires 3X the amount of memory for the gradients compared to
vanilla SGD.

Small random number sampling is reminiscent of early neural network research
(Rumelhart et al., 1986) before the advent of classical weight initialization
methods like Xavier init (Glorot & Bengio, 2010) and Kaiming init (He et al.,
2015). Since then, a big lesson learned by the neural network optimization
community is that architecture specific initialization schemes are important
to the robust training of deep networks, as shown recently in the case of
residual networks (Zhang et al., 2019). In fact, weight initialization for
hypernetworks was recognized as an outstanding open problem by prior work
(Deutsch et al., 2019) that had questioned the suitability of classical
initialization methods for hypernetworks.

Our results We show that when classical methods are used to initialize the
weights of hypernetworks, they fail to produce mainnet weights in the correct
scale, leading to exploding activations and losses. This is because classical
network weights transform one layer’s activations into another, while hypernet
weights have the added function of transforming the hypernet’s activations
into the mainnet’s weights. Our solution is to develop principled techniques
for weight initialization in hypernetworks based on variance analysis. The
hypernet case poses unique challenges. For example, in contrast to variance
analysis for classical networks, the case for hypernetworks can be
asymmetrical between the forward and backward pass. The asymmetry arises when
the gradient flow from the mainnet into the hypernet is affected by the
biases, whereas in general, this does not occur for gradient flow in the
mainnet. This underscores again why architecture specific initialization
schemes are essential. We show both theoretically and experimentally that our
methods produce hypernet weights in the correct scale. Proper initialization
mitigates exploding activations and gradients or the need to depend on Adam.
Our experiments reveal that it leads to more stable mainnet weights, lower
training loss, and faster convergence.

Section 2 briefly covers the relevant technical preliminaries, and Section 3
reviews problems with the ad-hoc methods currently deployed by hypernetwork
practitioners. We derive novel weight initialization formulae for
hypernetworks in Section 4, empirically evaluate our proposed methods in
Section 5, and finally conclude in Section 6.

##  2 Preliminaries

######  Definition.

A hypernetwork is a meta neural network $H$ with its own parameters $\phi$
that generates the weights of a main network $\theta$ from some embedding $e$
in a differentiable manner: $\theta=H_{\phi}(e)$. Unlike a classical network,
in a hypernetwork, the weights of the main network are not model parameters.
Thus the gradients $\Delta\theta$ have to be further backpropagated to the
weights of the hypernetwork $\Delta\phi$, which is then trained via gradient
descent $\phi_{t+1}=\phi_{t}-\lambda\Delta\phi_{t}$.

This fundamental difference suggests that conventional knowledge about neural
networks may not apply directly to hypernetworks and novel ways of thinking
about weight initialization, optimization dynamics and architecture design for
hypernetworks are sorely needed.

###  2.1 Ricci Calculus

We propose the use of Ricci calculus, as opposed to the more commonly used
matrix calculus, as a suitable mathematical language for thinking about
hypernetworks. Ricci calculus is useful because it allows us to reason about
the derivatives of higher-order tensors with notational ease. For readers not
familiar with the index-based notation of Ricci calculus, please refer to Laue
et al. (2018) for a good introduction to the topic written from a machine
learning perspective.

For a general nth-order tensor $T^{i_{1},\ldots,i_{k},\ldots,i_{n}}$, we use
$\text{d}_{i_{k}}$ to refer to the dimension of the index set that $i_{k}$ is
drawn from. We include explicit summations where the relevant expressions
might be ambiguous, and use Einstein summation convention otherwise. We use
square brackets to denote different layers for added clarity, so for example
$W[t]$ denotes the $t$-th weight layer.

###  2.2 Xavier Initialization

Glorot & Bengio (2010) derived weight initialization formulae for a
feedforward neural network by conducting a variance analysis over activations
and gradients. For a linear layer ${y^{i}}={W^{i}_{j}}{x^{j}}+{b^{i}}$,
suppose we make the following Xavier Assumptions at initialization: (1) The
${W^{i}_{j}}$, ${x^{j}}$, and ${b^{i}}$ are all independent of each other. (2)
$\forall{i,j}:\mathbb{E}[{W^{i}_{j}}]=0$. (3)
$\forall{j}:\mathbb{E}[{x^{j}}]=0$. (4) $\forall{i}:{b^{i}}=0$.

Then, $\mathbb{E}[{y^{i}}]=0$ and
$\text{Var}({y^{i}})=\text{d}_{j}\text{Var}({W^{i}_{j}})\text{Var}({x^{j}}).$
To keep the variance of the output and input activations the same, i.e.
$\text{Var}({y^{i}})=\text{Var}({x^{j}})$, we have to sample ${W^{i}_{j}}$
from a distribution whose variance is equal to the reciprocal of the fan-in:
$\text{Var}({W^{i}_{j}})=\frac{1}{\text{d}_{j}}$.

If analogous assumptions hold for the backward pass, then to keep the variance
of the output and input gradients the same, we have to sample ${W^{i}_{j}}$
from a distribution whose variance is equal to the reciprocal of the fan-out:
$\text{Var}({W^{i}_{j}})=\frac{1}{\text{d}_{i}}$.

Thus, the forward pass and backward pass result in symmetrical formulae.
Glorot & Bengio (2010) proposed an initialization based on their harmonic
mean: $\text{Var}({W^{i}_{j}})=\frac{2}{\text{d}_{j}+\text{d}_{i}}$.

In general, a feedforward network is non-linear, so these assumptions are
strictly invalid. But odd activation functions with unit derivative at $0$
results in a roughly linear regime at initialization.

###  2.3 Kaiming Initialization

He et al. (2015) extended Glorot & Bengio (2010)’s analysis by looking at the
case of ReLU activation functions, i.e. ${y^{i}}={W^{i}_{j}}\
\text{ReLU}({x^{j}})+{b^{i}}$. We can write $z^{j}=\text{ReLU}(x^{j})$ to get

$\begin{split}\text{Var}(y^{i})&=\sum_{j}\mathbb{E}[({z^{j}})^{2}]\text{Var}({W^{i}_{j}})=\sum_{j}\frac{1}{2}\
\mathbb{E}[(x^{j})^{2}]\text{Var}(W^{i}_{j})=\frac{1}{2}\
\text{d}_{j}\text{Var}({W^{i}_{j}})\text{Var}({x^{j}}).\\\ \end{split}$

This results in an extra factor of $2$ in the variance formula. $W^{i}_{j}$
have to be symmetric around $0$ to enforce Xavier Assumption 3 as the
activations and gradients propagate through the layers. He et al. (2015)
argued that both the forward or backward version of the formula can be
adopted, since the activations or gradients will only be scaled by a depth-
independent factor. For convolutional layers, we have to further divide the
variance by the size of the receptive field.

‘Xavier init’ and ‘Kaiming init’ are terms that are sometimes used
interchangeably. Where there might be confusion, we will refer to the forward
version as fan-in init, the backward version as fan-out init, and the harmonic
mean version as harmonic init.

##  3 Review of Current Methods

In the seminal Ha et al. (2016) paper, the authors identified two distinct
classes of hypernetworks: dynamic (for recurrent networks) and static (for
convolutional networks). They proposed Orthogonal init (Saxe et al., 2013) for
the dynamic class, but omitted discussion of initialization for the static
class. The static class has since proven to be the dominant variant, covering
all kinds of non-recurrent networks (not just convolutional), and thus will be
the central object of our investigation.

Through an extensive literature and code review, we found that hypernet
practitioners mostly depend on the Adam optimizer, which is invariant to and
normalizes the scale of gradients, for training and resort to one of four
weight initialization methods:

  1. M1

Xavier or Kaiming init (as found in Pawlowski et al. (2017); Balazevic et al.
(2018); Serrà et al. (2019); von Oswald et al. (2019)).

  2. M2

Small random values (as found in Krueger et al. (2017); Lorraine & Duvenaud
(2018)).

  3. M3

Kaiming init, but with the output layer scaled by $\frac{1}{10}$ (as found in
Ukai et al. (2018)).

  4. M4

Kaiming init, but with the hypernet embedding set to be a suitably scaled
constant (as found in Meyerson & Miikkulainen (2019)).

M1 uses classical neural network initialization methods to initialize
hypernetworks. This fails to produce weights for the main network in the
correct scale. Consider the following illustrative example of a one-layer
linear hypernet generating a linear mainnet with $T+1$ layers, given
embeddings sampled from a standard normal distribution and weights sampled
entry-wise from a zero-mean distribution. We leave the biases out for now, and
assume the input data $x[1]$ is standardized.

$\begin{split}x[t+1]^{i_{t+1}}&=W[t]^{i_{t+1}}_{i_{t}}x[t]^{i_{t}},\qquad
W[t]^{i_{t+1}}_{i_{t}}=H[t]^{i_{t+1}}_{i_{t}k_{t}}e[t]^{k_{t}},\qquad 1\leq
t\leq T.\\\
\text{Var}(x[T+1]^{i_{t+1}})&=\text{Var}(x[1]^{i_{1}})\prod_{t=1}^{T}\text{d}_{i_{t}}\text{Var}(W[t]^{i_{t+1}}_{i_{t}})=\text{Var}(x[1]^{i_{1}})\prod_{t=1}^{T}\text{d}_{i_{t}}\text{d}_{k_{t}}\text{Var}(H[t]^{i_{t+1}}_{i_{t}k_{t}}).\end{split}$
(1)

In this case, if the variance of the weights in the hypernet
$\text{Var}(H[t]^{i_{t+1}}_{i_{t}k_{t}})$ is equal to the reciprocal of the
fan-in $\text{d}_{k_{t}}$, then the variance of the activations
$\text{Var}(x[T+1]^{i_{t+1}})=\prod_{t=1}^{T}\text{d}_{i_{t}}$ explodes. If it
is equal to the reciprocal of the fan-out
$\text{d}_{i_{t}}\text{d}_{i_{t+1}}$, then the activation variance
$\text{Var}(x[T+1]^{i_{t+1}})=\prod_{t=1}^{T}\frac{\text{d}_{k_{t}}}{\text{d}_{i_{t}+1}}$
is likely to vanish, since the size of the embedding vector is typically small
relatively to the width of the mainnet weight layer being generated.

Where the fan-in is of a different scale than the fan-out, the harmonic mean
has a scale close to that of the smaller number. Therefore, the fan-in, fan-
out, and harmonic variants of Xavier and Kaiming init will all result in
activations and gradients that scale exponentially with the depth of the
mainnet.

M2 and M3 introduce additional hyperparameters into the model, and the ad-hoc
manner in which they work is reminiscent of pre deep learning neural network
research, before the introduction of classical initialization methods like
Xavier and Kaiming init. This ad-hoc manner is not only inelegant and consumes
more compute, but will likely fail for deeper and more complex hypernetworks.

M4 proposes to set the embeddings $e[t]^{k_{t}}$ to a suitable constant
($\text{d}_{i_{t}}^{-1/2}$ in this case), such that both
$W[t]^{i_{t+1}}_{i_{t}}$ and $H[t]^{i_{t+1}}_{i_{t}k_{t}}$ can seem to be
initialized with the same variance as Kaiming init. This ensures that the
variance of the activations in the mainnet are preserved through the layers,
but the restrictions on the embeddings might not be desirable in many
applications.

Luckily, the fix appears simple — set
$\text{Var}(H[t]^{i_{t+1}}_{i_{t}k_{t}})=\frac{1}{\text{d}_{i_{t}}\text{d}_{k_{t}}}$.
This results in the variance of the generated weights in the mainnet
$\text{Var}(W[t]^{i_{t+1}}_{i_{t}})=\frac{1}{\text{d}_{i_{t}}}$ resembling
conventional neural networks initialized with fan-in init. This suggests a
general hypernet weight initialization strategy: initialize the weights of the
hypernet such that the mainnet weights approximate classical neural network
initialization. We elaborate on and generalize this intuition in Section 4.

##  4 Hyperfan Initialization

Most hypernetwork architectures use a linear output layer so that gradients
can pass from the mainnet into the hypernet directly without any non-
linearities. We make use of this fact in developing methods called _hyperfan-
in init_ and _hyperfan-out init_ for hypernetwork weight initialization based
on the principle of variance analysis.

###  4.1 Hyperfan-in

######  Proposition.

Suppose a hypernetwork comprises a linear output layer. Then, the variance
between the input and output activations of a linear layer in the mainnet
$y^{i}=W^{i}_{j}x^{j}+b^{i}$ can be preserved using fan-in init in the
hypernetwork with appropriately scaled output layers.

Case 1. The hypernet generates the weights but not the biases of the mainnet.
The bias in the mainnet is initialized to zero. We can write the weight
generation in the form $W^{i}_{j}=H^{i}_{jk}h(e)^{k}+\beta^{i}_{j}$ where $h$
computes all but the last layer of the hypernet and $(H,\beta)$ form the
output layer. We make the following Hyperfan Assumptions at initialization:
(1) Xavier assumptions hold for all the layers in the hypernet. (2) The
${H^{i}_{jk}}$, $h(e)^{k}$, $\beta^{i}_{j}$, ${x^{j}}$, and ${b^{i}}$ are all
independent of each other. (3) $\forall{i,j,k}:\mathbb{E}[{H^{i}_{jk}}]=0$.
(4) $\mathbb{E}[{x^{j}}]=0$. (5) $\forall{i}:{b^{i}}=0$.

Use fan-in init to initialize the weights for $h$. Then,
$\text{Var}(h(e)^{k})=\text{Var}(e^{l})$. If we initialize $H$ with the
formula
$\text{Var}(H^{i}_{jk})=\frac{1}{\text{d}_{j}\text{d}_{k}\text{Var}(e^{l})}$
and $\beta$ with zeros, we arrive at
$\text{Var}(W^{i}_{j})=\frac{1}{\text{d}_{j}}$, which is the formula for fan-
in init in the mainnet. The Hyperfan assumptions imply the Xavier assumptions
hold in the mainnet, thus preserving the input and output activations.

$\begin{split}\text{Var}({y^{i}})&=\sum_{j}\text{Var}({W^{i}_{j}})\text{Var}({x^{j}})=\sum_{j}\sum_{k}\text{Var}(H^{i}_{jk})\text{Var}(h(e)^{k})\text{Var}({x^{j}})\\\
&=\sum_{j}\sum_{k}\frac{1}{\text{d}_{j}\text{d}_{k}\text{Var}(e^{l})}\text{Var}(e^{l})\text{Var}({x^{j}})=\text{Var}({x^{j}}).\end{split}$
(2)

Case 2. The hypernet generates both the weights and biases of the mainnet. We
can write the weight and bias generation in the form
$W^{i}_{j}=H^{i}_{jk}h(e[1])^{k}+\beta^{i}_{j}$ and
$b^{i}=G^{i}_{l}g(e[2])^{l}+\gamma^{i}$ respectively, where $h$ and $g$
compute all but the last layer of the hypernet, and $(H,\beta)$ and
$(G,\gamma)$ form the output layers. We modify Hyperfan Assumption 2 so it
includes $G^{i}_{l}$, $g(e[2])^{l}$, and $\gamma^{i}$, and further assume
$\text{Var}(x^{j})=1$, which holds at initialization with the common practice
of data standardization.

Use fan-in init to initialize the weights for $h$ and $g$. Then,
$\text{Var}(h(e[1])^{k})=\text{Var}(e[1]^{m})$ and
$\text{Var}(g(e[2])^{l})=\text{Var}(e[2]^{n})$. If we initialize $H$ with the
formula
$\text{Var}(H^{i}_{jk})=\frac{1}{2\text{d}_{j}\text{d}_{k}\text{Var}(e[1]^{m})}$,
$G$ with the formula
$\text{Var}(G^{i}_{l})=\frac{1}{2\text{d}_{l}\text{Var}(e[2]^{n})}$, and
$\beta,\gamma$ with zeros, then the input and output activations in the
mainnet can be preserved.

$\begin{split}\text{Var}({y^{i}})&=\sum_{j}\left[\text{Var}({W^{i}_{j}})\text{Var}({x^{j}})\right]+\text{Var}({b^{i}})\\\
&=\sum_{j}\left[\sum_{k}\text{Var}(H^{i}_{jk})\text{Var}(h(e[1])^{k})\text{Var}({x^{j}})\right]+\sum_{l}\text{Var}(G^{i}_{l})\text{Var}(g(e[2])^{l})\\\
&=\sum_{j}\left[\sum_{k}\frac{1}{2\text{d}_{j}\text{d}_{k}\text{Var}(e[1]^{m})}\text{Var}(e[1]^{m})\text{Var}({x^{j}})\right]+\sum_{l}\frac{1}{2\text{d}_{l}\text{Var}(e[2]^{n})}\text{Var}(e[2]^{n})\\\
&=\frac{1}{2}\text{Var}({x^{j}})+\frac{1}{2}=\text{Var}({x^{j}}).\end{split}$
(3)

If we initialize $G^{i}_{j}$ to zeros, then its contribution to the variance
will increase during training, causing exploding activations in the mainnet.
Hence, we prefer to introduce a factor of $1/2$ to divide the variance between
the weight and bias generation, where the variance of each component is
allowed to either decrease or increase during training. This becomes a problem
if the variance of the activations in the mainnet deviates too far away from
$1$, but we found that it works well in practice.

###  4.2 Hyperfan-out

Case 1. The hypernet generates the weights but not the biases of the mainnet.
A similar derivation can be done for the backward pass using analogous
assumptions on gradients flowing

$\begin{split}\text{in the mainnet:}\quad&\frac{\partial L}{\partial
x[t]^{i_{t}}}=\frac{\partial L}{\partial
x[t+1]^{i_{t+1}}}W[t]^{i_{t+1}}_{i_{t}},\\\ \text{through mainnet
weights:}\quad&\frac{\partial L}{\partial
W[t]^{i_{t+1}}_{i_{t}}}=\frac{\partial L}{\partial
x[t+1]^{i_{t+1}}}x[t]^{i_{t}},\frac{\partial L}{\partial
h[t](e)^{k_{t}}}=\frac{\partial L}{\partial
W[t]^{i_{t+1}}_{i_{t}}}H[t]^{i_{t+1}}_{i_{t}k_{t}},\\\ \text{and through
mainnet biases:}\quad&\frac{\partial L}{\partial
b[t]^{i_{t+1}}}=\frac{\partial L}{\partial x[t+1]^{i_{t+1}}},\frac{\partial
L}{\partial g[t](e)^{l_{t}}}=\frac{\partial L}{\partial
b[t]^{i_{t+1}}}G[t]^{i_{t+1}}_{l_{t}}.\end{split}$ (4)

If we initialize the output layer $H$ with the analogous hyperfan-out formula
$\text{Var}(H[t]^{i_{t+1}}_{i_{t}k_{t}})=\frac{1}{\text{d}_{i_{t+1}}\text{d}_{k_{t}}\text{Var}(e^{k_{t}})}$
and the rest of the hypernet with fan-in init, then we can preserve input and
output gradients on the mainnet: $\text{Var}(\frac{\partial L}{\partial
x[t]^{i_{t}}})=\text{Var}(\frac{\partial L}{\partial x[t+1]^{i_{t+1}}})$.
However, note that the gradients will shrink when flowing from the mainnet to
the hypernet: $\text{Var}(\frac{\partial L}{\partial
h[t](e)^{k_{t}}})=\frac{\text{d}_{i_{t}}}{\text{d}_{k_{t}}\text{Var}(e^{k_{t}})}\text{Var}(\frac{\partial
L}{\partial W[t]^{i_{t+1}}_{i_{t}}})$, and scaled by a depth-independent
factor due to the use of fan-in rather than fan-out init.

Case 2. The hypernet generates both the weights and biases of the mainnet. In
the classical case, the forward version (fan-in init) and the backward version
(fan-out init) are symmetrical. This remains true for hypernets if they only
generated the weights of the mainnet. However, if they were to also generate
the biases, then the symmetry no longer holds, since the biases do not affect
the gradient flow in the mainnet but they do so for the hypernet (c.f.
Equation 4). Nevertheless, we can initialize $G$ so that it helps hyperfan-out
init preserve activation variance on the forward pass as much as possible
(keeping the assumption that $\text{Var}(x^{j})=1$ as before):

$\begin{split}\text{Var}(y^{i})&=\sum_{j}\bigl{[}\text{Var}(W^{i}_{j}x^{j})\bigr{]}+\text{Var}(b^{i})\\\
&=\text{d}_{j}\text{d}_{k}\text{Var}(e[1]^{m})\text{Var}(H[\text{hyperfan-
out}]^{i}_{jk})\text{Var}(x^{j})+\text{d}_{l}\text{Var}(e[2]^{n})\text{Var}(G^{i}_{l})\\\
&=\text{d}_{j}\text{d}_{k}\text{Var}(e[1]^{m})\text{Var}(H[\text{hyperfan-
in}]^{i}_{jk})\text{Var}(x^{j})\\\ &\text{Plugging in the formulae for
Hyperfan-in and Hyperfan-out from above, we get}\\\
\implies\text{Var}(G^{i}_{l})&=\frac{(1-\frac{\text{d}_{j}}{\text{d}_{i}})}{\text{d}_{l}\text{Var}(e[2]^{n})}.\\\
\end{split}$ (5)

We summarize the variance formulae for hyperfan-in and hyperfan-out init in
Table 1. It is not uncommon to re-use the same hypernet to generate different
parts of the mainnet, as was originally done in Ha et al. (2016). We discuss
this case in more detail in Appendix Section A.

Table 1: Hyperfan-in and Hyperfan-out Variance Formulae for $W^{i}_{j}=H^{i}_{jk}h(e[1])^{k}+\beta^{i}_{j}$. If $y^{i}=\text{ReLU}(W^{i}_{j}x^{j}+b^{i})$, then $\mathbbm{1}_{\text{ReLU}}=1$, else if $y^{i}=W^{i}_{j}x^{j}+b^{i}$, then $\mathbbm{1}_{\text{ReLU}}=0$. If $b^{i}=G^{i}_{l}g(e[2])^{l}+\gamma^{i}$, then $\mathbbm{1}_{\text{HBias}}=1$, else if $b^{i}=0$, then $\mathbbm{1}_{\text{HBias}}=0$. We initialize $h$ and $g$ with fan-in init, and $\beta^{i}_{j},\gamma^{i}=0$. For convolutional layers, we have to further divide $\text{Var}(H^{i}_{jk})$ by the size of the receptive field. Uniform init: $X\sim\mathcal{U}(-\sqrt{3\text{Var}(X)},\sqrt{3\text{Var}(X)})$. Normal init: $X\sim\mathcal{N}(0,\text{Var}(X))$. Initialization | Variance Formula | Initialization | Variance Formula
---|---|---|---
Hyperfan-in | $\text{Var}(H^{i}_{jk})=\frac{2^{\mathbbm{1}_{\text{ReLU}}}}{2^{\mathbbm{1}_{\text{HBias}}}\text{d}_{j}\text{d}_{k}\text{Var}(e[1]^{m})}$ | Hyperfan-out | $\text{Var}(H^{i}_{jk})=\frac{2^{\mathbbm{1}_{\text{ReLU}}}}{\text{d}_{i}\text{d}_{k}\text{Var}(e[1]^{m})}$
Hyperfan-in | $\text{Var}(G^{i}_{l})=\frac{2^{\mathbbm{1}_{\text{ReLU}}}}{2\text{d}_{l}\text{Var}(e[2]^{n})}$ | Hyperfan-out | $\text{Var}(G^{i}_{l})=\text{max}(\frac{2^{\mathbbm{1}_{\text{ReLU}}}(1-\frac{\text{d}_{j}}{\text{d}_{i}})}{\text{d}_{l}\text{Var}(e[2]^{n})},0)$

##  5 Experiments

We evaluated our proposed methods on four sets of experiments involving
different use cases of hypernetworks: feedforward networks, continual
learning, convolutional networks, and Bayesian neural networks. In all cases,
we optimize with vanilla SGD and sample from the uniform distribution
according to the variance formula given by the init method. More experimental
details can be found in Appendix Section B.

###  5.1 Feedforward Networks on MNIST

As an illustrative first experiment, we train a feedforward network with five
hidden layers ($500$ hidden units), a hyperbolic tangent activation function,
and a softmax output layer, on MNIST across four different settings: (1) a
classical network with Xavier init, (2) a hypernet with Xavier init that
generates the weights of the mainnet, (3) a hypernet with hyperfan-in init
that generates the weights of the mainnet, (4) and a hypernet with hyperfan-
out init that generates the weights of the mainnet.

The use of hyperfan init methods on a hypernetwork reproduces mainnet weights
similar to those that have been trained from Xavier init on a classical
network, while the use of Xavier init on a hypernetwork causes exploding
activations right at the beginning of training (see Figure 1). Observe in
Figure 2 that when the hypernetwork is initialized in the proper scale, the
magnitude of generated weights stabilizes quickly. This in turn leads to a
more stable training regime, as seen in Figure 3. More visualizations of the
activations and gradients of both the mainnet and hypernet can be viewed in
Appendix Section B.1. Qualitatively similar observations were made when we
replaced the activation function with ReLU and Xavier with Kaiming init, with
Kaiming init leading to even bigger activations at initialization.

Suppose now the hypernet generates both the weights and biases of the mainnet
instead of just the weights. We found that this architectural change leads the
hyperfan init methods to take more time (but still less than Xavier init), to
generate stable mainnet weights (c.f. Figure 25 in the Appendix).

Figure 1: Mainnet Activations before the Start of Training on MNIST. Figure 2:
Evolution of Hypernet Output Layer Activations during Training on MNIST.
Xavier init results in unstable mainnet weights throughout training, while
hyperfan-in and hyperfan-out init result in mainnet weights that stabilize
quickly. Figure 3: Loss and Test Accuracy Plots on MNIST.

###  5.2 Continual Learning on Regression Tasks

Continual learning solves the problem of learning tasks in sequence without
forgetting prior tasks. von Oswald et al. (2019) used a hypernetwork to learn
embeddings for each task as a way to efficiently regularize the training
process to prevent catastrophic forgetting. We compare different
initialization schemes on their hypernetwork implementation, which generates
the weights and biases of a ReLU mainnet with two hidden layers to solve a
sequence of three regression tasks.

In Figure 4, we plot the training loss averaged over $15$ different runs, with
the shaded area showing the standard error. We observe that the hyperfan
methods produce smaller training losses at initialization and during training,
eventually converging to a smaller loss for each task.

Figure 4: Continual Learning Loss on a Sequence of Regression Tasks.

###  5.3 Convolutional Networks on CIFAR-10

Ha et al. (2016) applied a hypernetwork on a convolutional network for image
classification on CIFAR-10. We note that our initialization methods do not
handle residual connections, which were in their chosen mainnet architecture
and are important topics for future study. Instead, we implemented their
hypernetwork architecture on a mainnet with the All Convolutional Net
architecture (Springenberg et al., 2014) that is composed of convolutional
layers and ReLU activation functions.

After searching through a dense grid of learning rates, we failed to enable
the fan-in version of Kaiming init to train even with very small learning
rates. The fan-out version managed to begin delayed training, starting from
around epoch $270$ (see Figure 5). By contrast, both hyperfan-in and hyperfan-
out init led to successful training immediately. This shows a good init can
make it possible to successfully train models that would have otherwise been
unamenable to training on a bad init.

Figure 5: Loss and Test Accuracy Plots on CIFAR-10.

###  5.4 Bayesian Neural Networks on ImageNet

Bayesian neural networks improve model calibration and provide uncertainty
estimation, which guard against the pitfalls of overconfident networks. Ukai
et al. (2018) developed a Bayesian neural network by using a hypernetwork to
simulate an expressive prior distribution. We trained a similar hypernetwork
by applying Ukai et al. (2018)’s methods on ImageNet, but differed in our
choice of MobileNet (Howard et al., 2017) as a mainnet architecture that does
not have residual connections.

In the work of Ukai et al. (2018), it was noticed that even with the use of
batch normalization in the mainnet, classical initialization approaches still
led to diverging losses (due to exploding activations, c.f. Section 3). We
observe similar results in our experiment (see Figure 6) — the fan-in version
of Kaiming init, which is the default initialization in popular deep learning
libraries like PyTorch and Chainer, resulted in substantially higher initial
losses and led to slower training than the hyperfan methods. We found that the
observation still stands even when the last layer of the mainnet is not
generated by the hypernet. This shows that while batch normalization helps, it
is not the solution for a bad init that causes exploding activations. Our
approach solves this problem in a principled way, and is preferable to the
trial-and-error based heuristics that Ukai et al. (2018) had to resort to in
order to train their model.

Surprisingly, the fan-out version of Kaiming init led to similar results as
the hyperfan methods, suggesting that batch normalization might be sufficient
to correct the bad initializations that result in vanishing activations. That
being said, hypernet practitioners should not expect batch normalization to be
the panacea for problems caused by bad initialization, especially in memory-
constrained scenarios. In a Bayesian neural network application (especially in
hypernet architectures without relaxed weight-sharing), the blowup in the
number of parameters limits the use of big batch sizes, which is essential to
the performance of batch normalization (Wu & He, 2018). For example, in this
experiment, our hypernet model requires $32$ times as many parameters as a
classical MobileNet.

To the best of our knowledge, the interaction between batch normalization and
initialization is not well-understood, even in the classical case, and thus,
our findings prompt an interesting direction for future research.

Figure 6: Loss and Test Accuracy Plots on ImageNet.

In all our experiments, hyperfan-in and hyperfan-out both led to successful
hypernetwork training with SGD. We did not find a good reason to prefer one
over the other (similar to He et al. (2015)’s observation in the classical
case for fan-in and fan-out init).

##  6 Conclusion

For a long time, the promise of deep nets to learn rich representations of the
world was left unfulfilled due to the inability to train these models. The
discovery of greedy layer-wise pre-training (Hinton et al., 2006; Bengio et
al., 2007) and later, Xavier and Kaiming init, as weight initialization
strategies to enable such training was a pivotal achievement that kickstarted
the deep learning revolution. This underscores the importance of model
initialization as a fundamental step in learning complex representations.

In this work, we developed the first principled weight initialization methods
for hypernetworks, a rapidly growing branch of meta-learning. We hope our work
will spur momentum towards the development of principled techniques for
building and training hypernetworks, and eventually lead to significant
progress in learning meta representations. Other non-hypernetwork methods of
neural network generation (Stanley et al., 2009; Koutnik et al., 2010) can
also be improved by considering whether their generated weights result in
exploding activations and how to avoid that if so.

##  7 Acknowledgements

This research was supported in part by the US Defense Advanced Research
Project Agency (DARPA) Lifelong Learning Machines Program, grant
HR0011-18-2-0020. We thank Dan Martin and Yawei Li for helpful discussions,
and the ICLR reviewers for their constructive feedback.

## References

  * Balazevic et al. (2018) Ivana Balazevic, Carl Allen, and Timothy M Hospedales.  Hypernetwork knowledge graph embeddings.  _arXiv preprint arXiv:1808.07018_ , 2018.
  * Bengio et al. (2007) Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle.  Greedy layer-wise training of deep networks.  In _Advances in neural information processing systems_ , pp. 153–160, 2007.
  * Brock et al. (2017) Andrew Brock, Theodore Lim, James M Ritchie, and Nick Weston.  Smash: one-shot model architecture search through hypernetworks.  _arXiv preprint arXiv:1708.05344_ , 2017.
  * Deutsch et al. (2019) Lior Deutsch, Erik Nijkamp, and Yu Yang.  A generative model for sampling high-performance and diverse weights for neural networks.  _arXiv preprint arXiv:1905.02898_ , 2019.
  * Glorot & Bengio (2010) Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 249–256, 2010.
  * Ha et al. (2016) David Ha, Andrew Dai, and Quoc V Le.  Hypernetworks.  _arXiv preprint arXiv:1609.09106_ , 2016.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 1026–1034, 2015.
  * Henning et al. (2018) Christian Henning, Johannes von Oswald, João Sacramento, Simone Carlo Surace, Jean-Pascal Pfister, and Benjamin F Grewe.  Approximating the predictive distribution via adversarially-trained hypernetworks.  In _Bayesian Deep Learning Workshop, NeurIPS (Spotlight)_ , volume 2018, 2018.
  * Hinton et al. (2006) Geoffrey E Hinton, Simon Osindero, and Yee-Whye Teh.  A fast learning algorithm for deep belief nets.  _Neural computation_ , 18(7):1527–1554, 2006\.
  * Howard et al. (2017) Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  _arXiv preprint arXiv:1704.04861_ , 2017.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Klocek et al. (2019) Sylwester Klocek, Łukasz Maziarka, Maciej Wołczyk, Jacek Tabor, Marek Śmieja, and Jakub Nowak.  Hypernetwork functional image representation.  _arXiv preprint arXiv:1902.10404_ , 2019.
  * Koutnik et al. (2010) Jan Koutnik, Faustino Gomez, and Jürgen Schmidhuber.  Evolving neural networks in compressed weight space.  In _Proceedings of the 12th annual conference on Genetic and evolutionary computation_ , pp. 619–626. ACM, 2010.
  * Kristiadi & Fischer (2019) Agustinus Kristiadi and Asja Fischer.  Predictive uncertainty quantification with compound density networks.  _arXiv preprint arXiv:1902.01080_ , 2019.
  * Krueger et al. (2017) David Krueger, Chin-Wei Huang, Riashat Islam, Ryan Turner, Alexandre Lacoste, and Aaron Courville.  Bayesian hypernetworks.  _arXiv preprint arXiv:1710.04759_ , 2017.
  * Laue et al. (2018) Sören Laue, Matthias Mitterreiter, and Joachim Giesen.  Computing higher order derivatives of matrix and tensor expressions.  In _Advances in Neural Information Processing Systems_ , pp. 2755–2764, 2018.
  * Liu et al. (2019) Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang, Tim Kwang-Ting Cheng, and Jian Sun.  Metapruning: Meta learning for automatic neural network channel pruning.  _arXiv preprint arXiv:1903.10258_ , 2019.
  * Lorraine & Duvenaud (2018) Jonathan Lorraine and David Duvenaud.  Stochastic hyperparameter optimization through hypernetworks.  _arXiv preprint arXiv:1802.09419_ , 2018.
  * Meyerson & Miikkulainen (2019) Elliot Meyerson and Risto Miikkulainen.  Modular universal reparameterization: Deep multi-task learning across diverse domains.  _arXiv preprint arXiv:1906.00097_ , 2019.
  * Pan et al. (2018) Zheyi Pan, Yuxuan Liang, Junbo Zhang, Xiuwen Yi, Yong Yu, and Yu Zheng.  Hyperst-net: Hypernetworks for spatio-temporal forecasting.  _arXiv preprint arXiv:1809.10889_ , 2018.
  * Pawlowski et al. (2017) Nick Pawlowski, Andrew Brock, Matthew CH Lee, Martin Rajchl, and Ben Glocker.  Implicit weight uncertainty in neural networks.  _arXiv preprint arXiv:1711.01297_ , 2017.
  * Ratzlaff & Fuxin (2019) Neale Ratzlaff and Li Fuxin.  Hypergan: A generative model for diverse, performant neural networks.  _arXiv preprint arXiv:1901.11058_ , 2019.
  * Reddi et al. (2018) Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.  On the convergence of adam and beyond.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=ryQu7f-RZ.
  * Rumelhart et al. (1986) David E Rumelhart, Geoffrey E Hintont, and Ronald J Williams.  Learning representations by back-propagating errors.  _NATURE_ , 323:9, 1986.
  * Saxe et al. (2013) Andrew M Saxe, James L McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  _arXiv preprint arXiv:1312.6120_ , 2013.
  * Serrà et al. (2019) Joan Serrà, Santiago Pascual, and Carlos Segura.  Blow: a single-scale hyperconditioned flow for non-parallel raw-audio voice conversion.  _arXiv preprint arXiv:1906.00794_ , 2019.
  * Shen et al. (2017) Falong Shen, Shuicheng Yan, and Gang Zeng.  Meta networks for neural style transfer.  _arXiv preprint arXiv:1709.04111_ , 2017.
  * Springenberg et al. (2014) Jost Tobias Springenberg, Alexey Dosovitskiy, Thomas Brox, and Martin Riedmiller.  Striving for simplicity: The all convolutional net.  _arXiv preprint arXiv:1412.6806_ , 2014.
  * Stanley et al. (2009) Kenneth O Stanley, David B D’Ambrosio, and Jason Gauci.  A hypercube-based encoding for evolving large-scale neural networks.  _Artificial life_ , 15(2):185–212, 2009.
  * Suarez (2017) Joseph Suarez.  Language modeling with recurrent highway hypernetworks.  In _Advances in neural information processing systems_ , pp. 3267–3276, 2017.
  * Sun et al. (2017) Zhun Sun, Mete Ozay, and Takayuki Okatani.  Hypernetworks with statistical filtering for defending adversarial examples.  _arXiv preprint arXiv:1711.01791_ , 2017.
  * Ukai et al. (2018) Kenya Ukai, Takashi Matsubara, and Kuniaki Uehara.  Hypernetwork-based implicit posterior estimation and model averaging of cnn.  In _Asian Conference on Machine Learning_ , pp. 176–191, 2018.
  * von Oswald et al. (2019) Johannes von Oswald, Christian Henning, João Sacramento, and Benjamin F Grewe.  Continual learning with hypernetworks.  _arXiv preprint arXiv:1906.00695_ , 2019.
  * Wilson et al. (2017) Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht.  The marginal value of adaptive gradient methods in machine learning.  In _Advances in Neural Information Processing Systems_ , pp. 4148–4158, 2017.
  * Wu & He (2018) Yuxin Wu and Kaiming He.  Group normalization.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , pp. 3–19, 2018.
  * Zhang et al. (2018) Chris Zhang, Mengye Ren, and Raquel Urtasun.  Graph hypernetworks for neural architecture search.  _arXiv preprint arXiv:1810.05749_ , 2018.
  * Zhang et al. (2019) Hongyi Zhang, Yann N Dauphin, and Tengyu Ma.  Fixup initialization: Residual learning without normalization.  _arXiv preprint arXiv:1901.09321_ , 2019.

## Appendix

##  Appendix A Re-using Hypernet Weights

###  A.1 For Mainnet Weights of the Same Size

For model compression or weight-sharing purposes, different parts of the
mainnet might be generated by the same hypernet function. This will cause some
assumptions of independence in our analysis to be invalid. Consider the
example of the same hypernet being used to generate multiple different mainnet
weight layers of the same size, i.e.
$H[t]^{i_{t+1}}_{i_{t}k}=H[t+1]^{i_{t+2}}_{i_{t+1}k},\text{d}_{i_{t+1}}=\text{d}_{i_{t+2}}=\text{d}_{i_{t}}$.
Then,
$x[t+1]^{i_{t+1}}=H[t]^{i_{t+1}}_{i_{t}k_{t}}e[t]^{k_{t}}x[t]^{i_{t}}\not\\!\perp\\!\\!\\!\perp
W[t+1]^{i_{t+2}}_{i_{t+1}}=H[t+1]^{i_{t+2}}_{i_{t+1}k}e[t+1]^{k_{t+1}}$.

The relaxation of some of these independence assumptions does not always prove
to be a big problem in practice, because the correlations introduced by
repeated use of $H$ can be minimized with the use of flat distributions like
the uniform distribution. It can even be helpful, since the re-use of the same
hypernet for different layers causes the gradient flowing through the hypernet
output layer to be the sum of the gradients from the weights of these layers:
$\frac{\partial L}{\partial h(e)^{k}}=\sum_{t}\frac{\partial L}{\partial
W[t]^{i_{t+1}}_{i_{t}}}H^{i_{t+1}}_{i_{t}k}$, thus combating the shrinking
effect.

###  A.2 For Mainnet Weights of Different Sizes

Similar reasoning applies if the same hypernet was used to generate
differently sized subsets of weights in the mainnet. However, we encourage
avoiding this kind of hypernet architecture design if not otherwise essential,
since it will complicate the initialization formulae listed in Table 1.

Consider Ha et al. (2016)’s hypernetwork architecture. Their two-layer
hypernet generated weight chunks of size $(K,n,n)$ for a main convolutional
network where $K=16$ was found to be the highest common factor among the size
of mainnet layers, and $n^{2}=9$ was the size of the receptive field. We
simplify the presentation by writing $i$ for $i_{t}$, $j$ for $j_{t}$, $k$ for
$k_{t,m}$, and $l$ for $l_{t,m}$.

$\begin{split}W[t]^{i}_{j}&=\begin{cases}H^{i(\text{mod
}K)}_{k}\alpha[t][j+\lfloor\frac{i}{K}\rfloor\text{d}_{j}]^{k}+\beta^{i(\text{mod
}K)}\qquad\qquad\qquad\qquad\ \ \text{if}\ i\text{ is divisible by }K\\\
\delta_{j(\text{mod }K)j(\text{mod }K)}\bigl{[}H^{j(\text{mod
}K)}_{k}\alpha[t][i+\lfloor\frac{j}{K}\rfloor\text{d}_{i}]^{k}+\beta^{j(\text{mod
}K)}\bigr{]}\quad\text{if}\ j\text{ is divisible by }K\end{cases}\\\
\alpha[t][m_{t}]^{k}&=G[t][m_{t}]^{k}_{l}e[t][m_{t}]^{l}+\gamma[t][m_{t}]^{k}\end{split}$
(6)

Because the output layer $(H,\beta)$ in the hypernet was re-used to generate
mainnet weight matrices of different sizes (i.e. in general, $i_{t}\neq
i_{t+1},j_{t}\neq j_{t+1}$), $G$ effectively becomes the output layer that we
want to be considering for hyperfan-in and hyperfan-out initialization.

Hence, to achieve fan-in in the mainnet
$\text{Var}(W[t]^{i}_{j})=\frac{1}{\text{d}_{j}}$, we have to use fan-in init
for $H$ (i.e. $\text{Var}(H^{i(\text{mod
}K)}_{k})=\frac{1}{\text{d}_{k}}\neq\frac{1}{\text{d}_{j}\text{d}_{k}\text{Var}(e[t][m_{t}]^{l})}$),
and hyperfan-in init for $G$ (i.e.
$\text{Var}(G[t][m_{t}]^{k}_{l})=\frac{1}{\text{d}_{j}\text{d}_{l}\text{Var}(e[t][m_{t}]^{l})}$).

Analogously, to achieve fan-out in the mainnet
$\text{Var}(W[t]^{i}_{j})=\frac{1}{\text{d}_{i}}$, we have to use fan-in init
for $H$ (i.e. $\text{Var}(H^{i(\text{mod
}K)}_{k})=\frac{1}{\text{d}_{k}}\neq\frac{1}{\text{d}_{i}\text{d}_{k}\text{Var}(e[t][m_{t}]^{l})}$),
and hyperfan-out init for $G$ (i.e.
$\text{Var}(G[t][m_{t}]^{k}_{l})=\frac{1}{\text{d}_{i}\text{d}_{l}\text{Var}(e[t][m_{t}]^{l})}$).

##  Appendix B More Experimental Details

###  B.1 Feedforward Networks on MNIST

The networks were trained on MNIST for $30$ epochs with batch size $10$ using
a learning rate of $0.0005$ for the hypernets and $0.01$ for the classical
network. The hypernets had one linear layer with embeddings of size $50$ and
different hidden layers in the mainnet were all generated by the same hypernet
output layer with a different embedding, which was randomly sampled from
$\mathcal{U}(-\sqrt{3},\sqrt{3})$ and fixed. We use the mean cross entropy
loss for training, but the summed cross entropy loss for testing.

We show activation and gradient plots for two cases: (i) the hypernet
generates only the weights of the mainnet, and (ii) the hypernet generates
both the weights and biases of the mainnet. (i) covers Figures 3, 1, 7, 8, 9,
10, 11, 12, 2, 13, 14, 15, and 16. (ii) covers Figures 17, 18, 19, 20, 21, 22,
23, 24, 25, 26, 27, 28, and 29.

The activations and gradients in our plots were calculated by averaging across
a fixed held-out set of $300$ examples drawn randomly from the test set.

In Figures 1, 8, 9, 11, 12, 13, 14, 16, 18, 20, 21, 23, 24, 26, 27, and 29,
the y axis shows the number of activations/gradients, while the x axis shows
the value of the activations/gradients. The value of activations/gradients
from the hypernet output layer correspond to the value of mainnet weights.

In Figures 2, 7, 10, 15, 19, 22, 25, and 28, the y axis shows the mean value
of the activations/gradients, while each increment on the x axis corresponds
to a measurement that was taken every $1000$ training batches, with the bars
denoting one standard deviation away from the mean.

####  B.1.1 Hypernet Generates Only the Mainnet Weights

Figure 7: Evolution of Mainnet Activations during Training on MNIST. Figure 8:
Mainnet Activations at the End of Training on MNIST. Figure 9: Mainnet
Gradients before the Start of Training on MNIST. Figure 10: Evolution of
Mainnet Gradients during Training on MNIST. Figure 11: Mainnet Gradients at
the End of Training on MNIST. Figure 12: Hypernet Output Layer Activations
before the Start of Training on MNIST. Figure 13: Hypernet Output Layer
Activations at the End of Training on MNIST. Figure 14: Hypernet Output Layer
Gradients before the Start of Training on MNIST. Figure 15: Evolution of
Hypernet Output Layer Gradients during Training on MNIST. Figure 16: Hypernet
Output Layer Gradients at the End of Training on MNIST.

####  B.1.2 Hypernet Generates Both Mainnet Weights and Biases

Figure 17: Loss and Test Accuracy Plots on MNIST. Figure 18: Mainnet
Activations before the Start of Training on MNIST. Figure 19: Evolution of
Mainnet Activations during Training on MNIST. Figure 20: Mainnet Activations
at the End of Training on MNIST. Figure 21: Mainnet Gradients before the Start
of Training on MNIST. Figure 22: Evolution of Mainnet Gradients during
Training on MNIST. Figure 23: Mainnet Gradients at the End of Training on
MNIST. Figure 24: Hypernet Output Layer Activations before the Start of
Training on MNIST. Figure 25: Evolution of Hypernet Output Layer Activations
during Training on MNIST. Figure 26: Hypernet Output Layer Activations at the
End of Training on MNIST. Figure 27: Hypernet Output Layer Gradients before
the Start of Training on MNIST. Figure 28: Evolution of Hypernet Output Layer
Gradients during Training on MNIST. Figure 29: Hypernet Output Layer Gradients
at the End of Training on MNIST.

####  B.1.3 Remark on the Combination of Fan-in and Fan-out Init

Glorot & Bengio (2010) proposed to use the harmonic mean of the two different
initialization formulae derived from the forward and backward pass. He et al.
(2015) commented that either version suffices for convergence, and that it
does not really matter given that the difference between the two will be a
depth-independent factor.

We experimented with the harmonic, geometric, and arithmetic means of the two
different formulae in both the classical and the hypernet case. There was no
indication of any significant benefit from taking any of the three different
means in both cases. Thus, we confirm and concur with He et al. (2015)’s
original observation that either the fan-in or the fan-out version suffices.

###  B.2 Continual Learning on Regression Tasks

The mainnet is a feedforward network with two hidden layers ($10$ hidden
units) and the ReLU activation function. The weights and biases of the mainnet
are generated from a hypernet with two hidden layers ($10$ hidden units) and
trainable embeddings of size $2$ sampled from
$\mathcal{U}(-\sqrt{3},\sqrt{3})$. We keep the same continual learning
hyperparameter $\beta_{output}$ value of $0.005$ and pick the best learning
rate for each initialization method from
$\\{10^{-2},10^{-3},10^{-4},10^{-5}\\}$. Notably, Kaiming (fan-in) could only
be trained from learning rate $10^{-5}$, with losses diverging soon after
initialization using the other learning rates. Each task was trained for
$6000$ training iterations using batch size $32$, with Figure 4 plotted from
losses measured at every $100$ iterations.

###  B.3 Convolutional Networks on CIFAR-10

The networks were trained on CIFAR-10 for $500$ epochs starting with an
initial learning rate of $0.0005$ using batch size $100$, and decaying with
$\gamma=0.1$ at epochs $350$ and $450$. The hypernet is composed of two layers
($50$ hidden units) with separate embeddings and separate input layers but
shared output layers. The weight generation happens in blocks of $(96,3,3)$
where $K=96$ is the highest common factor between the different sizes of the
convolutional layers in the mainnet and $n=3$ is the size of the convolutional
filters (see Appendix Section A.2 for a more detailed explanation on the
hypernet architecture). The embeddings are size $50$ and fixed after random
sampling from $\mathcal{U}(-\sqrt{3},\sqrt{3})$. We use the mean cross entropy
loss for training, but the summed cross entropy loss for testing.

###  B.4 Bayesian Neural Network on ImageNet

Ukai et al. (2018) showed that a Bayesian neural network can be developed by
using a hypernetwork to express a prior distribution without substantial
changes to the vanilla hypernetwork setting. Their methods simply require
putting $\mathcal{L}_{2}$-regularization on the model parameters and sampling
from stochastic embeddings. We trained a linear hypernet to generate the
weights of a MobileNet mainnet architecture (excluding the batch normalization
layers), using the block-wise sampling strategy described in Ukai et al.
(2018), with a factor of $0.0005$ for the $\mathcal{L}_{2}$-regularization. We
initialize fixed embeddings of size $32$ sampled from
$\mathcal{U}(-\sqrt{3},\sqrt{3})$, and sample additive stochastic noise coming
from $\mathcal{U}(-{0.1},{0.1})$ at the beginning of every mini-batch
training. The training was done on ImageNet with batch size $256$ and learning
rate $0.1$ for $25$ epochs, or equivalently, $125125$ iterations. The testing
was done with $10$ Monte Carlo samples. We omit the test loss plots due to the
computational expense of doing $10$ forward passes after every mini-batch
instead of every epoch.


(eccv) Package eccv Warning: Package ‘hyperref’ is loaded with option
‘pagebackref’, which is *not* recommended for camera-ready version

11institutetext: School of Computer Science and Technology, Huazhong
University of Science and Technology, Wuhan, China
11email<EMAIL_ADDRESS>22institutetext:
Department of Computing, The Hong Kong Polytechnic University, Hongkong, China
22email<EMAIL_ADDRESS>
# Personalized Federated Domain-Incremental Learning
based on Adaptive Knowledge Matching

Yichen Li 11 0009-0009-8630-2504 Wenchao Xu 22 1111-2222-3333-4444 Haozhao
Wang Haozhao Wang, Yining Qi and Ruixuan Li are corresponding authors.11
0000-0002-7591-5315
Yining Qi 11 0009-0001-3685-5989 Jingcai Guo 22 0000-0002-0449-4525 Ruixuan Li
11 0000-0002-7791-5511

###### Abstract

This paper focuses on Federated Domain-Incremental Learning (FDIL) where each
client continues to learn incremental tasks where their domain shifts from
each other. We propose a novel adaptive knowledge matching-based personalized
FDIL approach (pFedDIL) which allows each client to alternatively utilize
appropriate incremental task learning strategy on the correlation with the
knowledge from previous tasks. More specifically, when a new task arrives,
each client first calculates its local correlations with previous tasks. Then,
the client can choose to adopt a new initial model or a previous model with
similar knowledge to train the new task and simultaneously migrate knowledge
from previous tasks based on these correlations. Furthermore, to identify the
correlations between the new task and previous tasks for each client, we
separately employ an auxiliary classifier to each target classification model
and propose sharing partial parameters between the target classification model
and the auxiliary classifier to condense model parameters. We conduct
extensive experiments on several datasets of which results demonstrate that
pFedDIL outperforms state-of-the-art methods by up to 14.35% in terms of
average accuracy of all tasks.

##  1 Introduction

Federated Learning (FL) has become a basic paradigm that facilitates
collaborative training of deep neural networks (DNNs) from distributed clients
while preserving their data privacy[mcmahan2017communication,
zhang2021survey]. To deploy FL in real-world applications, it is practical
that each client can continuously learn new incremental tasks while retaining
the accuracy for the previous tasks [yoon2021federated]. However, history data
may often be out of reach with privacy constraints like GDPR and many studies
have shown that DNNs will forget the knowledge of previous tasks when learning
new tasks [van2022three, tankard2016gdpr]. Such a phenomenon is known as
’Catastrophic Forgetting’, which is a major challenge of incremental learning
in the centralized setting. Especially, this challenge is even more formidable
to solve in the FL scenario, where the limited previous data are stored in
distributed clients and cannot be accessed by each other [zhang2023target,
bakman2023federated].

Many recent studies have been proposed to solve the catastrophic forgetting
challenge in federated incremental learning. For instance, [yoon2021federated]
aims to learn a personalized model for each client by decomposing the model
parameters into shared parameters and adaptive parameters and transferring
common knowledge between similar tasks across clients. FCIL
[dong2022federated] addresses the federated class-incremental learning and
trains a global model by computing additional class-imbalance losses. Authors
in [ijcai2022p0303] utilize extra data at both the server and client sides to
conduct knowledge distillation to alleviate forgetting. [qi2023better]
optimize the generative network and reconstruct previous samples for replay.
Recently, Re-Fed is proposed in [Li_2024_CVPR] to prevent catastrophic
forgetting with synergistic replay by caching important samples which
contribute to both the local and global understanding.

While these approaches have achieved great success, they are generally
designed for the scenario where the new tasks are incremental classes, i.e.,
Federated Class-Incremental Learning (FCIL). These methods may lose their
effectiveness when the new tasks have the same classes but their domain shift
with different feature distributions, i.e., Federated Domain-Incremental
Learning (FDIL). To elaborate further, the FCIL[dong2022federated] method
relies on changes in the new data category label in incremental tasks to
alleviate forgetting of previous tasks; FOT[bakman2023federated] projects
different incremental tasks into different orthogonal subspaces, which means
feature isolation for different tasks, but if suffers from the intersected
feature space boundaries of incremental tasks. The primary challenge in
implementing FCIL methods in the context of FDIL is that FCIL methods
typically depend on the connection between previous and new classes, which is
not applicable in FDIL scenarios where the classes remain consistent.

In this paper, we seek to tackle the challenge of catastrophic forgetting in
the FDIL scenario. A simple way to address this problem is to train different
models for each incremental task. However, this requires extensive storage and
time-consuming re-training. Additionally, all these models need to be saved
for future use which is laborious because we need to manually select the
appropriate mode for each inference data. Instead of training a separate model
for each different task, a better solution is to ideally have only a single
model that can work in all incremental tasks. However, naively training a
model for all tasks can lead to terrible catastrophic forgetting when the
unique knowledge in the new task overlaps the ones in the previous tasks.

Unlike in the FCIL scenario, where incremental tasks involve data from
different categories, In the FDIL scenario, the incremental tasks involve data
from different domains, i.e., handwritten digit data and cartoon digit data.
Compared to the data from different categories in FCIL. i.e., cat data and dog
data, there is more shared feature information between the different
incremental tasks, which can facilitate the overall FDIL task[liu2019compact].
The key to efficiently learning new tasks while memorizing previous tasks in
the FDIL scenario is to identify the correlation between incremental tasks and
then share the useful common knowledge in different tasks. Like
[yoon2021federated] in FCIL, we consider each client maintains their
personalized models while the number of personalized models is under a given
threshold due to the limited capacity of each client. Given the correlation
between tasks, each client can maintain several personalized models and share
the same model between similar tasks. In this way, the client can choose a
model with tasks most similar to the new task, enabling the acquisition of new
task knowledge while retaining existing task knowledge. This dual benefit
permits the new task to also capitalize on previous knowledge, ultimately
enhancing performance. Nonetheless, considering the tasks come incrementally,
identifying the inter-task correlations and subsequently training the selected
model is challenging.

To tackle these limitations, we in this paper propose a novel Personalized
Federated Domain-Incremental Learning (pFedDIL) approach which adaptively
matches and migrates knowledge among different incremental tasks. More
specifically, we propose that each client employs a small auxiliary classifier
for each personalized model which is trained to distinguish itself from other
tasks. Using the auxiliary classifiers of previous tasks, each client can
adaptively identify the correlation of the new task to all previous tasks.
Based on these correlations, each client can match the new task with a
suitable model, i.e., a new initialized model or a previous model but with
similar knowledge. To maximize the training efficiency of the matched model
over the new task, we further allow each client to migrate knowledge from
other not matched but similar personalized models in a weighted manner, where
the weights are the correlations obtained previously. Furthermore, to reduce
the size of auxiliary classifiers while enhancing their performance, we
propose sharing partial parameters between the target personalized model and
the auxiliary classifier. Finally, considering the presence of multiple
personalized local models within each client that could lead to additional
inference costs, we apply a weighted ensemble distillation over these models
where the weights are also obtained from the output of the auxiliary
classifiers.

Through extensive experiments on several datasets and different settings, we
show that the proposed pFedDIL significantly improves the model accuracy as
compared to state-of-the-art approaches. The major contributions of this paper
are summarized as follows:

  * •

We are the first to formulate the problem of federated domain-incremental
learning. To address this issue, we propose a novel method named pFedDIL which
enables each client to match and migrate knowledge from previous tasks to
alleviate catastrophic forgetting.

  * •

To adaptively employ incremental task learning strategy and migrate knowledge,
we define the knowledge matching intensity by the auxiliary classifier to
discern the correlation between the new task and previous tasks.

  * •

To facilitate the adaptive selection of inference model among a set of
personalized local models, we coordinate all personalized local models to
predict weighted by normalized output by auxiliary classifiers. Besides, we
further propose sharing partial parameters between the auxiliary classifier
and the target classification model to condense model parameters.

  * •

We carry out extensive experiments on various datasets and different settings.
Experimental results illustrate that our proposed model outperforms the state-
of-the-art methods by up to 14.35% in terms of average accuracy on all tasks.

##  2 Related Work

Federated Learning. Federated learning (FL) is a distributed learning approach
that trains a global model by merging multiple locally trained models on
private datasets[li2020review]. A significant challenge in FL is data
heterogeneity, where data distribution among clients is non-identically and
independently distributed (Non-IID), negatively impacting model
performance[Jeong2018CommunicationEfficientOM, wang2023fedcda,
Shen2020FederatedML]. FedAvg is a notable FL method that refines the global
model by incorporating parameters trained on local private
data[mcmahan2017communication]. [li2020federated] propose a proximal term to
align closely with the global model. The federated distillation
method[wang2023dafkd] suggests distilling multiple local models into a global
one by only aggregating soft output predictions. However, existing FL
approaches typically assume that the data in each client is static, which
cannot account for incremental data with domain shift.

Incremental Learning. Incremental learning (IL) is a learning paradigm that
empowers a model to learn sequentially from a series of tasks while preserving
the knowledge acquired from previous tasks[wang2024comprehensive,
zhao2023does]. IL can be segmented into three primary variants: task-
incremental learning[Maltoni2018ContinuousLI], class-incremental
learning[rebuffi2017icarl] and domain-incremental
learning[mirza2022efficient]. Existing strategies can further be grouped into
three categories: replay-based methods[rolnick2019experience], regularization-
based methods[Jung2020ContinualLW] and parameter isolation
methods[chen2020mitigating]. Most IL methods are designed for task-IL or
class-IL and are typically server-based. This work focuses on the federated
domain-incremental learning scenario and proposes the pFedDIL technique.

Figure 1: Illustration of the pFedDIL framework. In the training stage, each
client randomly receives one domain dataset without replacement to form a
task, where different colors represent different domains. Clients can
adaptively employ an appropriate incremental task learning strategy and
migrate the knowledge from previous tasks with the knowledge matching
intensity.

Federated Incremental Learning. Federated Incremental Learning (FIL), an
approach that targets the dynamic adaptation of streaming tasks within each
client’s environment, emphasizes the seamless integration of new data into the
global model while preserving the acquired knowledge from previous data
streams. Despite its crucial role, FIL has garnered attention relatively
recently, with [yoon2021federated] pioneering the field by focusing on Task-
Incremental Learning (Task-IL) in FL. In parallel, studies such as
[yu2024personalized] takes the multi-granularity expression of knowledge into
account, promoting the spatial-temporal integration of knowledge in FCIL.
Conversely, [Dong_2023_ICCV, 10323204] simplify the problem’s complexity by
presupposing clients possess sufficient storage capacity to retain historical
examples and engage in data sharing, a scenario that deviates from the
conventional federated learning paradigm. Beyond image classification,
investigations like [Jiang_2021, 10204005] extend the horizons of FIL into
diverse domains. Our research focuses on the domain shift between different
tasks and the transfer of the common knowledge.

##  3 Methodology

We first formulate our problem definition and then present our pFedDIL
approach. Figure 1 gives an overview of our method.

###  3.1 Problem Definition

In incremental learning, each model learns from a sequence of streaming tasks
{$\mathcal{T}_{(1)},\mathcal{T}_{(2)}\\\ ,\ldots\mathcal{T}_{(n)}$} where
$\mathcal{T}_{(t)}$ denotes the dataset of the $t$-th task;
$\mathcal{T}_{(t)}=\\{x_{i}^{t},y_{i}^{t}\\}_{i=1}^{N^{t}}$, which has $N^{t}$
pairs of sample data $x_{i}^{t}\in X^{t}$ and corresponding label
$y_{i}^{t}\in Y^{t}$. $X^{t}$ and $Y^{t}$ represent the domain space and label
space for the $t$-th task. In this paper, we focus on the variant of domain
incremental learning, where all tasks share the same numbers of classes i.e,
$\mathcal{Y}^{1}=\mathcal{Y}^{t},\forall t\in[1,n]$. As the sequence of tasks
arrives, the client needs to learn the new task while their domain and data
distribution changes, i.e, $\mathcal{X}^{1}\neq\mathcal{X}^{t},\forall
t\in[1,n]$.

We consider the federated domain-incremental learning scenario. We aim to
train a global model for $K$ total clients in the federated domain-incremental
learning setting. We consider the local client $k$ can only access the local
private streaming tasks
{$\mathcal{T}^{k}_{(1)},\mathcal{T}^{k}_{(2)},\dots,\mathcal{T}^{k}_{(n)}$}.
Now our goal is to train a set of personalized local models
$\widetilde{w}^{k}=[w^{k}_{1},w^{k}_{2},\ldots,w^{k}_{d}]$ for each client $k$
over all existing $t$ tasks
$\mathcal{T}=\sum_{n=1}^{t}\sum_{k=1}^{K}\mathcal{T}^{k}_{(n)}$. Here each
personalized local model corresponds to at least one task, and possibly
multiple tasks (i.e. client $k$ may train the task $\mathcal{T}^{k}_{(1)}$ and
$\mathcal{T}^{k}_{(2)}$ with the model $w^{k}_{1}$). Then, we can formulate
the goal as :

$\displaystyle\mathop{\arg\min}_{\widetilde{w}}\mathcal{L}(\widetilde{w})=\sum_{n=1}^{t}\sum_{k=1}^{K}\frac{1}{|\mathcal{T}|}\mathcal{L}_{(n)}^{k}(\widetilde{w}).$
(1)

where
$\mathcal{L}_{(n)}^{k}(\widetilde{w})=\mathbb{E}_{(x,y)\sim\mathcal{T}^{k}_{(n)}}[l(\widetilde{w};(x,y))]$
is the empirical loss of client $k$ on $n$-th task with the personalized local
model.

###  3.2 The Proposed pFedDIL Approach

The key idea of the pFedDIL is to identify the correlation between the new and
previous tasks such that clients can adaptively migrate knowledge from
previous tasks to improve performance. Specifically, in each communication
round, the client in pFedDIL trains the local model with its private sequence
of tasks and the server aggregates local models from all participating
clients. Then, when it comes to new tasks, each client first employs the
existing personalized models’ auxiliary classifiers to calculate the
correlations denoted as knowledge matching intensity and applies the
appropriate incremental task learning strategy: training the new task with a
previous personalized local model which has the most similar knowledge or
training with a new initial model. Simultaneously, each client can migrate
knowledge from other personalized local models with the knowledge matching
intensity to to assist in training. During the inference stage, the final
inference result is aggregated by all personalized local models weighted by
the corresponding auxiliary classifiers. Moreover, we propose sharing partial
parameters between the auxiliary classifier and the target personalized model
to condense model parameters. The workflow of the pFedDIL is shown in
Algorithm 1.

Knowledge Matching Intensity. To assess the similarity between different
incremental tasks, where a higher similarity indicates a greater amount of
common knowledge, we propose the knowledge matching intensity. When it comes
to the new incremental task: $(t+1)$-th task, each client first needs to
compute the knowledge matching intensity between the $(t+1)$-th task and
previous $t$ tasks. To achieve the adaptive knowledge matching intensity
estimation for each client, we define the auxiliary classifier, a binary
classification deep neural network $\theta$ for the corresponding local task
$T$, which is trained by the discrimination of the sample belonging to the
task $T$. As a new task emerges, the auxiliary classifier can output the
similarity between different tasks based on the extracted feature
representation. Supposed that the client $k$ has converged on $t$ tasks within
$d$ local existing task models when it comes to the incremental $(t+1)$-th
task, the client needs to calculate the knowledge matching intensity
$\widetilde{\rho}_{k}=[\rho^{k}_{1},\rho^{k}_{2},\ldots,\rho^{k}_{d}]$ with
each auxiliary classifier
$\widetilde{\theta}^{k}=[\theta^{k}_{1},\theta^{k}_{2},\ldots,\theta^{k}_{d}]$:

$\displaystyle\widetilde{\rho}^{k}=\frac{1}{N^{t+1}}\sum_{i=1}^{N^{t+1}}{f(x^{t+1}_{i};\widetilde{\theta}^{k})}.$
(2)

where the $f(\cdot)$ denotes the output of the auxiliary classifier,
$\widetilde{\theta}^{k}$ denotes a set of personalized local auxiliary
classifiers in client $k$. Each auxiliary classifier is trained by minimizing
the total empirical loss over the corresponding task’s local dataset
$\mathcal{T}$, which can be formulated as:

$\displaystyle\mathcal{L}^{k}(\theta)=\frac{1}{\mathcal{T}}\sum_{i=1}^{\mathcal{T}}\mathcal{L}_{CE}(\theta;x_{i},y_{i}).$
(3)

where $\mathcal{L}^{k}(\theta)$ is the local loss in the client $k$ and
$\mathcal{L}_{CE}$ is the cross-entropy loss function that measures the
difference between the prediction and the ground truth labels.

Input : $T$: the communication round; $K$: client number;

$C$: the fraction of active client in each round;

$\\{\mathcal{T}_{(t)}\\}_{t=1}^{n}$: the distributed dataset with $n$ tasks;

$w$: the parameter of the target classification model;

$\theta$: the parameter of the auxiliary classifier.

1

2Initialize the parameter $w$ and $\theta$;

3for  _$t=1$ to $T$_ do

4 $S_{t}\leftarrow$ (random set of $[C\cdot K]$ clients); for  _each selected
client $k\in S_{t}$ in parallel_ do

5 receives $w_{t}$ from the server;

6 calculate knowledge matching intensity $\widetilde{\rho}^{k}$ with (2);

7 select learning strategy $w^{k}_{t+1}$ with (4);

8 set local models $w_{t+1}^{k}$ and $\theta^{k}_{t+1}$;

9 for  _$e=1$ to $E$_ do

10 update $\theta^{k}_{t+1}$ with (3);

11 update $w_{t+1}^{k}$ through adaptive knowledge migration with
(LABEL:localloss);

12 end for

13 pushes $w_{t+1}^{k}$ to the server.

14

15 end for

16 $w_{t}\leftarrow\text{ServerAggregation}(\\{w^{k}\\}_{k\in S_{t}})$

17 end for

Algorithm 1 pFedDIL

Incremental Task Learning Strategy. After acquiring the knowledge matching
intensity $\widetilde{\rho}_{k}$ with equation (2), the client can select the
appropriate incremental task learning strategy. An intuition is that the
client will prefer to train with previous personalized local models when the
knowledge matching intensity exceeds the $\lambda$, which means the high-level
similarity between the new task and the previous tasks so that the model can
converge on the new task with slight adjustments. Besides, the client will
sacrifice some storage by training the new task with the new initial model
parameters to ensure accuracy for both the new task and previous tasks. Here
we define the incremental task learning strategy as follows:

$w^{k}_{t+1}=OPT\begin{cases}\hat{w}&\text{ $if\
max(\widetilde{\rho}_{k})<\lambda$ }\\\ \widetilde{w}^{k}[m]&\text{ $else\
max(\widetilde{\rho}_{k})=\rho_{k}^{m}\geq\lambda$ }.\end{cases}$ (4)

where $\lambda$ is a hyper-parameter,
$\widetilde{w}^{k}=[w^{k}_{1},w^{k}_{2},\ldots,w^{k}_{d}]$, $\hat{w}$ denotes
new initial model parameters, and $OPT$ means the alternative of existed
personalized local models and a new initial model.

Knowledge Migration. Then, during the training of $(t+1)$-th task, each client
respectively migrates the knowledge from other existing personalized local
models to assist the training with the calculated knowledge matching intensity
$\widetilde{\rho}_{k}$, minimize the following objective: