Layerwise Bregman Representation Learning of Neural Networks with Applications to Knowledge Distillation
Authors: Ehsan Amid
Rohan Anil
Christopher Fifty
Manfred K Warmuth
Venue: Accepted by TMLR ()

Abstract: We propose a new method for layerwise representation learning of a trained neural network that conforms to the non-linearity of the layer's transfer function. In particular, we form a Bregman divergence based on the convex function induced by the layer's transfer function and construct an extension of the original Bregman PCA formulation by incorporating a mean vector and revising the normalization constraint on the principal directions. These modifications allow 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 is then 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 that uses the teacher's soft labels. 

Published in Transactions Machine Learning Research (02/2023)
Layerwise Bregman Representation Learning of Neural Networks with Applications to Knowledge Distillation
Ehsan Amid', Rohan Anill, Christopher Fifty?, Manfred K. Warmuth
'Google Research, 2 Stanford University
Reviewed on OpenReview:
https: // openreview. net/ forum? id= 6dsvH7pQHH
Abstract
We propose new method for layerwise representation learning of trained neural network that conforms to the non-linearity of the layer's transfer function. In particular, we form a Bregman divergence based on the convex function induced by the layer's transfer function and construct an extension of the original Bregman PCA formulation by incorporating a mean vector and revising the normalization constraint on the principal directions. These modifications allow exporting the learned representation as a fixed layer with a non-linearity. As an application to knowledge distillation, we cast the lcarning problem for the student network predicting the compression coefficients of the teacher's representations, which is then 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 that uses the teacher's soft labels.
1 Introduction
Representation lcarning (Bengio al.] 2013) is at the core of Machine Learning (ML) applications. Larger MT. models generally extract better explanatory factors from the data, resulting in improved generalization. Such large-scale ML systems train
vast amount of data. A concrete example is the Ads CTR prediction model (Anil et al.] 2022). To effectively learn the intrinsic characteristics of the data (such as user behaviors and seasonal patterns, in our example), such systems rely on training large models with billions of parameters. Yet, such outrageously large models are cumbersome when serving live traffic; more lightweight models, which allow serving thousands of users simultancously, are necessary for deployment. However, the improved capability of the model is implicitly encoded in its learned weights, making it infeasible to transfer directly to smaller network architectures. A common alternative strategy is to use the learned representation (i.c., activations) of possibly larger models as guidance for training smaller networks (Gou et al.) 2021 From this perspective, knowledge distillation is the critical component in compressing larger models with better representations into more compact variants without sacrificing much of their quality.
Knowledge distillation refers to a set of techniques used for transferring information from typically a larger trained model, called the teacher, a smaller model, called the student (Hinton al.)
2015 Anil et al. 2018 Xie et al. 2020). The goal of distillation is to improve the performance of

---

Published in Transactions Machine Learning Research (02/2023)
the student model by augmenting the knowledge of the teacher model with the raw information provided by the set of train examples. Since its introduction, several previous works have applied knowledge distillation to obtain improved results for language modeling (Sanh et al., 2019: image classification (Beyer et al. 2022). and robustness against adversarial attacks (Papernot et al. 2016)
The teacher's knowledge is lypically 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 (Hinton et al. 2015 Müller al.) 2020). Other approaches consider matching the teacher's representations for a given input, typically in the penultimate layer, by the stu- dent (Romero et al. 2014]. However, direct transfer of knowledge in the form of compressed representations has not been sufficiently explored before.
In this paper, we first introduce a technique for extracting compressed representations from an ar- bitrary layer of a trained model. Our technique conforms to the non-linearity of the layer's transfer function. Considering the geometry induced by the non-linearity allows us to learn more compact representations than the baseline approach of using standard PCA (Pearson 1901]. As an applica- tion to our layerwise representation learning approach, we explore the idea of directly transferring information from teacher to student in the form of learned (fized) principal directions of ar- bitrary layers of the teacher model. Our focus for representation learning will be generalized form of the PCA method based on the broader class of Bregman divergences. The Bregman di- vergence (Bregman 1967) induced by the strictly convex and differentiable function F Rd + between u, v E dom F defined as
v) F(u) F(v) f(u) (u-
(1)
where f V 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 F*(u") = sup, 2 . =-" F(z) be the Legendre dual (Hiriart-Urruty Lemaréchal, 2001) of F. Then, we can write Dr(u, v) DF-(v*, u²) where u* = f(u) and e" = f(v) are the pair of dual variables. When F is strictly convex and diffcrentiable, we have f" = f-1 and u u= f"(") and f=()). Also, the derivative of Bregman divergence with respect to the first argunent takes the following simple form
VuDr(u, = f(u) - f(v).
(2)
The loss construction in Amid et al. 2022 provides a natural way of generating layerwise Bregman divergences for deep neural networks as line integrals of the strictly monotonic transfer functions. dual view of such Bregman divergences in one dimension is illustrated in Figure I(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 to Bregman divergences was done by Collins et al. (2001) and extended in several previous works (Roy & Gordon 2002; Acharyya 2006 Chiquet et al. 2018 However, our generalization differs largely from the previous approaches by introducing a mean vector for handling non-centered data. Also, the orthogonality constraint in terms of Euclidean geometry extended to orthonormal k-frames with respect the Riemannian metric evaluated at the mean. To efficiently apply this constraint, we introduce variant of the well-known QR
2

---

Published in Transactions Machine Learning Research (02/2023)
163
1. (Pro Poj) Dri8.s)
Cutgg Ree) ar Dris.a)
Layer Throngh
Teachur Network Layer
omtput W.
pratt
[16]
pre-activation
Student Netwoek
(a)
Figure 1: (a) A dual view of Bregman divergences the area under the curves of the transfer function f and ils inverse I" s- (Helmbold et al. 1999 Amid el al. 2022). (b) An illustration of knowledge distillation using Bregman representation learning. Given input example æ,, the output representation y of the teacher network at layer é L is approximated as of = (VV++tm where the principal directions V and the mean vector m are learned using Bregman PCA on the full train set. The compression coefficients C; are predicted by the student network, and the approximate representation of passed instead of y through the rest of the network from layer f 1 to L. The teacher component inside the dashed bounding box can be discarded once the student network trained.
decomposition. We apply layerwise representation learning approach to distilling knowledge from larger teacher network smaller student. Our extension of the Bregman PCA formulation can be viewed as adding an extra layer to the student network with fixed weights and biases. The distillation problem thus reduces to learning the corresponding compression coefficienss of given example by the student, which is passed as input to this layer. Our knowledge distillation approach of using Bregman representation learning summarized in Figure] [b).
1.1 Related Work
Representation Learning and 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 Romero et al. (2014) where they add regression loss between the representation of the teacher and thinner (in the number units) student representation. 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 Zagoruyko & Komodakis (2016) where the students are less deep and use the teacher's spatial attention map guidance. Additionally, Czarnecki et al. (2117) propose incorporating the derivative of the teacher's prediction w.r.t to the imput to be matched by the student, and this additional information has been shown to improve the distillation procedure, More recently, Tian et al. (2019) propose adding contrastive losses to transfer representation between two networks and show improvement over vanilla distillation. However, their procedure requires extra computation push apart the teacher's representation randomly sampled inputs from the student's representation of the actual inputs. Finally, Müüller et al. (2020) tackle the task of extracting more information from the teacher model when the
their work, they learn a
3

---

Published in Transactions Machine Learning Research (02/2023)
number of classes small. They propose creating sub-classes corresponding to every class, followed by soft-label training for distillation.
Generalized PCA: Collins et al. (2001) 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 (Chiquet et al. 2018 and compositional data (Avalos et al., 2018 Landgraf & Lee (020b) apply logistic PCA for binary data for dimensionality reduction. Landgraf & Lee (2020a) also propose a ittioorratioommmmmmiaattt approach for solving the generalized PCA problem. Despite the success of earlier PCA extensions, a direct generalization of the original PCA formulation that includes mean vector and more principled orthogonality constraint the directions missing; extension adds mean vector while taking the local geometry of the space around the mean into account.
2 Extended Bregman PCA
Principal Component Analysis (PCA) Pearson 1901 Shlens 2014) is perhaps one of the most com- monly 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 X' {æz e [R4) and k < d, the PCA problem can be formulated as finding the mean vector m Rd and principal directions V E Raxk where V Is such that the compression loss,
m, V, (ci) = arg min
Σ < - (m vm)))
(3)
[MERR) [ VEStakk (2TER))
is minimized. Here, Sta,k (U E Rack U [x] denotes the Stiefel manifold of k-frames in Rd (Hatcher 2000) and C, T corresponds to the compression coefficients of x;, The problem in Eq. (3) can be solved effectively in two steps. First, we note that m can be viewed as a constant shared representation for all points in X' (for which the code length is zero) that minimizes the total compression loss. With this interpretation, the mean vector can be written as the minimizer of
m- arg min Σ 1| -1 ml mERd
(4)
for which, the solution corresponds to the geometric mean m কर. E. By fixing m, the solution for V and {c:} can be obtained by enforcing the orthonormality constraints using a set of Lagrange multipliers and setting the derivatives to zero. The solution to V amounts to the top-k eigenvectors of the covariance matrix 1/|x1 (,) m)) and c; = V (a,-n) m) corresponds to the projection of the centered point onto the column space of V. Noticeably, the online variants of PCA, such as Oja's algorithm (Oja 1982 alternatively apply a gradient stop V and project the update onto Sta,ı by an application of QR decomposition (Golub & Van Loan 1996).
Given a strictly convex and differentiable function F: Rd + R with link function f = VF, we cast the generalized Brogman PCA problem as approximating 2; E X as a lincar combination of a set of orthonormal principal directions in the dual space. This formulation reduces to Eq. (3) for the choice of f = idg, which corresponds to the squared Euclidean divergence. However, before defining the objective function formally, consider the problem of finding a generalized mean in the dual
4

---

Published in Transactions Machine Learning Research (02/2023)
space as follows. Let X fa, dom F*} be a set of given data points, We define the generalized mean vector m the minimizer of the following objective,
m arg min
Σ Dro f(m)). nn F
(5)
Note that the above Eq. 5 is a direct generalization of Eq. in terms of finding a shared constant representation for all points in 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
m = f
X Σxae xz)
(6)
Thus, the dual mean simply corresponds the dual of the arithmetic mean of the data points. When = ida, ie., the identity function, 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. 3 to the class of Bregman divergences. First, we note that the geometry of the space of principal directions is altered when switching from the squared loss to a more general Bregman divergence. Specifically, given the convex function F of a Bregman divergence, the inner product locally governed by the Riemannian metric
(Lee 2006).
Dr(u + ou, u) (F, a + ōu) 1/2 Hr(u)ōu,
where du is a small perturbation and Hr = V2F the Hessian of F. Thus, the definition of orthonormality needs conform the new geometry imposed by the Bregman divergence. In the following, extend the definition of a Stiefel manifold to include Riemannian metric. Recall that for strictly convex function F, we have HF 54+ where S" denotes the set of symmetric positive definite matrices. We now formally define the Riemann-Stiefé manifold.
Definition 1. The Riemann-Stiefe manifold of k-frames in Rd with respect to the Riemannian metric M M(v) E gd 0 Rd defined as
St = IU E pdxk UUMU- -U--K] Is).
Note that for the Euclidean geometry, M(v) = Ia for all v E Rdd Thus, the definition of (az) recovers St(fa) Sty, and we arrive at the orthonormality in the sense of the Euclidean geometry.
We now formulate the generalized Bregman PCA problem following the local geometry of the strictly convex function at m. Let X' (æ, dom F*) be set of given points. We define the generalized Bregman PCA as finding a linear combination of set of generalized principal directions that minimizes the Bregman compression loss:
(ci) arg min (VSt (
Σ Dr-(2n, J(m + Vε.)),
{
where m given via Eq. (6) and HF(m)= F(m). Note that the constraint E Sty, Eq. (3) is now replaced with Stak (m)))))) i.e., using the Riemannian metric induced by the Hessian of the convex function F evaluated at the dual mean m (Amari 2016).
5

---

Published in Transactions Machine Learning Research (02/2023)
2.1 Optimization
The generalized Bregman PCA objective in Eq. 7 does not yield closed-form solution in terms of V and {c}) However, the problem is convex in both V and {cs) and can be solved iteratively by applications of gradient descent steps. Let d; == f(m + Vea) denote the approximation of æq. For the compression coefficients {cih, we apply
G = Ta V (#: -x)),
where To >> 0 denotes the learning rate. Updating V involves two steps: gradient updates followed by projection onto St Stikk We apply gradient descent updates,
Vuew = V - Σ (á. xz)
(9)
where Fy > 0 denotes the learning rate. The next step involves projecting Voew back to (He)) (mn)) d,k As we shall see, this projection needs to be applied only once at the end of optimization since both V and {cz} are trained using gradient descent (any intermediate factor can be absorbed into the gradients). For the vanilla PCA where Hp(m)= = Ia. this can be applied easily by an application QR docomposition. We provide a simple modification of the standard QR dccomposition algorithm that achieves this for any H (m(m) E 87+ with almost no additional overhead in practice for our application.
2.2 QR Decomposition on the Riemann- Seeenen--Stefel Stiefel Manifold
A QR decomposition is factorization of trix A E Rmxn where I into product A QR where Q E Smm.n) and R E Rax= is upper-triange trianglri matrix. The first factor Q can be viewed as an orthonormalizaion of columns of A, similar to the result of Gram-Schmidt procedure. However, QR decomposition provides more nu- merically stable procedure in general. The method of Houscholder reflections (Golub & Van Loan 1996) is the most common algorithm for QR de- composition.
Algorithm RS-QR(A, M): QR Decomposi-
tion the Eiemann-Stteffe Manifold
Input: matrix A E RR s.t. S m, Rie- mannian metric M E ST.
Output: Q E mR and R E RRxn factors such that = QR and Q MQ = I, Q. R. + QR( QR(vMA A)
Q+ VM- Return: Q, R
The following theorem provides a procedure that extends the standard QR decomposition to produce conjugate factors Q TM I, for given M E +
Theorem 1. Let QR denote the procedure that returns the QR factors. Ginen M $4 and
A E Ruxn, let Q QR(VM Then, the malria Q = VM-1Q corresponds to the first factor of the QR decomposition of A on the Riemann-Ster Stiefel manifold such that A = QR and Q MQ = -m-
The QR decomposition on the Riemann-Sticfel manifold imposes almost no extra overhead com- pared lo standard QR when the matrix M is diagonal. As we shall see, this is in fact the case for the local metric induced by the majority of the commonly used elementwise transfer functions such as leaky ReLU, sigmoid, and tanh. We defer the special case of the QR decomposition for the
6

---

Published in Transactions Machine Learning Research (02/2023)
Dhipptt Pribabiitien
Resemttrrtedd Aeprosmtationn (
Recoretructed Regrommtation
Recondracied Seproettation (R 110)
(a)
(b)
(d)
Figure 2: Visualization of (a) the output probabilities of the teacher model {y and (b)-(d) the reconstructed representation of the training examples at the penultimate leaky ReLU layer (0, using TriMap (Amid & Warmuth 2019). As the number of components k = 20, 50, 100 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 of the original teacher model in (a).
Algorithm 2 Bregman PCA with Mean
Input: X' {a, dom F Ç Rd), number of components l < d Output: dual mean m dom F,, V St (Hr(m)) {a R*}
m < ( Σ, m:) initialize V and {c}} repeat
for E [[]]] do
i; 4 s'm + Vc;)
C c: Ta V (#: x;
V V - (á; - c.
until V, {c,} not converged
V,T + RS-QR(V, (m)) return m, V. (Te}}
softmax function Appendix B Our generalized Bregman PCA algorithm with mean is given in Algorithm 2 We omit the case of the softmax function in the main algorithm for simplicity (see Appendix C for the full algorithm).
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 deep neural network, each layer transforms the representa- tion 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 reprosentations of all training examples that layer. Although the vanilla PCA might be the naïve choice for this purpose, we will consider learning more effective representations using our extended Bregman PCA approach.
7

---

Published in Transactions Machine Learning Research (02/2023)
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 (Amid et al. 2022 In Bregman PCA, we essentially want minimize the matching loss of the transfer function f instead of the quadratic compression loss used for the vanilla PCA. The matching loss was introduced by Helmbold et al.] (1999); Kivinen & Warmuth| (2001) and is also the main tool in the recent work on training deep neural networks (Amid al. 2019 2022). Specifically, let a E R and (8)
(4)
y, E Rd respectively be the pre and post (transfer function) activations of a neural network for a given input example æ, at layer E [L]. The layer applies an (elementwise) strictly increasing (a,"", Let F()) denote convex integral function of f(0, ie., f(() V Then, for given set of input examples == {a,} having post-activation representations y(t) = {v, ()) we can cast the Bregman PCA problem in layer l E [L] as learning v(0) Si (o and {ct ) with (!) = - DIOTT
transfer function f(o), i.e.; y: (f)
Σ: E.y!")
( that minimize the objective
Dpriolus (0) ('m( + v())()).
(10)
We defer the technical details of constructing the matching loss of transfer function for a given layer to Appendix [D] 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 (penultimate) leaky ReLU layer before the final softmax layer of ResNet-18 model. Nonetheless, our constouction is applicable to other types of layers, such convolutions. We then consider knowledge distillation using the representations learned from the ResNet-18 teacher model to train smaller convolutional student model on the same dataset. Our knowledge distillation approach is illustrated in Figure [ Our approach consists of learning the dual mean m and the principal directions V in the ε-th layer of the teacher network. The representation yf of a training example æ; is then approximated by d = f' + m) where the compression coefficients C, E R² are predicted by a smaller student network. The approximated representation of can then be passed through the rest of the pre- trained teacher network or 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. 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 50,000 train and 10,000 test images of size 32 x 32 from, respectively, 10 and 100 classes 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 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 leaky ReLU with a small negative slope of 3 = le-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
| In order
Available at hattp:://uuu.co http:://www.ce.ttooooo.eeeee twrontorooooooooooooo www.:oooooooo..tkiii///// wwt.:eeedddkiiiiiaar.t hhml html

---

Published in Transactions Machine Learning Research (02/2023)
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-I0 and CIFAR-100, respectively) of each example. In the following, consider several experiments for compressing each of these representations. We use normalized SGD with exponential moving average (EMA) momentum for optimizing our Bregman PCA problem. At last, we also provide results on the ImageNet- datasets (Deng et al. 2009). Further details about the experiments are provided in Appendix
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 the vanilla PCA (Eq. B in the pre-activations (i.e., logits) domain. Note that PCA not directly applicable to the post-activation probabilitics, as the reconstructed examples may not correspond valid probability distributions. We also consider CoDA PCA (Avalos et al., 2018) by augmenting constant bias of to the coefficient to handle the non-centered case, as suggested in the paper. We 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 Lest. sets as the performance measure.
The average KL divergence values between the original output and the reconstructed probabilities are shown in Figure 3 (b)) As can be seen from the figure, Bregman PCA yields much lower loss compared to vanilla PCA and CoDA PCA. Also, notice that CoDA-PCA becomes inefficient. 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 (i.e., the penultimate layer). The convex conjugate function F3 corresponding to leaky ReLU transfer function fz( max(u, 0) - 3 max( (x, 0) with slope B > 0 amounts to (u) = 1/2 u @ (z) where @ denotes Hadamard product. For this problem, we consider the 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 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 cc)-(d). Our Bregman PCA works significantly better than vanilla PCA for reconstructing the representation and yields much higher classification accuracy at the output. In order to visualize these representations, we illustrate the 2-D TriMap (Amid & Warmuth, 2019) projections of the CIFAR-10 output probabilities as well as the reconstructed representations of the examples at the leaky ReLU layer using our method in Figure X We can see from the figure that. the two representations, although one layer apart, form very similar clusters. Also, as the number of components & increases, the leaky ReLU representations become better separable.
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 [b). The student network consists of a small convolutional neural network with 5 convolutional layers of size [128 (x2),

---

Published in Transactions Machine Learning Research (02/2023)
CIFAR-10- Outpuls
CIFAB-ICO PCA Ontputs
CIFAB-1O- PCA Lanky Bell Rape
CIFAR-100 (00 PCA Laiky Rege
10a
10
10
80
R diveggmce
30
10-2
10
PCA
Pre PCA
CoDA PCA
Post PCA
Bregman PCA
40
Brogmon PCA
10-1
trii
train
10-3
lest
12
50
75
(a)
(b)
(c)
(d)
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 much higher classification accuracy.
256 (x)), 512] followed by a dense linear layer which outputs & coefficients. These coefficients are then passed through dense layer (with weights V and bias 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) Bascline 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) Bascline 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 V and m, which are set to the values obtained by applying our Bregman PCA on the Icaky ReLU representations of the training examples, as in Section 1.2 5) BRL + Readout, which is similar to BRL, but we set the readout layer weights to the teacher weights. Fach model is trained for 51 epochs using a batch size of 128 using SGD with a Nesterov momentun optimizer with a linear decay schedule. In each case, tune the learning rate and momentum hyper- parameters.
To train BRI, 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 Breg- man PCA the teacher's representations (as in the previous section). This essentially reduces the training of the student network from classification to regression problem. We find that this pre-training the tcacher coefficients significantly improves the convergence of the student model. Afler training 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 discntangle 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 (Zhang et al. 2018). ctc.), which are standard for training these models. This resembles a limited sample
10

---

Published in Transactions Machine Learning Research (02/2023)
CIFAR 10 (Full)
CIFAR 10 (Tall)
CIFAR 100 (Full)
CIFAR 100 (1aaf)
80
43
10st SOUSEWKY
10
20
Todios
30
Sadder
Figure 4: Results of training different basclines, including soft-label teachor-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.
regime where the number of training examples is relatively small. We conduct additional ex- periments by incorporating these augmentations and observe improvements for all baselines (sce Appendix F.. In order to highlight the limited samples aspect even further, we also repeat, each experiment using only half the training dataset, which is uniformly sampled from the full dataset (and fixed for all methods).
The results of training the smaller student model shown in Figure We observe significant improvement in accuracy using our BRL distillation approach across different numbers of com- ponents (k). For CIFAR-10, even with small k, BRL shows major advantage over soft-label teacher-student training. For CIFAR-100, these improvements are minor for smaller k values as the leaky ReLU reprosentations require a relatively larger mumber of components to be reconstructed efficiently compared to CIFAR-10 (see Figure F())reB(c)-(d).. (c)). 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 V) for the baseline distillation approach using soft labels (56.5% top-1 accuracy for BRL compared to 54.1% for Bascline Distillation).
Additionally, in Figure g 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 datasct). This shows that BRL 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 the ImageNet- 1k dataset into PreAct ResNet-18 student. Similar to the previous experiments, we replace the transfer function of the penultimate layer with leaky ReLU with a negative slope of 3 = 1e-2. We train all models on an 8 x 8 TPU v3 hardware. The teacher model is trained for 90 epochs using SGD with momentum with a piecowise constant schodule 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 (Hinton et al. 2015) and tune the soft-label ratio and the logit temperature, and train the model for 90
11

---

Published in Transactions Machine Learning Research (02/2023)
Table 1:
ImageNet-1k distillation results using ResNet-50 teacher and ResNet-18 student. Each model is trained for 90 epochs. We also report the number of parameters of cach model and the runtime on the same 8 x 8 TPU v3 hardware.
Model
Top-1 Accuracy
Top-5 Accuracy
#Params
Rumtime
ResNet-50 Teacher
76.2%
93.0%
25.6M
1.7h
ResNet-18 Baseline
70.1%
89.4%
11.7M
1.3h
ResNet-18 Soft-label Distillation
71.5%
90.4%
11.7M
2.4h
ResNet-18 BRT. (Ours)
72.6%
91.1%
15.6M
3.9h
epochs using the same optimization hyperparameters as the ResNet-18 baseline. The soft-label distillation baseline achieves 71.5% top-1 Lest accuracy.
The teacher's penultimate layer has dimension d = 2048, which 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 EMA.) We then import those reprosentations into the student model and keep them fixed. 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 (Kingma & Ba 2014) with 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 bet ween the teacher's s target and the student's '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 Table [ We also report the number of parameters of each model and the runtime on the same hardware. summary, our BRL distillation significantly improves the test accuracy while introducing slight overhead in terms of the number of parameters and the runtime.
5 Conclusion and Future Work
We presented a new direction for knowledge distillation based on directly transferring the com- pressed Bregman representations of the teacher network into student network. Our work signif- icantly 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 intermediale layers. Our construction is more flexible and can be exlended to training a cascade of student networks. Specifically, a teacher network can be split: into non-overlapping sub-networks where each portion is approximated by a smaller student. However, other than the first student, each student receives the output of the previous student as its input, thus, forming a cascade of networks to downsize the teacher model. Such training strategies are interesting future directions for technique.
Limitations and Broader Impact One limitation of our current approach is that it only applies to layers with strictly monotonic transfer functions. The extension of our work to monoooooiic transfer functions is an interesting open problem. In addition, our layerwise representation learning improves the quality of the distilled models. Reducing the size of the model while maintaining its quality decreases the power usage of the model at the serving time. However, our method requires a longer training time compared standard soft-label distillation. Such overhead is only reasonable if the model is utilized for inference for reasonably long period of time.
12

---

Published in Transactions Machine Learning Research (02/2023)
References
Sreangsu Acharyya. Transductive de-noising and dinensionality reduction using total Bregman regression. In Proceedings of the 2006 SIAM International Conference on Data Mining, pp. 514-518. SIAM, 2006.
Shun-ichi Amari. Information geometry and its applications, volume 194. Springer, 2016.
Ehsan Amid and Manfred K. Warmuth. TriMap: Large-scale Dimensionality Reduction Using
Triplets. arXic preprint arXiv:1910 00204, 2019.
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.
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.
Rohan Anil Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George Dahl, and Geoffrey Hinton. Large scale distributed neural network training through online distillation. arXiv preprint or Xrxiv: 1814.1800.038 03235, 2018.
Rohan Anil, Sandra Gadanho, Da Huang, Nijith Jacob, Zhuoshu Li, Dong Lin, Todd Phillips, Cristina Pop, Kevin Regan, Gil Shamir, al. On the factory floor: ML engineering for industrial-scale Ads recommendation models. arXiv preprint arXiv:2209.0533 2022.
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.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: review and new perspectives. IEFE transactions on pattern analysis and machine intelligence, 35(8):1798-1825 2013.
Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov. Knowledge distilaaton:: A good teacher is patient and consistent. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 10925-10934, 2022.
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-25) 217, 1967.
Julien Chiquct, Mahendra Mariadassou, and Stéphane Robin. Variational inference for proba- bilistic Poisson PCA. The Annals of Applied Statistics, 12:2674 2698, 2018. doi: 10.1214/18%
2DA0AS1177.
Michael Collins, Sanjoy Dasgupta, and Robert E. Schapire. A generalization of principal component analysis to the exponential family. Tn Proceedings of the 14th International Conference Neural Information Processing Systems: Natural and Synthetic, NIPS'01, pp. 617-624, 2001.
Wojciech Marian Czarnecki, Simon Osindero, Max Jaderberg, Grzegorz Swirszez, and Razvan Pas- camu. Sobolev training for neural networks. arXiv preprint arxi:: 070..89906,485 2017.
13

---

Published in Transactions Machine Learning Research (02/2023)
J. Deng, W. Dong, R. Socher, L.-J. Li, Li, and Fei-Fei. ImageNet: Large-Scale Hierarchical Image Databasc. In CVPR09, 2009.
Gene H. Golub and Charles F. Van Loan. Matrir Computations. The Johns Hopkins University Press, third edition, 1996.
Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision, 129(6):1789-1819, 2021.
Allen Hatcher. Algebraic topology. Cambridge Univ. Press, Cambridge, 2000. URL https://cos
cern. cern. wwtp:ccccrr/e//00000 478079
D. P. Helmbold, Kivinen, and M. K. Warmuth. Relative loss bounds for single neurons. IEFE Transactions Neural Networks, 10(6):1291 1101991-144,,,, 1999.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint urXix:1503.02591 2015. URL
02531v1
Jean- Baptiste Hiriart-Urruts and Claude Lemaréchal. Fundamentals of Conver Anolysis. Springer- Verlag Berlin Heidelberg, first edition, 2001.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ar preprint arXie:1412 6980, 2014.
Kivinen and M. K. Warmuth. Relative loss bounds for multidimensional regression problems. Journal of Machine Learning, 45(3):301 329, 2001.
Andrew J. Landgraf and Yoonkyung Lee. Generalized principal component analysis: Projection of saturated model parameters. Technometrics, 62:459 472, 2020a.
Andrew J. Landgraf and Yoonkyung Lee. Dimensionality reduction for binary data through the projection of natural parameters. Journal of Multivariate Analysis, 180:104668, 2020b.
John M Lec. Riemannian manifolds: on introduction to curvature, volume 176. Springer Science & Business Media, 2006.
Rafael Müller, Simon Kornblith, and Geoffrey Hinton. Subclass distillation. arXiv preprint arXi::2002.0399 039336, 2020.
Erkki Oja. Simplified neuron model as principal component analyzer. Journal of mathematical biology, 15(3):267 273, 1982.
Nicolas Papernot, Patrick Medaniel, Xi Wu, Somesh Jha, and Ananthram Swami. Distillation as a defense to adversarial perturbations against deep neural networks. 2016 IEEE Symposium on Securily and Privacy (SP), pp. 582 597, 2016.
Karl Pearson. On lines and planes of closest fit to systems of points in space. Philosophical Magazine Series 1, 2:559-5555 2-559-572, 1901.
Ralph Rockafellar. Characterization of the subdifferentials of convex functions. Pacific Journal of Mathematics, 17(3):497-510, 1966.
14

---

Published in Transactions Machine Learning Research (02/2023)
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and
Yoshua Bengio. Fitncts: Hints for thin deep nets. arXiv preprint arXi::14226666 2014.
Nicholas Roy and Geoffrey J Gordon. Exponential family PCA for belief compression in POMDPs.
Advances in Neural Information Processing Systems, 15:1667 1674, 2002.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint 1rXiv:19100010000 2019.
Jonathon Shlens. A lutorial principal component analysis. arXiv preprint orXio: 2014.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive represenlation distillation. arxiv preprint arXiv:1910. 10699, 2019.
Qizhe Xie, Minh- Tian-TTaan Luong, Eduard Hovy, and Quoc V Le. Self-training with noisy student improves imagenct classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 10687 10698, 2020.
Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the perfor- mance of convolutional neural networks via attention transfer. arXiv preprint arXiv:1612.0800 2016.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez- Paz. mixup: Beyond empirical risk minimization. In ICLR, 2018.
15

---

Published in Transactions Machine Learning Research (02/2023)
A
Omitted Proofs
Proposition 1. The generalized dual mean in Eq. 5 can be written
Σ m4)
(6)
Proof. Using the duality property of Bregman divergences, Eq. 55 can be written
min mEdom
Dr-(aa ( f(m))
min
(mm,
medo
Applying Eq. and setting the derivative of above with respect to m to zero, have
(((m))-(((())) = | [(m) Σ x; 0.
Rearranging the term and applying the inverse yields the result.
Theorem 1. Let QR denole the procedure that returns the QR factors. Given M E 84+ and M A). Then, the malrie Q VM- corresponds the first
A Ruxn, let R QR(VMA) factor of the QR decomposition of A the Riemann-Stieff manifold such that A QR and Q MQ = Im-
Proof. Since QR- v A, we have
QR= QR-VM-QR-VMM--AA A,
and Q satisfies
Q MQ VM- Im
B The Case of the Softmax Transfer Function
Consider the case where the input examples are probability distributions {a; E Ad-1) belonging to the (d (d-12 11-smmippplex A = {u E R /u 1. 1}. The transfer function in this case corresponds to Js = softmax which induces the KT. Bregman divergence Dr (u, v) = Ey log 41 Uj + uy- Requiring fam to inevitable imposes the constraint dom fex = Rd {+cld, CE R+} as fex (u + e ela) = fem(u)) (u) for e e R. Thus, for the principal directions, we have 1 É CS(V) where CS denotes column span. The above constraint can be easily incorporated into the QR decomposition on the Riemann-Stiefel manifold in Algorithm I when using the Houscholder method for the internal QR step.
The Householder method applies a series of reflections P. = Im - 2 b. on matrix A such that
b.b
s.t. b, I and e [n [n-1] -n],
R= P ... P.PA
is upper-triangular. The orthonormal matrix Q then can be written as
P.P.p... PP
16

---

Published in Transactions Machine Learning Research (02/2023)
The following proposition shows that, in order to obtain Q Stat from A Q M1 0 using Houscholder reflections, it suffices to augment A from left by column of all ones and apply Algorithm I
The resulting matrix corresponds to the first factor when the first column is
dropped.
Proposition 2. Given E Ruxn s.t. nt < m. Q. RS-QR([1m.. A), M) by applying Al- gorithm using Householder reflections. The Q and R factors, QR A s.t. Q St and Q M1m 0, cun be oblained from Q and R, respectively, by dropping the first column of Q and the first row and column of R.
Proof. The first reflection Pi of the Householder method transforms the first column of the input matrix into a unit vector S1e1 where $1 R. Thus, the first column of the resulting Q term is just a rescaling of the first column the input matrix. Thus, for the augmented matrix [1m. A], the first column of Q corresponds to Im with ri E R s.l. rii, M1. 1. The remaining columns of Q, i.e. the matrix Q, correspond the column space of A minus the direction 1m Also, we have QTMQ function, i.e., softmax(QR) soltmax(A) where the softmax function is applied to each column and R corresponds the upper- matrix by removing the first row and column of R.
MQ In and Q M1m O. Note that the direction 1m is redundant for the softmax
C The Full Bregman PCA with 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 Mean (Including Softmax)
Input: X' (æx dom F Rd), number of components k < d Output: dual mcan dom F, V &t {c; R$}
(m))
m - f Σ, x;) initialize V and {c}} repeat
for E [[]]]] do à; + f(m + Vc,)
+ c Na V
( x;)
V V V Th Σ, c x; c.
until V, [c,c not converged
if f is softmax then
V (1,,
> Augment a column of all ones.
V, T + RS-QR(V, Hr(m))
if f is softmax then
V - V.,. V...
Drop the first column. Drop the first row and column.
T Th:1:
return m, V, {Tc}
17

---

Published in Transactions Machine Learning Research (02/2023)
D
Matching Loss of Transfer Function
Given an (element-wise) striclly-increasing (or cyclically monotone, in the more general case (Rock-
afellar 1966)) transfer function f Id > IRd the matching loss between the (post-)targer y and
the prediction ÿ = f(ā) is defined as the line integral (Helmbold et al. 1999 Kivinen & Warmmuhh
2001
[ly.f(ā)) = (f(()) (f(x)-((a)) f(a)) dz
((nia)))(((f (f(z)
(11)
where a = (z) is the pre-target. Note that the inverse function f-1 always exists since f is strictly increasing. By simplifying the integral, we can write the expression in Eq. 11 in termss of the convex integral function F Rd R s.t. f VF.
F - à = (y, Dr
(12)
Eq. 12) shows the equivalence between the matching loss and the Bregman divergence induced by the convex integral function F (or its Fenchel dual F*). Also, the last expression indicates that T( (a)) is always convex with respect to
Further Details on the Experiments
CIFAR-10/100 Experiments: For each experiment, we tune the learning rate and the momen- tum 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, use V-100 GPUs. For the ImageNet-lk experiments, we use 8 x 8 TPU v3s.
F Experimental Results with Data Augmentation and Mixup
Data augmentation is 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 (Zhang et al., (018)) with B on the CIFAR- 10/100 datasets. Table 2 presents the results of different methods with k 100 components and 51 training epochs with batch size of 128.
From the table, see that the performance of all methods consistently improves with data augmen- lation. 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 sufficient amount of data (i.e., training examples) for distillation is provided (Beyer et al.) 2022). As show in the main experiments, BRL delivers clear advantage in setup with limited number of samples.
18

---

Published in Transactions Machine Learning Research (02/2023)
Table 2: Top-1 accuracy for different methods with number of components A = 100 and using data augmentation combined with Mixup (Zhang et al.
2018).
Dataset
Method
CIFAR-10
CIFAR-100
Baseline
83.65 + 0.41
54.25 + 0.35
Baseline + Readout
83.34 + 0.51
55.11 0.34
Baseline Distillation
86.00 + 0.22
62.25 0.15
BRL
87.03 + 0.12
62.72 0.28
BRL Readout
86.24 + 0.29
61.11 + 0.29
19


Reviews:
[{"review_id": "HeDRxAQvvR", "reviewer": "Reviewer_YJiT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "First, thank you to the authors for pursuing an interesting direction, and bringing interesting mathematical ideas to the problem of model distillation. I see promise in this work and believe that it will of interest to the community. So I am recommending acceptance. However this verdict is very much giving the benefit of the doubt, and based on the TMLR philosophy that \"only some readers need be interested\" in order to justify acceptance. I have a number of concerns with the current manuscript that make it hard for me to determine the significance of the mathematical contribution, and hard to interpret the experimental results. I have attempted to give some critical, but constructive suggestions on how to set about framing this work more clearly. \n\n---\n\n> **Strengths**\n\n- Introduction of a mathematically interesting variation on the Bregman divergence that is also allowed to minimize over a mean vector. \n- The authors apparently closely follow the lines of prior Bregman divergence computation algorithms, making the necessary steps in order to re-derive an algorithm that also accounts for their desired mean vector. \n- Method doesn’t appear to introduce additional hyperarameters above and beyond what were already needed for computing Bregman divergence. \n- Method performs well against simple (yet sensible) baselines like PCA. \n\n---\n\n> **Weaknesses**\n\n- The generalized Bregman PCA is introduced ‘to handle the case of non-centered data’. Clearly non-centered data is important, but what is wrong with just centering the data, computing the usual Bregman PCA, then uncentering again? Discussion of this point, or experiments showing it, would be very valuable to explaining the significance of this work.\n- The contributions on page 3 are not made understandable. They are written as if the reader has already read the paper, as opposed to a the reader still trying to grasp the high-level contours. For instance in the first bullet point: why is it important that you added a generalized mean vector? Its it essential to have in order to formulate your distillation method, or just a nice-to-have? The second bullet: you “improve the orthonormality constraint… to conjugacy” - but what has this improved? Why is this better? And what is this used for? The third: what is QR decomposition and how is it related to the conjugacy constraint? A construction suggestion would be to include a high-level sketch of the algorithm structure (instead of only citing previous works Roy & Gordon and Chiquet).\n\n---\n\n> **General Assessment** \n\nTo conclude, this paper introduces a new Bregman PCA formulation and algorithm to compute it, which appears to be mathematically interesting. It appears to have some promising applications to distillation of neural networks - beating PCA alternatives, and the prototypical Soft-label Distillation method. Giving the benefit of the doubt I am therefore inclined to accept on the basis that this might be useful to some. However I would like to clearly register my main concern: the entire methods section is devoted to developing a ‘generalized’ Bregman PCA algorithm, following similar lines to the original. However nowhere in the manuscript (that I saw) was it explained or motivated why this generalized Bregman PCA was a better thing to be doing than usual Bregman PCA. And there are no experiments comparing generalized vs standard Bregman PCA. And I see no experiments comparing generalized vs standard Bregman. This leaves me quite uncertain as to what the value in the generalized method is, and whether perhaps existing Bregman divergence computation methods would work just as well. So I am quite skeptical, and if other reviewers have doubts about this paper then I would prefer to join a consensus to reject than to champion this papers acceptance. \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes using a variant of PCA—the so-called generalized Bregman PCA—as an approach for projecting neural network representations into lower dimensions. The generalized Bregman PCA method is the same as the original Bregman PCA, with the addition of a optimizable mean vector added to the reconstruction (resulting differences in the algorithm are listed in the introduction). The authors apply this method to distillation of neural networks, finding that their approach outperforms using vanilla PCA, and outperforms the common distillation method of training the student model to match the logits of the teacher model. \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Generally speaking I would highly recommend revising Sec 2 and the latter part of Sec 1 to make the ideas and the new contributions (and their significance) clearer. For example, the following key points could be make much clearer: 1) what is the idea behind your algorithm for computing ‘generalized Bregman divergences’?, assume the reader doesn’t already know how to compute Bregman divergence PCA 2) why is this new generalized Bregman formulation a good idea for distilling neural network representations? What is it giving us above and beyond usual Bregman? Equivalently, why is it important to deal with non-centered data when we could just center data, compute Bregman PCA, then uncenter again? This could also be addressed with additional experiments comparing to this as a baseline.\n- More specifically, remove the list on page 3, as it is not informative to the reader. Replace it instead with a schematic explanation of the new PCA method (projected gradient descent —> gradient computation easy and analytic, but how to project? Why does data lying on a different manifold (e.g., the simplex) require a change to the algorithm, and where is the change made? \n- I hate to be that person asking for extra experiments. But it would be really nice if you could add runtime comparisons for your method versus e.g., usual PCA, and the Soft-label method in the case of distillation. ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gKFw6cvoTJ", "reviewer": "Reviewer_aEXP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nI like the technical cleanness of the approach and of the modification to the QR algorithm. \n\nThe application of PCA to distillation is compelling.\n\nThe explanations of the technical parts are well explained overall.\n\nWeaknesses: \n\nThe necessity of including the mean term in the PCA output is unclear and ill defined to me. In the other generalized PCA cases mentioned, it’s not especially limiting to assume a zero mean. In general, why can’t we just demean the features (or enforce that type of constraint during training)? There’s no space savings here because were storing a mean vector still?\n\nThe overall message of the paper is a bit confused – the experiments are written as if it was a distillation paper (but these experiments are limited), while the methodology is very heavy (as if it was a numerical analysis paper). What else can be done with this non-linear bregman PCA beyond just distillation of neural networks?\n\nWhile not pitched as a “state of the art” paper, it seems very limiting to not compare to other reasonable compression / distillation approaches, for these experiments, e.g. adversarial ones [1] or other approaches [2,3]. I think the baselines of Coda PCA and simple soft label based distillation are a bit too weak for the lack of other experiments because it’s not clear if this is really a practical method for distillation / compression or just an interesting artifact of the math you did.\n\nThere’s no convergence analysis of the gradient descent algorithm for your approach – why should it converge to a stationary point? And, furthermore, the experiments mention an “online” version, why should this converge?\n\nReferences: \n[1] MEAL V2: Boosting vanilla Resnet-50 to 80%+ top-1 accuracy on ImageNet with-\nout tricks., Shen and Savvides, 2020.\n[2] Does knowledge distillation really work?, Stanton et al, Neurips, 2021.\n[3] Knowledge distillation: a good teacher is patient and consistent, Beyer et al, CVPR, 2022.\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an extension of Bregman PCA which estimates non-zero means and projection coefficients and can be extended to include constraints from non-linear activations. They then use Bregman PCA to compress and distill neural nets on the output layers. Their set of experiments demonstrate that the distillation performance is better than a soft-label baseline in terms of accuracy.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major: the construction of QR in section 2.1 definitely should be renamed to something that is not “generalized QR”. “Generalized QR” is a separate technique for separately generating orthogonal matrices (see https://netlib.org/lapack/lug/node46.html for a quick introduction and [1] for more details). I’d suggest renaming this to something like “QR on the stiefel manifold”.\n\nAs a broader point, I believe that there should be some more discussion as to what exactly the transformed part of the QR decomposition is doing here. I think the transform $\\sqrt{M} A$ is a retraction onto the manifold defined by $M$ but state this with low confidence. [2] may be helpful here.\n\nMinor: spell out how you’d apply the constraints in either the more general case (rather than just softmax) and to a QR approach that isn’t simply householder (e.g. with Givens rotations, or realistically in the general case). \n\nMinor: the paper would benefit from a definiation of the “convex integral function of the transfer function” – on pg 7.\n\nMedium: while I quickly verified the algebra in Prop. 1, I’m a bit unclear about the intuition as to why the solution to the mean vector should just be the inverse map applied to the arithmetic mean? This also seems to assume function invertibility (which I see why there’s the restriction to monotonic transfer functions).\n\nmedium: Would we expect to see good results with transfer learning across datasets, or just good results from distillation?\n\nIdeally, there’d be some discussion of “why” there’s better results on knowledge distillation? Is the optimization process made easier from this type of compression or is it mapping the data somehow? It seems like this intuition is a bit missing currently and should be explored a bit more, see [3,4] for some interesting metrics and experiments that can be studied.\n\nTypos, etc:\n-\tPg 6: change “decent” to descent\n-\tPg 16: please just use two lines instead of the brackets in the proof of prop 1.\n-\tPg 8: alg 1. “until … not converged”. I think you mean “while … not converged”\n-\tPg 8: “we essentially minimize the matching loss of the transfer function”  unclear, are you doing this exactly or not?\n\nReferences:\n[1] Generalized QR Factorization and its Applications, Anderson, Bai, and Dongarra, Linear Algebra and Its Applications, 1992.\n\n[2] Cholesky Based QR retraction on the generalized Stiefel manifold, Sato and Aihara, computational Optimization and its Applications, 2019.\n\n[3] Does knowledge distillation really work?, Stanton et al, Neurips, 2021.\n\n[4] Knowledge distillation: a good teacher is patient and consistent, Beyer et al, CVPR, 2022.\n\n", "broader_impact_concerns": "Actually, some discussion of broader impacts, aka, “why do we care about this work?” would be helpful.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "r72AQ3wmoM", "reviewer": "Reviewer_CEhv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Methodologically interesting, might be applicable more broadly (I would appreciate discussion on this though).\n* Clear technical exposition and background discussion. I found that starting with an explicit and relatively verbose discussion of PCA effective in making the paper quite self-contained and easy to follow. I feel quite confident that implementing the method just based on the paper would be fairly straight-forward.\n* The method appears to consistently perform better than the distillation baseline. There are large-scale results on ImageNet.\n\nWeaknesses:\n* No clearly stated claims. Essentially this paper is written like a conference submission (\"we propose a new method and it works better than the baselines under certain metrics and datasets\") rather than working out explicit claims and supporting them with evidence. E.g. \"Our empirical findings indicate that our approach is substantially more effective for transferring information...\" as in the abstract is extremely vague. I found this quite surprising as this is one of the two central acceptance criteria for this venue, although I will say that this is also my first time reviewing for TMLR, so I'm open to down-weighing this concern.\n* Lack of high-level discussion. The introduction is essentially a technical background section, and provides no context on why this research is relevant/timely, what its overarching goal is or how the proposed approach is motivated.\n* Related to the previous point, I find the positioning of the paper a bit odd. The title and abstract strongly emphasize the methodological contribution and argue that distillation is only an example application of the method, however there is no discussion of further applications (let alone demonstrations of these). \n* The experiments do not discuss parameter counts at all, i.e. size of student network + $|m|, |V|$. My hunch is that for the proposed method they don't differ too much from classical knowledge distillation, however I haven't calculated this precisely -- these should really be provided in the paper. The storage complexity increases linearly in $k$, so I think these should be stated explicitly along all results (Figs 3 & 4, Tab 1). Given that (loosely speaking) the goal of distilling a larger network into a smaller one is to reduce storage/compute complexity under acceptable accuracy loss, I find it crucial to show some related metrics, parameter count would be the most obvious one.\n\nMinor suggestions:\n* In Fig 4, mark the distillation baseline with a horizontal line.\n* I found the paragraph with the different variants at the bottom of page 10 a bit difficult to parse. Stating the hypothesis that you're testing with each approach/what ablation this corresponds to in (relatively) plain English would probably make it easier to read.\n* I don't really follow the motivation for the cascade of student networks to distill each layer independently. Do you have reason to believe this would lead to an overall smaller approximation error while still being parameter efficient? I would like to see a little bit more detailed discussion on this and ideally a concrete experiment (can be exploratory, e.g. distilling individual layers vs entire ResNet blocks and comparing errors).\n* Typo in Table 1, the first Accuracy column should say \"Top-1 Accuracy\" instead of \"Top-5 Accuracy\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a PCA algorithm with Bregman divergences under a generalized mean, which it utilizes to propose Bregman Representation Learning, a distillation method for approximating a large neural network with a smaller one. It demonstrates that this approach results in better predictive accuracies on image classification benchmarks (CIFAR, Imagenet) compared to some baselines, most notably knowledge distillation.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical:\n* Work out clear and explicit claims and discuss how they are supported by the theoretical and empirical results.\n* Provide parameter counts for all methods as discussed in weaknesses.\n\nOptional:\n* I would suggest either editing this paper to be a more clearly focused neural distillation paper or adding discussion (ideally alongside empirical demonstrations) of other applications of the proposed method if the authors prefer to maintain their current general perspective.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose an extension of Bregman PCA, incorporating a mean vector and revising the normalization constraint on the principal directions. They then use Bregman PCA to compress and distill neural networks on the output layers.

The paper has a clear technical exposition and background discussion, as well as compelling experimental results on large-scale data. 

Initial concerns by the reviewers about the clarity and scope of the work have been successfully addressed in the review period, which resulted in an extensive and constructive dialogue between the authors and reviewers. The revised paper now addresses technical details (e.g., highlighting the importance of information transfer between models as opposed to achieving SOTA distillation performance and other details, such as discussing the role of the mean vector for non-centered PCA).

Overall, this is a very interesting paper that will find its readership. I, therefore, recommend acceptance. 
Decision: Accept (poster)