Efficient Maximal Coding Rate Reduction by Variational Forms

The original MCR² objective mcr takes the following form: Given $m$ training samples $\mathbf{X}=[\mathbf{X}_{1},\ldots,\mathbf{X}_{m}]\in\mathbb{R}^{D\times m}$ belonging to $k$ classes¹¹ 1 Here we adopt the notation that an upper case letter $\mathbf{X}$ represents a matrix and $\mathbf{X}_{i}$ denotes the $i^{\text{th}}$ column of a matrix., let $\mathbf{Z}_{\theta}=[f_{\theta}(\mathbf{X}_{1}),...,f_{\theta}(\mathbf{X}_{m})]\in\mathbb{R}^{d\times m}$ be the latent representation where recall $f_{\theta}$ is the featurizer parameterized by $\theta$, and let $\mathbf{\Pi}\in\mathbb{R}^{m\times k}$ define the class membership, where $\mathbf{\Pi}_{i,j}$ denotes the probability²² 2 Note that if the labels are known exactly then the entries of $\mathbf{\Pi}$ are binary $\{0,1\}$ with each row of $\mathbf{\Pi}$ summing to one. Notice that our notation of $\mathbf{\Pi}$ is slightly different from that adopted in mcr. Our choice is more compact for optimization purposes. that $\mathbf{X}_{i}$ is in class $j$. Then, $\text{MCR}^{2}$ aims to learn a feature representation $\mathbf{Z}_{\theta}$ that maximizes the following coding rate reduction $\Delta R(\mathbf{Z}_{\theta})$:

where $\mathbf{\Pi}_{j}$ denotes the $j^{\text{th}}$ column of $\mathbf{\Pi}$, $\mathrm{Diag}(\mathbf{\Pi}_{j})$ denotes a diagonal matrix with $\mathbf{\Pi}_{j}$ along the diagonal, $\alpha=d/(m\epsilon^{2})$, $\alpha_{j}=d/(\langle\mathbf{1},\mathbf{\Pi}_{j}\rangle\epsilon^{2})$, $\gamma_{j}=\langle\mathbf{1},\mathbf{\Pi}_{j}\rangle/m$, $\mathcal{S}$ is the set of all matrices whose columns all have unit $\ell_{2}$ norm³³ 3 Note the constraint that $\mathbf{Z}_{\theta}$ has unit norm columns is often achieved by simply having the final operation of the network $f_{\theta}$ be a normalization. and $\epsilon>0$ is a prescribed precision error. Roughly speaking, $R(\mathbf{Z}_{\theta})$, known as the expansion term, captures the dimension (or the volume) of the space spanned by $\mathbf{Z}_{\theta}$ while $R_{c}(\mathbf{Z}_{\theta},\mathbf{\Pi})$, or the compression term, measures the sum of the dimensions/volumes of the data from each class. From an information-theoretic point of view, $R(\mathbf{Z}_{\theta})$ estimates the coding rate, or the number of binary bits required to encode $\mathbf{Z}_{\theta}$, through $\epsilon$-ball packing ma2007segmentation. The terms are called expansion and compression terms respectively, since by maximizing $\Delta R$, the first coding rate term is maximized, which seeks to expand the overall volume of the embedded features, while the second coding rate term is minimized, which seeks to compress the volume of the embedded features from each class.

By assessing the MCR² objective (1), one can already observe a potential drawback of MCR² for optimization. In particular, note that each $\log\det$ term requires $O(\min\{d^{3},m^{3}\})$ operations to compute (and similarly to back-propagate through). While $d\ll D$ can often be made reasonably small for many high-dimensional datasets which have an underlying low-dimensional structure, $R_{c}(\mathbf{Z}_{\theta},\mathbf{\Pi})$ in particular is still often expensive to compute because it requires $k$ computations of $\log\det$. This severely limits the application of MCR² for datasets with even moderate numbers, say hundreds, of classes as the objective becomes computationally infeasible on common machines.

To avoid this computational bottleneck, here we propose instead a formulation which takes advantage of variational forms of spectral functions of a matrix. Specifically, for a given positive semi-definite (PSD) matrix $\mathbf{M}$ and any scalar $c\geq 0$, note that

where $r$ denotes the rank of $\mathbf{M}$ and $\sigma_{i}(\mathbf{M})$ is the $i^{th}$ singular value of $\mathbf{M}$. Here note that $\log(1+c\sigma)$ is a non-decreasing, concave function of $\sigma$, so we can exploit known variational forms of spectral functionsgiampouras2020novel; ornhag2020bilinear. In particular note the following result in ornhag2020bilinear.

Using Proposition 3.1, we develop our formulation by replacing the $-\log\det$ terms in $R_{c}$ of (1) by the above variational form. In particular, note that for each class $j$ we can eliminate the associated $-\log\det$ term in $R_{c}$ of (1) by introducing an additional matrix $\mathbf{U}^{(j)}$ subject to the constraint that $\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}=\mathbf{U}^{(j)}(\mathbf{U}^{(j)})^{\top}$. Further, due to the fact that each row of $\mathbf{\Pi}$ sums to one, we also have $\sum_{j=1}^{k}\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{T}=\mathbf{Z}_{\theta}\mathbf{Z}_{\theta}^{\top}$. As such, the variational form in Proposition 3.1 gives that the original MCR² objective in (1) is equivalent to the following constrained variational form:

From this form, we now reparameterize the $\mathbf{U}^{(j)}$ matrices as $\mathbf{U}^{(j)}=\Gamma\mathrm{Diag}(\mathbf{A}_{j})^{1/2}$ where $\Gamma\in\mathbb{R}^{d\times q}\cap\mathcal{S}$ is a dictionary with unit norm columns, and $\mathbf{A}_{j}\in\mathbb{R}_{+}^{q}$ is a (non-negative) encoding vector. Now let $\mathbf{A}\in\mathbb{R}_{+}^{q\times k}=[\mathbf{A}_{1},\ldots,\mathbf{A}_{k}]$ be a matrix of the concatenated encoding vectors and note that we trivially have $\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}=\mathbf{U}^{(j)}(\mathbf{U}^{j})^{\top}$ and $\|\mathbf{U}^{(j)}_{i}\|_{2}^{2}=\mathbf{A}_{i,j}$, which gives another equivalent formulation for the MCR² objective, provided the number of dictionary elements $q$ is sufficiently large so that each optimal $\mathbf{U}^{(j)}$ matrix in (5) can be encoded by $\Gamma$ (i.e., each column of $\mathbf{U}^{(j)}$ must be a column of $\Gamma$ within a scaling factor):

Finally, we relax the strict equality constraints $\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}=\mathbf{Z}_{\theta}\textrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$ with $\ell_{2}$ penalties $\frac{1}{\gamma_{j}}\|\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}-\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\|_{F}^{2}$ to arrive at our final proposed formulation, which we call V-MCR²:

such that $\mathbf{Z}_{\theta}\in\mathcal{S}$. Regularization parameter $\mu>0$ weights how strictly the equality constraints should be approximated, and the $\frac{1}{\gamma_{j}}$ terms roughly ensure class balance (recall, $\gamma_{j}=\langle\mathbf{1},\mathbf{\Pi}_{j}\rangle/m$). From this reformulation, we have significantly reduced the complexity of evaluating the objective function. The $\log\det$ terms in $R_{c}$ which take $O(k\min\{d^{3},m^{3}\})$ time to evaluate is now replaced by $O(qk)$ for the sum of the $\log(1+\alpha_{j}\mathbf{A}_{l,j})$ terms along with the cost of computing the $M$ term which scales as $O(kd^{2})$.

Besides the above computational advantages, we also discuss a few additional aspects of our formulation below.

For fair comparison across training objectives (original MCR², variational MCR²), we use a learning rate of $10^{-3}$ for the network optimizer and the same batch size. For CE, we use the same batch size, but a larger learning rate of $10^{-2}$. The batch size is 1000 for MNIST and CIFAR-10, and 2000 for CIFAR-100 and Tiny ImageNet. The network is optimized using stochastic gradient descent for all objectives. For the precision error $\epsilon$ of the MCR² objectives, we use $\epsilon^{2}=0.5$ for all datasets. $\epsilon$ and batch size are consistent with the experimental settings in the original MCR² work mcr for MNIST and CIFAR-10. For Variational MCR², the regularization constant $\mu=1$ and initial learning rates $\nu_{\Gamma}=5$, $\nu_{\mathbf{A}}=5$ across all experiments. We perform latching every 50 epochs for all experiments. See Appendix for precise details. The dictionary size $q$ and feature dimension $d$ varies across the datasets. For MNIST and CIFAR-10 we use $d=128$ and $q=20\cdot k$, and for CIFAR-100 and Tiny ImageNet we use $d=500$ and $q=10\cdot k$.

Across datasets, we compare the wall-clock time to train one epoch using the MCR² (1) and V-MCR² (7) formulations. The batch size is set to be the same for both models, and all our experiments are performed using PyTorch 1.9.0 and Python 3.8.11 on Nvidia A100-SXM4 GPUs with 40GB of CUDA memory for fair comparison. As shown in Table 1, V-MCR² training completes approximately $5\times$ faster on CIFAR-100 and $12\times$ faster on Tiny ImageNet. Even for datasets containing a small number of classes, i.e., MNIST and CIFAR-10, we observe a $1.5-2\times$ speedup. Note that the overhead for the original MCR² model escalates significantly as the number of classes increases, so we expect even greater improvements in training efficiency with datasets with more classes. Even for Tiny ImageNet, training until convergence using MCR² becomes nearly impractical, while it is easily handled by V-MCR².

Additionally, we compare the true $\Delta R(\mathbf{Z}_{\theta})$ (i.e., computing the original MCR² objective with the iterates of the V-MCR² model) over training epochs and observe that both models 1) take approximately the same number of training epochs to converge and 2) reach approximately the same final $\Delta R(\mathbf{Z}_{\theta})$ objective value at convergence. Thus, V-MCR² does not require additional epochs to obtain a good solution which might offset the increased efficiency per epoch. As shown in Figure 1, on MNIST and CIFAR-10, V-MCR² and MCR² follow a similar training loss trajectory across epochs. For CIFAR-100, we observe that the convergence rate depends on the dimension of the features/representations, $d$. For a feature dimension of $d=100$ for CIFAR-100, we observed similar number of epochs to convergence for MCR² and V-MCR². With a feature dimension of $d=500$, we observe across 5 seeds that training with the original MCR² leads to a rapid convergence to a poor local optima where the expansion term $R(\mathbf{Z})$ increases rapidly but the compression term $R_{c}(\mathbf{Z})$ remains the same. On the other hand, although V-MCR² requires more epochs to converge, we observe the standard expected behavior where $R(\mathbf{Z})$ increases and $R^{c}(\mathbf{Z})$ decreases and the final solution is of much higher quality (see 5.2). Experiments on Tiny ImageNet show similar behaviors.

There are two properties of the representation that we aim to attain by optimizing the MCR² objective. We say that a representation of the training data is of ‘high quality’ if points from different classes lie on separate, orthogonal subspaces, and the union of these subspaces span as many dimensions as possible. In particular, the orthogonal property is important in order to classify the points using the nearest subspace algorithm in Section 4.2. To check the orthogonality of subspaces learned by $f_{\theta}$, we report the inner product between every pair of training points as a heatmap in Figure 2. Namely, we sort the columns of $\mathbf{Z}_{\theta}$ by class and compute $|\mathbf{Z}_{\theta}^{\top}\mathbf{Z}_{\theta}|$. Ideally, we want $|\mathbf{Z}_{\theta}^{\top}\mathbf{Z}_{\theta}|$ to have a block diagonal structure, with $(\mathbf{Z}_{\theta})_{i}^{\top}(\mathbf{Z}_{\theta})_{j}\approx 0$ for $i,j$ notating points from different classes.

Figure 2 shows heatmaps for MNIST, CIFAR-10, and CIFAR-100 after 2000 training epochs, when $\Delta R$ has converged for both MCR² and V-MCR² (Figure 1). For MNIST and CIFAR-10, note that both MCR² and V-MCR² obtain a block diagonal structure. However, on CIFAR-100 and Tiny Imagenet, we observe no block diagonal structure after MCR² training, whereas we see a clear block diagonal structure for V-MCR². These findings suggest that V-MCR² training is more robust to avoiding poor local minima than training on the original MCR² model, particularly as the number of classes increase. We leave a rigorous study of these phenomena for a future work.

In Table 2, we present test accuracies on the four aforementioned datasets when trained under original MCR² and V-MCR² objectives. In addition, we also train a separate model for each dataset by using CE loss as a reference classifier and report its test accuracy. Notice that the goal of our study here is not about achieving the best possible classification accuracy on these datasets -- the training procedure and architectures used here are not optimal for that purpose⁷⁷ 7 We use simple training practices (e.g. input data downsampling, minimal data augmentation and training from scratch), leading to about 15% and 30% difference in performance on CIFAR-100 and Tiny ImageNet as reported in yun2020regularizing from much more carefully engineered training recipes. . Instead, we make fair comparison of all methods on the same networks and datasets to justify the computational efficiency and effectiveness of the proposed method. To ensure fairness, we initialize this reference model with the same architecture along with other hyperparameters⁸⁸ 8 For experiments using CE, we use a higher learning rate of $10^{-2}$ to improve convergence for CE training. as in V-MCR² experiments and attach a final linear classifier with output dimension corresponding to the number of classes.

As shown in Table 2, when trained on datasets with a small number of classes, all three training objectives can reach competitive classification performance. We can observe that both MCR² and V-MCR² objectives are comparable to CE in these small-scale datasets. Again on CIFAR-100 we observe that the poor local minima obtained by training the original MCR² objective results in a poor test accuracy, while the V-MCR² model achieves comparable performance to training the same network with CE. Due to limited resources and heavy computational requirements of MCR² training with a large number of classes, we can only report the result on Tiny ImageNet by optimizing MCR² after 200 epochs while the other two objectives complete the full training session of 2000 epochs.

To ensure stability of the gradient-based method, we scale the learning rate of $\Gamma$ and $\mathbf{A}$ by (the inverse of) approximate upper-bounds of the Lipschitz constants of the gradients, $\frac{1}{L_{\Gamma}}$ and $\frac{1}{L_{\mathbf{A}}}$, respectively. To simplify calculations, we only bound the Lipschitz constant of the gradients with respect to the matrix approximation term $M$, as we note this term typically dominates the Lipschitz constant of the gradient. Additionally, for notational simplicity we show the derivation for balanced classes which implies $\frac{1}{m\gamma_{j}}=\frac{k}{m}$.

Now, consider the function $f(\Gamma,\mathbf{A})=\frac{\mu k}{2m}\sum_{j=1}^{k}\left\|\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi})_{j}\mathbf{Z}^{\top}-\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\right\|_{F}^{2}$, and note that with some simple algebra one can show the following equivalence:

where $(\cdot)^{\odot 2}$ denotes raising to the second power entry-wise.

From this, the relevant gradients are

We first bound the Lipschitz constant of $\nabla_{\Gamma}f$. For $g_{j}$, we have

where the inequality is due to the sub-multiplicative property of the Frobenius norm and the operator norm inequality. Now we consider $h_{j}$. Here we make the simplifying assumption that in addition to each column of $\Gamma$ being a unit vector, $\Gamma$ also approximately satisfies $\Gamma^{\top}\Gamma=\mathbf{I}$, which is known to be true near the globally optimal solution from the analysis of mcr. Now, we have the following approximation

which implies,

Since $f(\Gamma,\mathbf{A})$ has a summation over class $j$, we get the following approximate upper bound for the Lipschitz constant

For the upper bound of the Lipschitz constant of $\nabla_{\mathbf{A}}f$ we simply have

We utilize the following architectures for the experiments in Section 5. We use a fairly simple architecture for MNIST, and for the other datasets, we use a slightly modified version of ResNet18. Note d is the feature dimension.

For CIFAR-10/100 and Tiny Imagenet, we use the Torchvision ResNet18 model as the featurizer, but we remove the final layer of the ResNet18 and replace it with the following to reshape the output into the desired feature dimension $d$. We also normalize the output at the end to fulfill the constraint that the features lie on the unit sphere in the MCR² objective.

For cross-entropy experiments, we add another linear layer on top to map the output of the featurizer to logits.

We utilize the following data augmentations for the experiments in Section 5.

The experiments conducted and presented in the main body, Table 2, were only meant to compare the computational efficiency for different methods under the same conditions. The settings however were not chosen to optimize the classification performance since we did not conduct data augmentation or other training recipes normally adopted, e.g. see yun2020regularizing.

Here we report experimental results on CIFAR-100 and Tiny ImageNet by training with a similar training recipe adopted in yun2020regularizing. Specifically, as in yun2020regularizing, we use PreAct ResNet-18 as the backbone and train with both cross-entropy loss and V-MCR². All networks are trained by SGD with momentum 0.9, and weight decay of $10^{-4}$. Similar to yun2020regularizing, we set the learning rate to be $10^{-1}$ for the first half of training epochs, divide it by $10$ for the next quarter of epochs, and finally divide it again by $10$ for the remaining iterations. Additionally, we utilize the same transformations to augment the data. Note that the only difference we make between our strategy and the one in yun2020regularizing is the choice of batch size and total number of training epochs (both specified in Section 4.3) to ensure fair comparison between CE and V-MCR². In addition to the training strategy, we set $\mu=10^{-1}$ for V-MCR² as we find a smaller value for $\mu$ helps stabilize training in the early stage. We keep all other hyperparameters identical to ones specified in Section 4.3.