# Layerwise Bregman Representation Learning
with Applications to Knowledge Distillation

Ehsan Amid
<EMAIL_ADDRESS>Rohan Anil
<EMAIL_ADDRESS>Christopher Fifty
<EMAIL_ADDRESS>Manfred K. Warmuth
<EMAIL_ADDRESS>
(Google Research, Brain Team)

###### Abstract

In this work, we propose a novel approach for layerwise representation
learning of a trained neural network. In particular, we form a Bregman
divergence based on the layer’s transfer function and construct an extension
of the original Bregman PCA formulation by incorporating a mean vector and
normalizing the principal directions with respect to the geometry of the local
convex function around the mean. This generalization allows exporting the
learned representation as a fixed layer with a non-linearity. As an
application to knowledge distillation, we cast the learning problem for the
student network as predicting the compression coefficients of the teacher’s
representations, which are passed as the input to the imported layer. Our
empirical findings indicate that our approach is substantially more effective
for transferring information between networks than typical teacher-student
training using the teacher’s penultimate layer representations and soft
labels.

##  1 Introduction

Principal Component Analysis (PCA) [26, 30] is perhaps one of the most
commonly used techniques for data compression, dimensionality reduction, and
representation learning. In the simplest form, the PCA problem is defined as
minimizing the _compression loss_ of representing a set of points as linear
combinations of a set of orthonormal _principal directions_. More concretely,
given $\mathcal{X}=\\{\mathbf{x}_{i}\in\mathbb{R}^{d}\\}$ and $k\leq d$, the
PCA problem can be formulated as finding the _mean vector_
$\bm{m}\in\mathbb{R}^{d}$ and principal directions
$\bm{V}\in\mathbb{R}^{d\times k}$ where $\bm{V}^{\top}\bm{V}=\bm{I}_{k}$ such
that the _compression loss_,

$\bm{m},\bm{V},\\{\bm{c}_{i}\\}=\\!\\!\\!\\!\\!\\!\\!\\!\\!\operatorname*{arg\,min}_{\Big{\\{}\begin{subarray}{c}\tilde{\bm{m}}\in\mathbb{R}^{d},\widetilde{\bm{V}}\in\text{St}_{d,k},\\\
\\{\tilde{\bm{c}}_{i}\in\mathbb{R}^{k}\\}\end{subarray}\Big{\\}}}\\!\\!\sum_{i}\|\mathbf{x}_{i}-(\tilde{\bm{m}}+\widetilde{\bm{V}}\tilde{\bm{c}}_{i})\|^{2}$
(1)

is minimized. Here, $\text{St}_{d,k}=\\{\bm{U}\in\mathbb{R}^{d\times
k}:\,\bm{U}^{\top}\bm{U}=\bm{I}_{k}\\}$ denotes the Stiefel manifold of
$k$-frames in $\mathbb{R}^{d}$ [14] and $\bm{c}_{i}\in\mathbb{R}^{k}$
corresponds to the _compression coefficients_ of $\mathbf{x}_{i}$. The problem
in Eq. (1) can be solved effectively in two steps. First, we note that
$\bm{m}$ can be viewed as a constant shared representation for all points in
$\mathcal{X}$ (for which the code length $k=0$) that minimizes the total
compression loss. With this interpretation, the mean vector can be written as
the minimizer of

$\bm{m}=\operatorname*{arg\,min}_{\tilde{\bm{m}}\in\mathbb{R}^{d}}\sum_{i}\|\mathbf{x}_{i}-\tilde{\bm{m}}\|^{2}\,,$
(2)

for which, the solution corresponds to the geometric mean
$\bm{m}=\frac{1}{|\mathcal{X}|}\sum_{i}\mathbf{x}_{i}$. By fixing $\bm{m}$,
the solution for $\bm{V}$ and $\\{\bm{c}_{i}\\}$ can be obtained by enforcing
the orthonormality constraints using a set of Lagrange multipliers and setting
the derivatives to zero. The solution to $\bm{V}$ amounts to the the top-$k$
eigenvectors of the covariance matrix
$\nicefrac{{1}}{{|\mathcal{X}|}}\sum_{i}(\mathbf{x}_{i}-\bm{m})(\mathbf{x}_{i}-\bm{m})^{\top}$
and $\bm{c}_{i}=\bm{V}^{\top}(\mathbf{x}_{i}-\bm{m})$ corresponds to the
projection of the centered point onto the column space of $\bm{V}$. Note that
the online variants of PCA, such as Oja’s algorithm [24] alternatively apply a
gradient step on $\bm{V}$ and project the update onto $\text{St}_{d,k}$ by an
application of QR decomposition [13].

Knowledge distillation refers to a set of techniques used for transferring
information from typically a larger trained model, called the _teacher_ , to a
smaller model, called the _student_ [16, 4, 32]. The goal of distillation is
to improve the performance of the student model by augmenting the knowledge
learned by the larger model with the raw information provided by the set of
train examples. Since its introduction, various approaches have applied
knowledge distillation to obtain improved results for language modeling [29],
image classification [6], and robustness against adversarial attacks [25].

The teacher’s knowledge is typically encapsulated in the form of (expanded)
soft labels, which are usually smoothened further by incorporating a
temperature parameter at the output layer of the teacher model [16, 23]. Other
approaches consider matching the teacher’s representations, typically in the
penultimate layer, for a given input by the student [27]. In this paper, we
explore the idea of directly transferring information from a teacher to a
student in the form of  _learned (fixed) principal directions_ in arbitrary
layers of the teacher model. Our focus for representation learning will be on
a generalized form of the PCA method based on the broader class of Bregman
divergences. The Bregman divergence [8] induced by the strictly-convex and
differentiable convex function $G:\mathbb{R}^{d}\rightarrow\mathbb{R}$ between
$\bm{u},\bm{v}\in\operatorname{dom}G$ is defined as

$D_{G}(\bm{u},\bm{v})=G(\bm{u})-G(\bm{v})-g(\bm{v})\cdot(\bm{u}-\bm{v})\,,$
(3)

where $g=\nabla G$ is the gradient function (commonly known as the _link_
function). Bregman divergence is always non-negative and zero iff the
arguments are equal. This broad class of divergences includes many well-known
cases, such as the squared Euclidean and KL divergences as special cases. In
addition, a Bregman divergence is not necessarily symmetric but satisfies a
duality property in terms of the Bregman divergence between the dual
variables. Let
$G^{*}(\bm{u}^{*})=\sup_{\bm{z}}\bm{z}\cdot\bm{u}^{*}-G(\bm{z})$ be the
Legendre dual [17] of $G$. Then, we can write
$D_{G}(\bm{u},\bm{v})=D_{G^{*}\\!}(\bm{v}^{*},\bm{u}^{*})$ where
$\bm{u}^{*}=g(\bm{u})$ and $\bm{v}^{*}=g(\bm{v})$ are the pair of dual
variables. When $G$ is strictly convex and differentiable, we have
$g^{*}=g^{-1}$ and $\bm{u}=g^{*}(\bm{u}^{*})$ and $\bm{v}=g^{*}(\bm{v}^{*})$.
Also, the derivative of a Bregman divergence with respect to the first
argument takes the following simple form

$\nabla_{\bm{u}}D_{G}(\bm{u},\bm{v})=g(\bm{u})-g(\bm{v})\,.$ (4)

Figure 1: (a) A dual view of Bregman divergences as area under the curves of
the transfer function $f$ and its inverse $f^{*}=f^{-1}$ [15, 1]. (b) An
illustration of knowledge distillation using Bregman representation learning.
Given an input example $\mathbf{x}_{i}$, the output representation
$\bm{y}_{i}^{\ell}$ of the teacher network at layer $\ell\in L$ is
approximated as $\hat{\bm{y}}_{i}^{\ell}=f^{\ell}(\bm{V}\bm{c}_{i}+\bm{m})$
where the principal directions $\bm{V}$ and the mean vector $\bm{m}$ are
learned using Bregman PCA on the full train set. The compression coefficients
$\bm{c}_{i}$ are predicted by the student network, and the approximate
representation $\hat{\bm{y}}_{i}^{\ell}$ is passed instead of
$\bm{y}_{i}^{\ell}$ through the rest of the network from layer $\ell+1$ to
$L$. The teacher component inside the dashed bounding box can be discarded
once the student network is trained.

The loss construction in [1] provides a natural way of generating layerwise
Bregman divergences for deep neural networks as line integrals of the strictly
monotonic transfer functions. A dual view of such Bregman divergences in one
dimension is illustrated in Figure 1(a) as areas under the curves of the
strictly monotonic transfer function $f$ and its inverse. We will utilize such
Bregman divergences for layerwise representation learning via an extension of
the Bregman PCA algorithm. Note that the generalization of PCA in Eq. (1) to
Bregman divergences was done in [10] and extended in several previous works
[28, 9]. However, our generalization differs largely from the previous
approaches in terms of the following:

  * •

We extend the formulation to include a generalized mean vector to handle the
non-centered data. To the best of our knowledge, our work is the first
formulation to tackle this case.

  * •

We improve the orthonormality constraint of the Euclidean geometry to
conjugacy in terms of the Riemannian metric [22] induced by the Bregman
divergence. Our extended formulation of mean and the conjugacy constraint
recovers the vanilla PCA problem in Eq. (1) in the Euclidean case.

  * •

We introduce a simple variant of QR decomposition to enforce the generalized
conjugacy constraint efficiently. We also provide a simple trick to handle the
case of constrained conjugate directions for the softmax link function.

  * •

We provide an application of our construction for layerwise representation
learning of deep neural networks with strictly monotonic transfer functions.
Our work builds on the layerwise loss construction approach proposed in [1].
We show that even low-rank approximations of the input examples maintain the
representativeness of the original network.

  * •

Finally, we propose a new approach for knowledge distillation by importing the
fixed principal directions learned from the teacher model into the student
model.

Our extension of the Bregman PCA formulation allows viewing the learned
representation as a fixed layer. Thus, the distillation problem reduces to
learning the corresponding compression coefficients of a given example by the
student, which is fed as input to this fixed layer. Our knowledge distillation
approach using Bregman representation learning is summarized in Figure 1(b).

###  1.1 Related Work

Generalized PCA: Collins et al. [10] propose a generalization of the original
PCA problem, defined for Gaussian models, to general exponential family
distributions. Several extensions of this work focus on applying the framework
to Poisson distributions [9] and compositional data [5]. In [20], the authors
apply logistic PCA for binary data to dimensionality reduction. Landgraf and
Lee [21] propose a majorization-minimization approach for solving the
generalized PCA problem. Despite the success of earlier generalized PCA
methods, a direct generalization of the original PCA formulation in Eq. (1)
with a mean vector, as well as a normalization that considers the local
geometry of the space around the mean, is lacking.

Knowledge Distillation: There have been several lines of work that go beyond
extracting knowledge from the model’s predictive distribution. These include
the work by [27] where they add a regression loss between the representation
of the teacher and a thinner (in the number of units) student representation.
In their work, they learn a projection matrix to match the entire
representation of the teacher rather than the relevant ones that are useful
for the task. They also require the student to be thinner and deeper, which
can result in difficulty in training. This line of work has been improved by
[33] where the students are less deep and make use of attention maps from the
teacher model from a modern convolutional network which performs well on the
problem. Additionally, Czarnecki et al. [11] propose incorporating the
derivative of the teacher’s prediction w.r.t to the input to be matched by the
student, and this additional information has been shown to improve the
distillation procedure. More recently, Tian et al. [31] propose adding
contrastive losses to transfer representation between two networks and show
improvement over vanilla distillation. However, their procedure requires extra
computation to push apart the teacher’s representation of randomly sampled
inputs from the student’s representation of actual inputs. Finally, Müller et
al. [23] tackle the task of extracting more information from the teacher model
when the number of classes is small. They propose creating sub-classes
corresponding to every class, followed by soft-label training for
distillation.

##  2 Extended Bregman PCA

Given a strictly convex and differentiable function
$F:\mathbb{R}^{d}\rightarrow\mathbb{R}$ with link function $f=\nabla F$, we
cast the generalized Bregman PCA problem as approximating
$\mathbf{x}_{i}\in|\mathcal{X}|$ as a linear combination of a set of
orthonormal principal directions in the dual space. This formulation reduces
to Eq. (1) for the choice of $f=\operatorname{id}_{d}$, which corresponds to
the squared Euclidean divergence. However, before defining the objective
function formally, we consider the problem of finding a generalized mean in
the dual space as follows. Let
$\mathcal{X}=\\{\mathbf{x}_{i}\in\operatorname{dom}F^{*}\\}$ be a set of given
data points. We define the generalized mean vector $\bm{m}$ as the minimizer
of the following objective,

$\bm{m}=\operatorname*{arg\,min}_{\tilde{\bm{m}}\in\operatorname{dom}F}\sum_{i}D_{F^{*}\\!}(\mathbf{x}_{i},f(\tilde{\bm{m}}))\,.$
(5)

Note that the above Eq. (5) is a direct generalization of Eq. (2) in terms of
finding a shared constant representation for all points in $\mathcal{X}$ that
minimizes a notion of _Bregman compression loss_. The following proposition
states the solution of the generalized mean in a closed form.

######  Proposition 1.

The generalized dual mean in Eq. (5) can be written as

$\bm{m}=f^{*}\big{(}\frac{1}{|\mathcal{X}|}\sum_{i}\mathbf{x}_{i}\big{)}\,.$
(6)

Thus, the dual mean simply corresponds to the dual of the arithmetic mean of
the data points. When $f=\operatorname{id}_{d}$, the dual mean reduces to the
arithmetic mean.

Given the definition of the dual mean in Eq. (6), we now extend the vanilla
PCA formulation in Eq. (1) to the class of Bregman divergences. First, we note
that the geometry of the space of principal directions is altered when
switching from a squared loss to a more general Bregman divergence.
Specifically, given the convex function $G$ of a Bregman divergence, the
inner-product is locally governed by the Riemannian metric [22],

$D_{G}(\bm{u}+\delta\bm{u},\bm{u})=D_{G}(\bm{u},\bm{u}+\delta\bm{u})\approx\nicefrac{{1}}{{2}}\,\delta\bm{u}^{\top}\,\bm{H}_{G}(\bm{u})\delta\bm{u}\,,$

where $\delta\bm{u}$ is a small perturbation and $\bm{H}_{G}=\nabla^{2}G$ is
the _Hessian_ of $G$. Thus, the definition of orthonormality needs to conform
with the new geometry imposed by the Bregman divergence. In the following, we
extend the definition of a Stiefel manifold to include a Riemannian metric.
Recall that for a strictly convex function $G$, we have
$\bm{H}_{G}\in\mathbb{S}_{+}^{n}$ where $\mathbb{S}_{+}^{n}$ denotes the set
of $n\times n$ symmetric positive-definite matrices.

######  Definition 1.

The generalized Stiefel manifold of $k$-frames in $\mathbb{R}^{d}$ with
respect to the Riemannian metric
$\bm{M}=\bm{M}(\bm{v})\in\mathbb{S}^{d}_{+},\bm{v}\in\mathbb{R}^{d}$ is
defined as

$\text{\emph{St}}^{\text{\scalebox{0.8}{$(\bm{M})$}}}_{d,k}=\\{\bm{U}\in\mathbb{R}^{d\times
k}:\,\bm{U}^{\top}\bm{M}\,\bm{U}=\bm{I}_{k}\\}\,.$

Note that for the Euclidean geometry, $\bm{M}(\bm{v})=\bm{I}_{d}$ for all
$\bm{v}\in\mathbb{R}^{d}$. Thus, the definition of
$\text{St}^{\text{\scalebox{0.8}{$(\bm{M})$}}}_{d,k}$ recovers
$\text{St}^{\text{\scalebox{0.8}{$(\bm{I}_{d})$}}}_{d,k}=\text{St}_{d,k}$ and
we arrive at the orthonormality in the sense of Euclidean geometry.

We now formulate the generalized Bregman PCA problem following the local
geometry of the strictly convex function at $\bm{m}$. Let
$\mathcal{X}=\\{\mathbf{x}_{i}\in\operatorname{dom}F^{*}\\}$ be a set of given
points. We define the generalized Bregman PCA as finding a linear combination
of a set of generalized principal directions that minimizes the Bregman
compression loss:

$\bm{V},\\{\bm{c}_{i}\\}=\\!\\!\\!\\!\operatorname*{arg\,min}_{\Big{\\{}\begin{subarray}{c}\widetilde{\bm{V}}\in\text{\scalebox{1.2}{$\text{St}$}}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k},\\\
\\{\tilde{\bm{c}}_{i}\in\mathbb{R}^{k}\\}\end{subarray}\Big{\\}}}\\!\\!\sum_{i}D_{F^{*}\\!}(\mathbf{x}_{i},f(\bm{m}+\widetilde{\bm{V}}\tilde{\bm{c}}_{i}))\,,$
(7)

where $\bm{H}_{F}(\bm{m})=\nabla^{2}F(\bm{m})$. Note that the constraint
$\bm{V}\in\text{St}_{d,k}$ in Eq. (1) is now replaced with
$\bm{V}\in\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k}$,
i.e., using the Riemannian metric induced by the Hessian of the convex
function $F$ evaluated at the dual mean $\bm{m}$.

###  2.1 Optimization

The generalized Bregman PCA objective in Eq. (7) does not yield a closed-form
solution in terms of $\bm{V}$ and $\\{\bm{c}_{i}\\}$. However, the problem can
be solved iteratively by applications of gradient descent steps on
$\\{\bm{c}_{i}\\}$ and $\bm{V}$. Let $\hat{\mathbf{x}}_{i}\coloneqq
f(\bm{m}+\bm{V}\bm{c}_{i})$ denote the approximation of $\mathbf{x}_{i}$. For
the compression coefficients $\\{\bm{c}_{i}\\}$, we apply

$\bm{c}_{i}^{\text{\tiny
new}}=\bm{c}_{i}-\eta_{a}\bm{V}^{\top}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})\,,$
(8)

where $\eta_{a}>0$ denotes the learning rate. Updating $\bm{V}$ involves two
stages: gradient updates followed by a projection onto
$\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k}$. We apply
gradient decent updates,

$\bm{V}^{\text{\tiny
new}}=\bm{V}-\eta_{\scalebox{0.5}{$V$}}\sum_{i}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})\bm{c}_{i}^{\top}\,.$
(9)

where $\eta_{\scalebox{0.5}{$V$}}>0$ denotes the learning rate. The final step
involves projecting $\bm{V}$ onto
$\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k}$. Notice
that this projection needs to be applied only once at the end of optimization
since both $\bm{V}$ and $\\{\bm{c}_{i}\\}$ are trained using gradient descent
(any intermediate factor can be absorbed into the gradients). For the vanilla
PCA where $\bm{H}_{F}(\bm{m})=\bm{I}_{d}$, this can be applied easily by an
application of QR decomposition. We provide a simple modification of the
standard QR decomposition algorithm that achieves this for any
$\bm{H}_{F}(\bm{m})\in\mathbb{S}_{+}^{n}$, with almost no additional overhead
in practice for our application.

###  2.2 Generalized QR Decomposition

A QR decomposition is a factorization of a matrix
$\bm{A}\in\mathbb{R}^{m\times n}$ where $n\leq m$ into a product
$\bm{A}=\bm{Q}\bm{R}$ where $\bm{Q}\in\text{St}_{\\{m,n\\}}$ and
$\bm{R}\in\mathbb{R}^{n\times n}$ is an upper-triangular matrix. The first
factor $\bm{Q}$ can be viewed as an orthonormalization of columns of $\bm{A}$,
similar to the result of a Gram-Schmidt procedure. However, QR decomposition
provides a more numerically stable procedure in general. The method of
Householder reflections [13] is the most common algorithm for QR
decomposition.

The following theorem provides a procedure that extends the standard QR
decomposition to produce conjugate factors
$\bm{Q}^{\top}\bm{M}\bm{Q}=\bm{I}_{n}$ for a given
$\bm{M}\in\mathbb{S}_{+}^{n}$.

######  Theorem 1.

Let QR denote the procedure that returns the QR factors. Given
$\bm{M}\in\mathbb{S}_{+}^{n}$ and $\bm{A}\in\mathbb{R}^{m\times n}$, let
$\widetilde{\bm{Q}},\bm{R}=\emph{\text{{QR}}}(\sqrt{\bm{M}}\bm{A})$. Then, the
matrix $\bm{Q}=\sqrt{\bm{M}^{-1}}\widetilde{\bm{Q}}$ corresponds to the
generalized QR decomposition of $\bm{A}$ such that $\bm{A}=\bm{Q}\bm{R}$ and
$\bm{Q}^{\top}\bm{M}\bm{Q}=\bm{I}_{m}$.

The generalized QR decomposition imposes almost no extra overhead compared to
standard QR when the matrix $\bm{M}$ is diagonal. As we will see, this is in
fact the case for the local metric induced by the majority of the commonly
used transfer functions such as leaky ReLU, sigmoid, and tanh.

Algorithm 1 $\text{{GQR}}(\bm{A},\bm{M})$: Generalized QR Decomposition

Input: matrix $\bm{A}\in\mathbb{R}^{m\times n}$ s.t. $n\leq m$, Riemannian
metric $\bm{M}\in\mathbb{S}^{m}_{+}$

Output: $\bm{Q}\in\mathbb{R}^{m\times n}$ and $\bm{R}\in\mathbb{R}^{n\times
n}$ factors such that $\bm{A}=\bm{Q}\bm{R}$ and
$\bm{Q}^{\top}\bm{M}\bm{Q}=\bm{I}_{n}$

$\widetilde{\bm{Q}},\bm{R}\leftarrow\text{{QR}}(\sqrt{\bm{M}}\bm{A})$

$\bm{Q}\leftarrow\sqrt{\bm{M}^{-1}}\widetilde{\bm{Q}}$

Return: $\bm{Q}$, $\bm{R}$

###  2.3 The Case of Softmax Transfer Function

Consider the case where the input examples are probability distributions
$\\{\mathbf{x}_{i}\in\Delta^{\\!d-1}\\}$ belonging to the $(d-1)$-simplex
$\Delta^{\\!d-1}=\\{\bm{u}\in\mathbb{R}^{d}_{+}|\bm{u}^{\top}\bm{1}_{d}=1\\}$.
The transfer function in this case corresponds to $f_{\text{\tiny
SM}}=\mathrm{softmax}$ which induces the KL Bregman divergence
$D_{F^{*}_{\text{\tiny}SM}}(\bm{u},\bm{v})=\sum_{j}(u_{j}\log\frac{u_{j}}{v_{j}})-u_{j}+v_{j}$.
Requiring $f_{\text{\tiny SM}}$ to be inevitable imposes the constraint
$\operatorname{dom}f_{\text{\tiny SM}}=\mathbb{R}^{d}-\\{\pm
c\bm{1}_{d},\,c\in\mathbb{R}_{+}\\}$ as $f_{\text{\tiny
SM}}(\bm{u}+c\,\bm{1}_{d})=f_{\text{\tiny SM}}(\bm{u})$ for $c\in\mathbb{R}$.
Thus, for the principal directions, we have
$\bm{1}_{d}\notin\operatorname{CS}(\bm{V})$ where $\operatorname{CS}$ denotes
column span. The above constraint can be easily incorporated into the
generalized QR decomposition in Algorithm 1 when using the Householder method
for the internal QR step.

The Householder method applies a series of _reflections_
$\bm{P}_{i}=\bm{I}_{m}-2\frac{\bm{b}_{i}\bm{b}_{i}^{\top}}{\bm{b}_{i}^{\top}\bm{b}_{i}}$
s.t. $\bm{b}\in\mathbb{R}^{m}$ and $i\in[n-1]$, on matrix $\bm{A}$ such that

$\bm{R}=\bm{P}_{n-1}\bm{P}_{n-2}\ldots\bm{P}_{2}\bm{P}_{1}\bm{A}$

is upper-triangular. The orthonormal matrix $\bm{Q}$ then can be written as

$\bm{Q}=\bm{P}_{1}\bm{P}_{2}\ldots\bm{P}_{n-2}\bm{P}_{n-1}\,.$

The following proposition shows that, in order to obtain
$\bm{Q}\in\text{St}^{\text{\scalebox{0.8}{$(\bm{M})$}}}_{d,k}$ from $\bm{A}$
s.t. $\bm{Q}^{\top}\bm{M}\bm{1}=0$ using Householder reflections, it suffices
to augment $\bm{A}$ from left by a column of all ones and apply Algorithm 1.
The resulting matrix $\bm{Q}$ corresponds to the first factor when the first
column is dropped. All proofs are relegated to the appendix.

######  Proposition 2.

Given $\bm{A}\in\mathbb{R}^{m\times n}$ s.t. $n<m$, let
$\widetilde{\bm{Q}},\widetilde{\bm{R}}=\emph{\text{{GQR}}}([\bm{1}_{m},\bm{A}],\bm{M})$
by applying Algorithm 1 using Householder reflections. The $\bm{Q}$ and
$\bm{R}$ factors, $\bm{Q}\bm{R}=\bm{A}$ s.t.
$\bm{Q}\in\text{\emph{\text{St}}}^{\text{\scalebox{0.8}{$(\bm{M})$}}}_{d,k}$
and $\bm{Q}^{\top}\bm{M}\bm{1}_{m}=0$, can be obtained from
$\widetilde{\bm{Q}}$ and $\widetilde{\bm{R}}$, respectively, by dropping the
first columns.

Our generalized Bregman PCA algorithm with a mean is given in Algorithm 2. We
omit the case of the softmax function in the main algorithm for simplicity
(see Appendix 3 for the full algorithm).

Figure 2: Visualization of (a) the output probabilities of the teacher model
$\\{\bm{y}_{i}^{L}\\}$ and (b)-(d) the reconstructed representation of the
training examples at the leaky ReLU layer $\\{\hat{\bm{y}}_{i}^{L-1}\\}$ using
TriMap [2]. As the number of components $k$ increases, the approximated
representations become better separable. Interestingly, TriMap reveals that
the representations at layer $L-1$ form similar clusters as the ones in layer
$L$. Algorithm 2 Bregman PCA with a Mean

Input: $\mathcal{X}=\\{\mathbf{x}_{i}\in\operatorname{dom}F^{*}\\}$, number of
components $k$

Output: dual mean $\bm{m}\in\operatorname{dom}F$,
$\bm{V}\in\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k}$,
$\\{\bm{c}_{i}\in\mathbb{R}^{k}\\}$

$\bm{m}\leftarrow
f^{*}\big{(}\frac{1}{|\mathcal{X}|}\sum_{i}\mathbf{x}_{i}\big{)}$

initialize $\bm{V}$ and $\\{\bm{c}_{i}\\}$

repeat

for $i\in[|\mathcal{X}|]$ do

$\hat{\mathbf{x}}_{i}\leftarrow f(\bm{m}+\bm{V}\bm{c}_{i})$

$\bm{c}_{i}\leftarrow\bm{c}_{i}-\eta_{a}\bm{V}^{\top}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})$

$\bm{V}\leftarrow\bm{V}-\eta_{\scalebox{0.5}{$V$}}\sum_{i}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})\bm{c}_{i}^{\top}$

until $\bm{V},\\{\bm{c}_{i}\\}$ not converged

$\bm{V},\bm{T}\leftarrow\text{{GQR}}(\bm{V},\bm{H}_{F}(\bm{m}))$

return $\bm{m}$, $\bm{V}$, $\\{\bm{T}\bm{c}_{i}\\}$

##  3 Representation Learning of Deep Neural Networks

One important application of our Bregman PCA is learning the representations
of a deep neural network in each layer. Specifically, in a deep neural
network, each layer transforms the representation that receives from the
previous layer and passes it to the next layer. In a given layer, we are
interested in learning the mean and principal directions that can encapsulate
the representations of all training examples in that layer. Although vanilla
PCA might be the naïve choice for this purpose, we will consider learning
better representations using our extended Bregman PCA approach.

A natural choice of a Bregman divergence for a layer having a strictly
increasing transfer function is the one induced by the convex integral
function of the transfer function [1]. In Bregman PCA, we essentially minimize
the matching loss of the transfer function $f$ instead of the quadratic
compression loss used for vanilla PCA. The matching loss was introduced in
[15, 19] and is also the main tool in the recent work on training deep neural
networks [3, 1]. Specifically, let $\bm{a}^{(\ell)}_{i}\in\mathbb{R}^{d}$ and
$\bm{y}^{(\ell)}_{i}\in\mathbb{R}^{d}$ respectively be the pre and post
(transfer function) activations of a neural network for a given input example
$\mathbf{x}_{i}$ at layer $\ell\in[L]$ having a (elementwise) strictly
increasing transfer function $f^{(\ell)}$ s.t.
$\bm{y}^{(\ell)}_{i}=f^{(\ell)}(\bm{a}^{(\ell)}_{i})$. Let $F^{(\ell)}$ denote
the convex integral function of $f^{(\ell)}$ where $f^{(\ell)}=\nabla
F^{(\ell)}$. Then, for a given set of input examples
$\mathcal{X}=\\{\mathbf{x}_{i}\\}$ having post-activation representations
$\mathcal{Y}^{(\ell)}=\\{\bm{y}^{(\ell)}_{i}\\}$, we can cast the Bregman PCA
problem in layer $\ell\in[L]$ as learning
$\bm{V}^{(\ell)}\in\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F^{(\ell)}}(\bm{m}^{(\ell)}))$}}}_{d,k}$
and $\\{\bm{c}_{i}^{(\ell)}\\}$ with
$\bm{m}^{(\ell)}=f^{*(\ell)}\big{(}\frac{1}{|\mathcal{Y}^{(\ell)}|}\sum_{i}\bm{y}^{(\ell)}_{i}\big{)}$
that minimize the objective

$\sum_{i}D_{F^{*(\ell)}\\!}(\bm{y}^{(\ell)}_{i},f^{\ell}(\bm{m}^{(\ell)}+\bm{V}^{(\ell)}\bm{c}^{(\ell)}_{i}))\,.$
(10)

We explore several ideas based on this layerwise construction, including
learning the principal directions in each layer in the next section.
Specifically, we consider representation learning in the final softmax layer
as well as the fully-connected leaky ReLU layer before the final softmax layer
of a ResNet-18 model. However, our construction is applicable to other types
of layers, such as convolutions. We then consider knowledge distillation using
the representations learned from the ResNet-18 teacher model to train a
smaller convolutional student model on the same dataset. Our knowledge
distillation approach is illustrated in Figure 1. Our approach consists of
learning the dual mean $\bm{m}$ and the principal $\bm{V}$ in the $\ell$-th
layer of the student network. The representation $\bm{y}_{i}^{\ell}$ of a
training example $\mathbf{x}_{i}$ is then approximated by
${\hat{\bm{y}}}_{i}^{\ell}=f^{\ell}(\bm{V}\bm{c}_{i}+\bm{m})$ where the
compression coefficients $\bm{c}_{i}\in\mathbb{R}^{k}$ are predicted by a
smaller student network. The approximated representation
${\hat{\bm{y}}}_{i}^{\ell}$ can then be passed through the rest of the pre-
trained teacher network or a smaller network that is trained from scratch to
predict the output labels. This approach can easily be extended to distilling
information from several different layers of the teacher model, where one can
use a cascade of student networks in which the approximate output
representation produced by the previous student network is passed as input to
the next student network. However, we defer exploring such extensions to
future work and focus on knowledge distillation using the representation of a
single layer of the teacher model.

##  4 Experiments

We conduct the first set of experiments on the CIFAR-10 and CIFAR-100
datasets, each consisting of $50000$ train and $10000$ test images of size
$32\times 32$ from, respectively, $10$ and $100$ classes.111Available at
https://www.cs.toronto.edu/~kriz/cifar.html. In order to extract
representations, we first train a PreAct ResNet-18 model using SGD with a
Nesterov momentum optimizer and a batch size of $128$. The only modification
we apply to the network is replacing the global average pooling layer before
the final dense layer with a flattening operator. This modification yields a
representation of dimension $8192$ before the final layer. We also change the
transfer function of that layer from a ReLU to a leaky ReLU with a small
negative slope of $\beta=1\mathrm{e}{-4}$ to obtain a strictly monotonic
transfer function. The network achieves $92.62\%$ and $70.44\%$ top-1 test
accuracy on CIFAR-10 and CIFAR-100, respectively. We then pass the train and
test examples of both datasets and store the leaky ReLU post-activations
($d=8192$) and the output softmax probabilities ($d=10$ and $d=100$ for
CIFAR-10 and CIFAR-100, respectively) of each example. In the following, we
consider several experiments for compressing each of these representations. We
use SGD with a heavy-ball momentum for optimizing our Bregman PCA problem. At
last, we also provide results on the ImageNet-1k datasets [12].

###  4.1 PCA on Output Probabilities

We first consider the problem of compressing the output probabilities. This
problem corresponds to a Bregman PCA with a softmax link function and a KL
divergence. For comparison, we consider vanilla PCA (Eq. (1)) in the pre-
activations (i.e. logits) domain. Note that vanilla PCA is not directly
applicable to the post-activation probabilities, as the reconstructed examples
may not correspond to valid probability distributions. We also consider CoDA
PCA [5] by augmenting a constant bias of $1$ to the coefficient to handle the
non-centered case, as suggested in the paper. We also use SGD with a heavy-
ball momentum for solving the CoDA PCA problem. We report the KL divergence
between the model probabilities and the reconstructed probabilities for the
train and test sets as the performance measure.

The average KL divergence values between the original output and the
reconstructed probabilities are shown in Figure 3(a)-(b). As can be seen from
the figure, Bregman PCA yields a much lower loss compared to vanilla PCA and
CoDA PCA. Also, notice that CoDA-PCA becomes inefficient as the dimension of
the problem ($d$ or $k$) increases.

###  4.2 PCA on Leaky ReLU Outputs

We consider the problem of compressing the representation in the leaky ReLU
layer before the final dense layer (a.k.a. the penultimate layer). The convex
conjugate function $F_{\beta}^{*}$ corresponding to a leaky ReLU transfer
function $f_{\beta}(\bm{u})=\max(\bm{u},\bm{0})-\beta\max(-\bm{u},\bm{0})$
with a slope $\beta>0$ amounts to
$F^{*}_{\beta}(\bm{u})=\nicefrac{{1}}{{2}}\,\bm{u}\odot
f_{\beta^{-1}}(\bm{u})$ where $\odot$ denotes Hadamard product. For this
problem, we consider vanilla PCA on the pre and post-leaky ReLU activations as
well as our Bregman PCA. In order to compare the methods, we pass the
reconstructed representation of each example through the final layer of the
trained ResNet-18 model and measure the top-1 accuracy.

The results are given in Figure 3(c)-(d). Our Bregman PCA works significantly
better than vanilla PCA for reconstructing the representation and yields a
much higher classification accuracy at the output. In order to visualize these
representations, we illustrate the 2-D TriMap [2] projections of the CIFAR-10
output probabilities as well as the reconstructed representations of the
examples at the leaky ReLU layer using our method. We can see from the figures
that the two representations, although one layer apart, form very similar
clusters. Also, as the number of components $k$ increases, the leaky ReLU
representations become better separable.

Figure 3: (a) and (b): KL divergence between the original and reconstructed
class probabilities of the train and test examples at the final layer. (Lower
values are better.) Our Bregman PCA approach significantly outperforms both
vanilla PCA on the logits as well as CoDA PCA. (c) and (d): Top-1
classification accuracy of the reconstructed leaky ReLU representations when
passed through the rest of the network. (Higher values are better.) Our
Bregman PCA works significantly better than vanilla PCA for reconstructing the
representation and yields a much higher classification accuracy.

###  4.3 Distillation Results on CIFAR-10/100

We conduct knowledge distillation experiments using our Bregman Representation
Learning (BRL) method. The architecture we consider for the experiments is
shown in Figure 1. The student network consists of a small convolutional
neural network with $5$ convolutional layers of size $[128\,(\times 2),$
$256\,(\times 2),$ $512]$ followed by a dense linear layer which outputs $k$
coefficients. These coefficients are then passed through a dense layer (with
weights $\bm{V}$ and bias $\bm{m}$), which outputs a representation of size
$8192$. The final dense (readout) layer with softmax transfer function
converts this representation into output class probabilities.

For comparison, we consider the following approaches: 1) Baseline model for
which all the weights are randomly initialized and are trained using the same
training examples as the teacher, 2) Baseline + Readout where all the weights
are initialized randomly except the readout layer, which is set to the pre-
trained weights of the teacher, 3) Baseline Distillation where we use a convex
combination of the teacher’s soft labels (for which we also tune a temperature
for the logits) with the one-hot training labels. The ratio of the teacher’s
soft labels to train labels is tuned for each case. 4) BRL, where we randomly
initialize all the weights except $\bm{V}$ and $\bm{m}$, which are set to the
values obtained by applying our Bregman PCA on the leaky ReLU representations
of the training examples, as in the previous section. 5) BRL + Readout, which
is similar to BRL, but we set the readout layer weights to the teacher
weights. Each model is trained for $51$ epochs using a batch size of $128$
using SGD with a Nesterov momentum optimizer with a linear decay schedule. In
each case, we tune the learning rate and momentum hyper-parameters.

To train BRL, we first train the student model by minimizing a squared loss
between the predicted coefficients of the student and the compression
coefficients obtained by directly applying our Bregman PCA to the teacher’s
representations (as in the previous section). This essentially reduces the
training of the student network from classification to a regression problem.
We find that this pre-training on the teacher coefficients significantly
improves the convergence of the student model. After training on the
compression coefficients for $25$ epochs, we switch to directly minimizing the
cross-entropy loss at the output layer for the rest of the iterations.

To disentangle the improvements due to better information transfer for
knowledge distillation from the gains obtained due to data augmentation, we
conduct the experiments on the original examples without any additional
augmentations (such as random horizontal flip, random crop, Mixup [34], etc.),
which are standard for training these models. This resembles a limited sample
regime where the number of training examples is relatively small. We conduct
additional experiments by adding these augmentations and observe improvements
for all baselines (see the appendix). In order to highlight the limited
samples aspect even further, we also repeat each experiment using only half of
the training dataset, which is uniformly sampled from the full dataset (and
fixed for all methods).

Figure 4: Results of training different baselines, including soft-label
teacher-student distillation and our BRL distillation approach on the full and
$50\%$ randomly subsampled (half) CIFAR-10 and CIFAR-100 datasets. In each
case, $+$ Readout indicates using the teacher’s pre-trained readout weights
for the student.

The results of training the smaller student model are shown in Figure 4. We
observe a significant improvement in accuracy using our BRL distillation
approach across different numbers of components ($k$). For CIFAR-10, even with
a small $k$, BRL shows a major advantage over soft-label teacher-student
training. For CIFAR-100, these improvements are minor for smaller $k$ values
as the leaky ReLU representations require a relatively larger number of
components to be reconstructed efficiently compared to CIFAR-10 (see Figure
3(c)-(d)). However, as $k$ increases, we observe that the added benefit from
more Bregman representations outweighs the benefit of a higher capacity model
(due to larger $\bm{V}$) for the baseline distillation approach using soft
labels ($56.5\%$ top-1 accuracy for BRL compared to $54.1\%$ for Baseline
Distillation).

Additionally, in Figure 4, we observe that the gap between BRL and soft-label
distillation becomes even more prominent as the number of training examples
decreases ($6.3\%$ gap for full vs. $9.0\%$ gap for half CIFAR-10 dataset).
This shows that BRL is a more sample-efficient approach for transferring
information between the teacher and the student models. Finally, in both
cases, we see a marginal improvement for using the pre-trained readout layer
weights over training the readout layer from scratch for larger $k$ values.

###  4.4 Distillation Results on ImageNet-1k

We consider the problem of distilling a PreAct ResNet-50 teacher model trained
on the ImageNet-1k dataset into a PreAct ResNet-18 student. Similar to the
previous experiments, we replace the transfer function of the penultimate
layer with a leaky ReLU with a negative slope of $\beta=1\mathrm{e}{-2}$. The
teacher model is trained for 90 epochs using SGD with momentum with a
piecewise constant schedule and achieves a $76.2\%$ top-1 test accuracy. The
baseline ResNet-18 model using the same training procedure achieves a $70.1\%$
top-1 test accuracy. As the distillation baseline, we consider soft-label
teacher-student training [16] and tune the soft-label ratio and the logit
temperature, and train the model for 90 epochs using the same optimization
hyperparameters as the ResNet-18 baseline. The soft-label distillation
baseline achieves a $71.5\%$ top-1 test accuracy.

Model | Top-1 Accuracy | Top-5 Accuracy
---|---|---
ResNet-50 Teacher | $76.2\%$ | $93.0\%$
ResNet-18 Baseline | $70.1\%$ | $89.4\%$
ResNet-18 Soft-label Distillation [16] | $71.5\%$ | $90.4\%$
ResNet-18 BRL (Ours) | $\bm{72.6}\%$ | $\bm{91.1}\%$
Table 1:  ImageNet-1k distillation results using a ResNet-50 teacher and a
ResNet-18 student. Each model is trained for $90$ epochs.

The teacher’s penultimate layer has dimension $d=$2048$$ which we compress
into $k=512$ principal directions. To extract the Bregman representations and
the mean vector for BRL, we simply apply our Bregman PCA algorithm in an
online fashion in the last 10 epochs of the teacher model’s training. (To find
the mean, we apply an exponential moving average.) We then import and fix
these representations into the student model. We replace the last ResNet-18
layer with a linear layer to produce the $512$ compression coefficients. We
also use the weights of the teacher’s readout layer at the student’s output.
We replace the optimizer with Adam [18] with a cosine learning rate schedule.
Similar to soft-label distillation, we consider the teacher’s soft labels with
the same ratio as the distillation baseline. We also add a squared regularizer
term between the teacher’s target and the student’s predicted compression
coefficients. Our BRL method combined with the teacher’s soft labels achieves
a $72.6\%$ top-1 test accuracy after 90 epochs. The results are summarized in
Table 1.

##  5 Conclusion and Future Work

We presented a new direction for knowledge distillation based on directly
transferring the Bregman representations learned by the teacher model into the
student model. Our work significantly advances the previous approaches for
knowledge distillation which rely on either training using soft teacher labels
or matching the representations of the teacher and the student in the
intermediate layers. Our construction is more flexible and can be extended to
training a cascade of student networks, in which each model constructs the
representation that is passed as the input to the next model. Such training
strategies are interesting future directions for our technique.

One limitation of our current approach is that it only applies to layers with
strictly monotonic transfer functions. The extension of our work to non-
monotonic transfer functions is an interesting open problem.

## References

  * [1] Ehsan Amid, Rohan Anil, and Manfred K. Warmuth.  LocoProp: Enhancing BackProp via local loss optimization.  In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
  * [2] Ehsan Amid and Manfred K. Warmuth.  TriMap: Large-scale Dimensionality Reduction Using Triplets.  arXiv preprint arXiv:1910.00204, 2019.
  * [3] Ehsan Amid, Manfred K. Warmuth, Rohan Anil, and Tomer Koren.  Robust bi-tempered logistic loss based on Bregman divergences.  In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NeurIPS, 2019.
  * [4] Rohan Anil, Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George E Dahl, and Geoffrey E Hinton.  Large scale distributed neural network training through online distillation.  arXiv preprint arXiv:1804.03235, 2018.
  * [5] Marta Avalos, Richard Nock, Cheng Soon Ong, Julien Rouar, and Ke Sun.  Representation learning of compositional data.  In Advances in Neural Information Processing Systems, 2018.
  * [6] Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov.  Knowledge distillation: A good teacher is patient and consistent.  arXiv preprint arXiv:2106.05237, 2021.
  * [7] Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov.  Knowledge distillation: A good teacher is patient and consistent.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10925–10934, 2022.
  * [8] Lev M Bregman.  The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming.  USSR computational mathematics and mathematical physics, 7(3):200–217, 1967.
  * [9] Julien Chiquet, Mahendra Mariadassou, and Stéphane Robin.  Variational inference for probabilistic poisson pca.  The Annals of Applied Statistics, 12:2674–2698, 2018.
  * [10] Michael Collins, Sanjoy Dasgupta, and Robert E. Schapire.  A generalization of principal component analysis to the exponential family.  In Proceedings of the 14th International Conference on Neural Information Processing Systems: Natural and Synthetic, NIPS’01, page 617–624, 2001.
  * [11] Wojciech Marian Czarnecki, Simon Osindero, Max Jaderberg, Grzegorz Świrszcz, and Razvan Pascanu.  Sobolev training for neural networks.  arXiv preprint arXiv:1706.04859, 2017.
  * [12] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei.  ImageNet: A Large-Scale Hierarchical Image Database.  In CVPR09, 2009.
  * [13] Gene H. Golub and Charles F. Van Loan.  Matrix Computations.  The Johns Hopkins University Press, third edition, 1996.
  * [14] Allen Hatcher.  Algebraic topology.  Cambridge Univ. Press, Cambridge, 2000.
  * [15] D. P. Helmbold, J. Kivinen, and M. K. Warmuth.  Relative loss bounds for single neurons.  IEEE Transactions on Neural Networks, 10(6):1291–1304, 1999.
  * [16] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [17] Jean-Baptiste Hiriart-Urruty and Claude Lemaréchal.  Fundamentals of Convex Analysis.  Springer-Verlag Berlin Heidelberg, first edition, 2001.
  * [18] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [19] J. Kivinen and M. K. Warmuth.  Relative loss bounds for multidimensional regression problems.  Journal of Machine Learning, 45(3):301–329, 2001.
  * [20] Andrew J. Landgraf and Yoonkyung Lee.  Dimensionality reduction for binary data through the projection of natural parameters.  Journal of Multivariate Analysis, 180:104668, 2020.
  * [21] Andrew J. Landgraf and Yoonkyung Lee.  Generalized principal component analysis: Projection of saturated model parameters.  Technometrics, 62:459 – 472, 2020.
  * [22] John M Lee.  Riemannian manifolds: an introduction to curvature, volume 176.  Springer Science & Business Media, 2006.
  * [23] Rafael Müller, Simon Kornblith, and Geoffrey Hinton.  Subclass distillation.  arXiv preprint arXiv:2002.03936, 2020.
  * [24] Erkki Oja.  Simplified neuron model as a principal component analyzer.  Journal of mathematical biology, 15(3):267–273, 1982.
  * [25] Nicolas Papernot, Patrick Mcdaniel, Xi Wu, Somesh Jha, and Ananthram Swami.  Distillation as a defense to adversarial perturbations against deep neural networks.  2016 IEEE Symposium on Security and Privacy (SP), pages 582–597, 2016.
  * [26] Karl Pearson.  On lines and planes of closest fit to systems of points in space.  Philosophical Magazine Series 1, 2:559–572, 1901.
  * [27] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets.  arXiv preprint arXiv:1412.6550, 2014.
  * [28] Nicholas Roy and Geoffrey J Gordon.  Exponential family pca for belief compression in pomdps.  Advances in Neural Information Processing Systems, 15:1667–1674, 2002.
  * [29] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  arXiv preprint arXiv:1910.01108, 2019.
  * [30] Jonathon Shlens.  A tutorial on principal component analysis.  arXiv preprint arXiv:1404.1100, 2014.
  * [31] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive representation distillation.  arXiv preprint arXiv:1910.10699, 2019.
  * [32] Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V Le.  Self-training with noisy student improves imagenet classification.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10687–10698, 2020.
  * [33] Sergey Zagoruyko and Nikos Komodakis.  Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer.  arXiv preprint arXiv:1612.03928, 2016.
  * [34] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  In ICLR, 2018.

##  Appendix A Omitted Proofs

See 1

###### Proof.

Using the duality property of Bregman divergence, Eq. (5) can be written as

$\min_{\tilde{\bm{m}}\in\operatorname{dom}F}\big{[}\sum_{i}D_{F^{*}\\!}(\mathbf{x}_{i},f(\tilde{\bm{m}}))=\sum_{i}D_{F}(\tilde{\bm{m}},f^{*}(\mathbf{x}_{i}))\big{]}\,.$

Applying Eq. (4) and setting the derivative of above wrt $\tilde{\bm{m}}$ to
zero, we have

$\sum_{i}\big{(}f(\bm{m})-f(f^{*}(\mathbf{x}_{i}))\big{)}=|\mathcal{X}|\cdot
f(\bm{m})-\sum_{i}\mathbf{x}_{i}=\bm{0}\,.$

Rearranging the term and applying the inverse yields the result. ∎

See 1

###### Proof.

Since $\widetilde{\bm{Q}}\bm{R}=\sqrt{\bm{M}}\bm{A}$, we have

$\bm{Q}\bm{R}=\sqrt{\bm{M}^{-1}}\widetilde{\bm{Q}}\bm{R}=\sqrt{\bm{M}^{-1}}\sqrt{\bm{M}}\bm{A}=\bm{A}\,,$

and $\bm{Q}$ satisfies

$\bm{Q}^{\top}\bm{M}\bm{Q}=\widetilde{\bm{Q}}^{\top}\sqrt{\bm{M}^{-1}}\bm{M}\sqrt{\bm{M}^{-1}}\widetilde{\bm{Q}}=\widetilde{\bm{Q}}^{\top}\widetilde{\bm{Q}}=\bm{I}_{m}\,.$

∎

See 2

###### Proof.

The first reflection $\bm{P}_{1}$ of the Householder method transforms the
first column of the input matrix into a unit vector $s_{1}\bm{e}_{1}$ where
$s_{1}\in\mathbb{R}$. Thus, the first column of the resulting
$\widetilde{\bm{Q}}$ term is just a rescaling of the first column of the input
matrix. Thus, for the augmented matrix $[\bm{1}_{m},\bm{A}]$, the first column
of $\widetilde{\bm{Q}}$ corresponds to $r_{1}\bm{1}_{m}$ with
$r_{1}\in\mathbb{R}$ s.t. $r_{1}^{2}\bm{1}_{m}^{\top}\bm{M}\bm{1}_{m}=1$. The
remaining columns of $\widetilde{\bm{Q}}$, i.e., the matrix $\bm{Q}$,
correspond to the column space of $\bm{A}$ minus the direction $\bm{1}_{m}$.
Also, we have $\bm{Q}^{\top}\bm{M}\bm{Q}=\bm{I}_{n}$ and
$\bm{Q}^{\top}\bm{M}\bm{1}_{m}=0$. Note that the direction $\bm{1}_{m}$ is
redundant for the softmax function, i.e.
$\mathrm{softmax}(\bm{Q}\bm{R})=\mathrm{softmax}(\bm{A})$ where the softmax
function is applied to each column and $\bm{R}$ corresponds to the upper-
triangular matrix by removing the first row and column of
$\widetilde{\bm{R}}$. ∎

##  Appendix B Full Bregman PCA with a Mean Algorithm

We provide the full procedure for the Bregman PCA algorithm with mean,
including the case of the softmax function.

Algorithm 3 Bregman PCA with a Mean (Including Softmax)

Input: $\mathcal{X}=\\{\mathbf{x}_{i}\in\operatorname{dom}F^{*}\\}$, number of
components $k$

Output: dual mean $\bm{m}\in\operatorname{dom}F$,
$\bm{V}\in\text{St}^{\text{\scalebox{0.8}{$(\bm{H}_{\\!F}(\bm{m}))$}}}_{d,k}$,
$\\{\bm{c}_{i}\in\mathbb{R}^{k}\\}$

$\bm{m}\leftarrow
f^{*}\big{(}\frac{1}{|\mathcal{X}|}\sum_{i}\mathbf{x}_{i}\big{)}$

initialize $\bm{V}$ and $\\{\bm{c}_{i}\\}$

repeat

for $i\in[|\mathcal{X}|]$ do

$\hat{\mathbf{x}}_{i}\leftarrow f(\bm{m}+\bm{V}\bm{c}_{i})$

$\bm{c}_{i}\leftarrow\bm{c}_{i}-\eta_{a}\bm{V}^{\top}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})$

$\bm{V}\leftarrow\bm{V}-\eta_{\scalebox{0.5}{$V$}}\sum_{i}(\hat{\mathbf{x}}_{i}-\mathbf{x}_{i})\bm{c}_{i}^{\top}$

until $\bm{V},\\{\bm{c}_{i}\\}$ not converged

if $f$ is $\mathrm{softmax}$ then

$\bm{V}\leftarrow[\bm{1}_{d},\bm{V}]$ $\triangleright$ Augment a column of all
ones.

$\bm{V},\bm{T}\leftarrow\text{{GQR}}(\bm{V},\bm{H}_{F}([0;\bm{m}]))$
$\triangleright$ Augment a zero element to the mean vector.

else

$\bm{V},\bm{T}\leftarrow\text{{GQR}}(\bm{V},\bm{H}_{F}(\bm{m}))$

if $f$ is $\mathrm{softmax}$ then

$\bm{V}\leftarrow\bm{V}[:,1\\!:]$ $\triangleright$ Drop the first column.

$\bm{T}\leftarrow\bm{T}[1\\!:,1\\!:]$ $\triangleright$ Drop the first row and
column.

return $\bm{m}$, $\bm{V}$, $\\{\bm{T}\bm{c}_{i}\\}$

##  Appendix C Further Details on the Experiments

CIFAR-10/100 Experiments: For each experiment, we tune the learning rate and
the momentum hyperparameters in the range $[10^{-5},0.2]$ and $[0.5,0.99]$,
respectively. The results are averaged over $5$ independent trials for each
experiment. The input images are normalized for each dataset by subtracting
the mean and dividing by the standard deviation over the whole training set.

Compute Resources: For the CIFAR-10/100 experiments, we use V-100 GPUs. For
the ImageNet-1k experiments, we use $8\times 8$ TPU v3s.

##  Appendix D Experimental Results with Data Augmentation and Mixup

Data augmentation is a standard procedure for training deep neural networks.
However, our main results in the paper focus on decoupling the information
transfer from the teacher due to the knowledge distillation technique from the
generalization improvement due to data augmentation. In this section, we
showcase results using standard augmentation procedures (random horizontal
flip and random crop) combined with Mixup [34] with $\alpha=\beta=1$. Table 2
presents the results of different methods with $k=100$ components and $51$
training epochs with a batch size of $128$.

Table 2: Top-1 accuracy for different methods with number of components $k=100$ and using data augmentation combined with Mixup [34]. Method |  Datset
---|---
CIFAR-10 | CIFAR-100
Baseline | $83.65\pm 0.41$ | $54.25\pm 0.35$
Baseline + Readout | $83.34\pm 0.51$ | $55.11\pm 0.34$
Baseline Distillation | $86.00\pm 0.22$ | $62.25\pm 0.15$
BRL | $\bm{87.03\pm 0.12}$ | $\bm{62.72\pm 0.28}$
BRL + Readout | $86.24\pm 0.29$ | $61.11\pm 0.29$

From the table, we see that the performance of all methods consistently
improves with data augmentation. Additionally, among all methods, BRL performs
the best. However, the difference between BRL and soft-label knowledge
distillation is not as prominent as before (see Figure 4). Nonetheless, such
behavior is not surprising as we expect soft-label knowledge distillation and
our method to eventually reach the teacher’s performance, given that a
sufficient amount of data (i.e., training examples) for distillation is
provided [7]. As we show in the main experiments, BRL delivers a clear
advantage in a setup with a limited number of samples.