# Active learning for efficient annotation in precision agriculture: a use-
case on crop-weed semantic segmentation

Bart M. van Marrewijk Wageningen University & Research, Wageningen, The
Netherlands  Corresponding author. Email<EMAIL_ADDRESS>Charbel
Dandjinou  Exxact Robotics, Épernay, France  Dan Jeric Arcega Rustia
Wageningen University & Research, Wageningen, The Netherlands  Nicolas Franco
Gonzalez  Exxact Robotics, Épernay, France  Boubacar Diallo  Exxact Robotics,
Épernay, France  Jérôme Dias  Exxact Robotics, Épernay, France  Paul Melki
Exxact Robotics, Épernay, France  IMS, CNRS, University of Bordeaux, France
Pieter M. Blok  Wageningen University & Research, Wageningen, The Netherlands

###### Abstract

Optimizing deep learning models requires large amounts of annotated images, a
process that is both time-intensive and costly. Especially for semantic
segmentation models in which every pixel must be annotated. A potential
strategy to mitigate annotation effort is active learning. Active learning
facilitates the identification and selection of the most informative images
from a large unlabelled pool. The underlying premise is that these selected
images can improve the model’s performance faster than random selection to
reduce annotation effort. While active learning has demonstrated promising
results on benchmark datasets like Cityscapes, its performance in the
agricultural domain remains largely unexplored. This study addresses this
research gap by conducting a comparative study of three active learning-based
acquisition functions: Bayesian Active Learning by Disagreement (BALD),
stochastic-based BALD (PowerBALD), and Random. The acquisition functions were
tested on two agricultural datasets: Sugarbeet and Corn-Weed, both containing
three semantic classes: background, crop and weed. Our results indicated that
active learning, especially PowerBALD, yields a higher performance than Random
sampling on both datasets. But due to the relatively large standard
deviations, the differences observed were minimal; this was partly caused by
high image redundancy and imbalanced classes. Specifically, more than 89% of
the pixels belonged to the background class on both datasets. The absence of
significant results on both datasets indicates that further research is
required for applying active learning on agricultural datasets, especially if
they contain a high-class imbalance and redundant images. Recommendations and
insights are provided in this paper to potentially resolve such issues.

##  1 Introduction

Deep learning has received much attention in the agricultural domain due to
its potential for resolving complex computing tasks [1]. Within the deep
learning subdomain, semantic segmentation stands out for its pivotal role in
precisely distinguishing class-specific pixels within agricultural imagery
[2]. Semantic segmentation can be used for all kinds of precision-related crop
monitoring tasks and targeted crop interventions. However, to achieve accurate
performance, the segmentation model must be trained on a significant number of
annotated images, a process that is time-consuming and expensive [3]. To
address this challenge, active learning has emerged as a promising approach to
optimize the annotation process by actively selecting the most informative
images for annotation [4]. Selection of those informative images is typically
done by uncertainty-based acquisition functions, which acquire images about
which the segmentation model is most uncertain. By selecting these uncertain
images and using them to retrain the segmentation model, performance can be
significantly improved with minimal annotation effort. Thus, active learning
is a promising solution to reduce annotation effort and improve efficiency in
precision farming applications [5, 6].

One of the remaining challenges in active learning is the design and use of
uncertainty-based acquisition functions. A crucial part in designing
acquisition functions is the concept of aleatoric and epistemic uncertainty.
Aleatoric uncertainty refers to the inherent randomness or variability in the
data itself [7]. It includes the uncertainty that cannot be reduced even if we
had access to more annotated images or improved models. Epistemic uncertainty,
on the other hand, refers to the uncertainty of the model arising from a lack
of knowledge or limited data [7]. Epistemic uncertainty can be reduced with
more annotated images and it is therefore crucial in active learning because
it provides essential information to select the most informative images.
Acquisition functions that incorporate epistemic uncertainty are therefore
preferred over functions that incorporate aleatoric uncertainty. One of the
acquisition functions that only work with epistemic uncertainty is Bayesian
Active Learning by Disagreement (BALD) [8]. Other hybrid functions of BALD,
such as PowerBALD, combines the concept of epistemic uncertainty and
stochastic sampling [9]. BALD, and to a lesser extent PowerBALD, have been
applied to well-known public datasets such as Cityscapes, but to the best of
our knowledge, these methods have not been tested for semantic segmentation on
agricultural datasets. As these datasets are often not as diverse and balanced
as other public datasets, the question is whether uncertainty-based active
learning is still able to improve annotation efficiency on agricultural
datasets.

The aim of this paper is to bridge this research gap through a comparative
study between BALD, PowerBALD and Random acquisition functions on agricultural
image datasets. Our goal is to provide insights into the challenges and
opportunities of active learning by testing different acquisition functions on
these datasets. We have conducted such experiments on two crop-weed datasets,
one of which is publicly available [10] and the other was used in a practical
production setup. The contributions of our work include: 1.) Detailed analysis
of how active learning, especially uncertainty sampling, works on agricultural
datasets; 2.) Identification of issues arising from applying active learning;
and 3.) Potential solutions to resolve such issues.

##  2 Active learning

The proposed active learning framework is based on uncertainty sampling (BALD)
and stochastic uncertainty sampling (PowerBALD). Both methods select images
the semantic segmentation model is most uncertain about, with the assumption
that these images will contribute most to the performance when retraining the
segmentation model. In this work, the level of uncertainty on each unlabeled
image was determined with Monte-Carlo dropout, a frequently used sampling
technique in active learning [11]. With Monte-Carlo dropout, unlabeled images
were repeatedly passed through the model in inference mode with dropout
activated, thus producing multiple segmentations for the same image. If the
outputs for the same image show large variation, this may indicate that the
model is uncertain about this image, making it a candidate for selection and
annotation. This approach is presented in more detail in the sections below.
First, an introduction to the active learning framework is presented in
Section 2.1, then the semantic segmentation model used in this work is
explained Section 2.2. In Section 2.3 the different uncertainty acquisition
functions are presented.

###  2.1 Active learning framework

Our active learning framework was developed using the Bayesian Active Learning
(Baal) software, which has been specially developed for active learning [12].
The pseudo-code is given in Algorithm 1. In the text below, the main
components of our active learning framework are described in italics.

Our active learning process starts with an initial dataset randomly sampled
from the training pool. The training pool is the dataset with unlabeled images
from which the active learning framework selects the most uncertain images. A
semantic segmentation model is trained on this initially sampled dataset.
After training, the model is used to calculate the uncertainty of each image
in the training pool, using one of the acquisition functions, which are
described in Section 2.3. Depending on the number of Monte-Carlo iterations,
each image is x times inferred through the model with dropout enabled. After
each iteration, the predicted model output for every pixel may be different,
since the model is forced to make its decision with a potentially different
set of model weights due to the dropout. The dropout intensity was made
adjustable through the dropout probability, whose value could be set between 0
(no dropout) and 1 (complete dropout). From the multiple model predictions,
the uncertainty is calculated based on the acquisition function. This outputs
a list of the most informative images ready to be annotated. After annotating
these images the segmentation model was trained again on the combined set of
initial images and most informative images. This loop was repeated depending
on the number of iterations. This process iteratively increases the number of
annotated images and thereby the performance of the segmentation model.

Inputs : acquisitionFunction : BALD, PowerBALD, Random
numberOfIterations : integer
sampleSize : integer
initialDatasetSize : integer
dropoutProbability : float
numberOfMonteCarloIterations : integer
trainingPool : dataset with all available images for training and sampling
valSet : dataset with images for validation

Outputs : mIoU : semantic segmentation performance

Function  _ActiveLearning(_acquisitionFunction , numberOfIterations,
sampleSize, initialDatasetSize, dropoutProbability,
numberOfMonteCarloIterations, trainingPool, valSet_)_:

mIoU $\leftarrow$ [ ] ;

initialDataset $\leftarrow$ Random(_trainingPool , initialDatasetSize_);

annotatedTrainSet $\leftarrow$ Annotate(_initialDataset_);

availableImages $\leftarrow$ trainingPool $-$ annotatedTrainSet ;

model $\leftarrow$ TrainModel(_annotatedTrainSet_);

mIoU.insert( EvaluateModel(_valSet_)) ;

for  _$i\leftarrow 1$ to numberOfIterations_ do

pool $\leftarrow$ Sampling(_model , availableImages, acquisitionFunction,
numberOfMonteCarloIterations, sampleSize, dropoutProbability_);

annotatedPool $\leftarrow$ Annotate(_pool_);

annotatedTrainSet $\leftarrow$ annotatedTrainSet $+$ annotatedPool ;

model $\leftarrow$ TrainModel(_annotatedTrainSet_);

mIoU.insert( EvaluateModel(_valSet_)) ;

availableImages $\leftarrow$ availableImages $-$ annotatedPool ;

return mIoU

Algorithm 1 Active learning

###  2.2 Segmentation model

Our active learning framework could potentially work with any segmentation
model as long as there is a dropout layer added to that model. In our
experiments, the Fully Convolutional HarDNet (FCHarDNet) was used as semantic
segmentation model. FCHarDNet is a semantic segmentation model that employs a
harmonic densely connected network [13], of which its architecture is shown in
Figure 1. Compared to other segmentation models, FCHardNet has a very low
memory traffic, allowing it to have a very low inference time, which is
relevant for potential real-time segmentation of agricultural images
containing crops and weeds. To perform active learning, a dropout (DO) layer
was added before the final convolutional layer of the FCHarDNet architecture
(Figure 1). By adding this dropout layer, Monte-Carlo iterations could be
performed. If an image is relatively new to the currently trained model, we
expect the pixel uncertainty to be relatively high when applying multiple
Monte-Carlo iterations. We hypothesized that images with high model
uncertainty contributed the most to the segmentation performance when
retraining FCHarDNet.

Figure 1: FCHarDNet network architecture [13] with dropout (DO). The numbers
in the figure indicate the number of strides for a given block.

###  2.3 Acquisition functions

After obtaining multiple model predictions with Monte-Carlo dropout, the
predictions were summarized by a single value that represents the model
uncertainty on that image. In our research, three different acquisition
functions were tested: BALD[8], PowerBALD [9], and Random acquisition.

  * •

BALD: BALD was calculated using Equation 1, which was obtained from [4, 8].
The first term of BALD, the marginal entropy $H(y|x,D)$, triggers on images
with high entropy. These are images that contain large fractions of pixels
with distributed prediction probabilities across all classes (indicating that
the model is not able to make a confident prediction on a single class per
pixel). The second term of BALD, the conditional uncertainty
$E_{p(\omega|D)}[H(y|x,\omega,D)]$, penalizes images with aleatoric
uncertainty. Because of BALD’s two terms, we were able to get better estimates
of epistemic uncertainty rather than aleatoric uncertainty.

${BALD(y;\omega|x,D)=H(y|x,D)-E_{p(\omega|D)}[H(y|x,\omega,D)]}$ (1)

  * •

PowerBALD: A limitation of using BALD is that it can sample a set of uncertain
images with large visual similarities between them. This can especially occur
with agricultural datasets that are obtained with slow-moving platforms,
because then there can be a high visual similarity between the image at time t
and the image at time t+1. Instead of relying on the single-most uncertain
images from the training pool when using BALD, we also used PowerBALD [9].
PowerBALD uses stochastic sampling to sample a more diverse set of images from
the uncertainty distribution of BALD. PowerBALD uses a temperature factor T to
conduct this sampling (Equation 2). In our experiments, T was set to 1.0.

${PowerBALD(y;\omega|x,D)=\frac{BALD(y;\omega|x,D)^{\frac{1}{T}}}{\sum
BALD(y;\omega|x,D)}}$ (2)

  * •

Random acquisition: This approach selects images randomly from the training
pool. It is used as the default method for comparison with BALD and PowerBALD.

##  3 Datasets

In this research the active learning framework was evaluated on two
agricultural datasets: ”Sugarbeet” and ”Corn-Weed”. These datasets consisted
of images which can be segmented pixel-wise according to three classes:
background, weed, and crop (sugarbeet for Sugarbeet, corn for Corn-Weed). The
Sugarbeet dataset is a publicly available dataset [10] that was collected from
a single field at multiple acquisition moments during three months. An example
image and its ground truth annotation is shown in Figures 2a and 2b. The Corn-
Weed dataset is a private dataset. This dataset is from a production setup
representing images from corn fields with weeds. It consisted of images
captured at a slight angle as visible in Figures 3a and 3b. In Table 1, a
summary of both datasets is given. Both datasets are heavily imbalanced.
Especially the Sugarbeet dataset has a minimal amount of diversity and high
image redundancy with 98.5% of the pixels belonging to the background class.
The Corn-Weed is also imbalanced, but it has more pixels related to crop class
(6.2%). Each dataset represents distinctive challenges and conditions
corresponding to a variety of agricultural landscapes to which the semantic
segmentation model was applied in real-world scenarios.

Dataset | Image resolution | Train | Validation | Test | Background % | Crop % | Weed %
---|---|---|---|---|---|---|---
Sugarbeet | 1296 × 966 | 9287 | 1155 | 1155 | 98.5% | 1.3% | 0.2%
Corn-Weed | 1710 x 2045 | 331 | 117 |  | 89.8% | 6.2% | 4.0%

Table 1: Active learning dataset statistical summary

In a production setup, annotated data is often already available. This
available data can be used to pre-train a model beforehand to improve
uncertainty estimation and active learning sampling (detailed information can
be found in Section 4). Therefore, in this research we have included two
additional datasets for pre-training: ”PhenoBench” and ”Corn-Weed pre-train”.
The PhenoBench dataset is a publicly available dataset [14] with sugarbeets
and weeds. Images were acquired on multiple days in two different years,
resulting in a dataset with different growth stages and lighting conditions.
An example image and its annotation is shown in Figures 2c and 2d. Corn-Weed
pre-train is a diverse dataset consisting of 1190 images from 17 fields spread
over 3 countries (Figure 3c and 3d). Compared with the Corn-Weed dataset the
acquisition angle is slightly different. In Table 2 a summary of both pre-
trained datasets are given. In particular, the Phenobench dataset stands out,
with 11.8% of the pixels annotated as crop.

Dataset | Image resolution | Train | Validation | Test | Background % | Crop % | Weed %
---|---|---|---|---|---|---|---
PhenoBench | 1024 x 1024 | 1407 | 772 |  | 87.7% | 11.8% | 0.5%
Corn-Weed pre-train | 1710 x 2045 | 1190 |  |  | 90.9% | 5.4 % | 3.7 %

Table 2: Pre-training dataset statistical summary

(a)

(b)

(c)

(d)

Figure 2: Example image from the Sugarbeet dataset with (a) input image and
(b) corresponding annotation, where red=sugarbeet, blue=weed (only a few
pixels) and transparent=background. Similar visualization for the PhenoBench
dataset (c, d).

(a)

(b)

(c)

(d)

Figure 3: Example image from the Corn-Weed dataset with (a) input image and
(b) corresponding annotation, where red=corn, purple=weed (only a few pixels)
and transparent=background. Similar visualization for the Corn-Weed pre-
trained dataset (c, d).

##  4 Experiments

This section describes the experiments conducted on the selected datasets. In
all experiments, the active learning process was evaluated based on mean
Intersection-over-Union (mIoU). mIoU is a metric for the number of overlapping
pixels between the predicted segmentation masks and the ground truth
annotations. In this work, the mIoU value represents the average mIoU across
all classes (crop, weed, and background).

###  4.1 Effect of model pre-training on the active learning performance on
SugarBeet

Active learning is based on the concept of sampling images about which the
model is most uncertain. The concept of uncertainty depends on the knowledge
of the model. In the initialization phase of active learning, this knowledge
is usually brought in through transfer learning, a process of transferring the
weights of the same model trained on a public dataset, such as Cityscapes.
These large datasets contain many general classes, such as human and car, but
no agricultural classes, such as crop and weed. In this work, a hypothesis is
that the uncertainty determination and active learning significantly benefit
when transfer learning is done on agricultural images. To test this, the first
experiment was tested under two scenarios:

  1. 1.

Active learning sampling starting with the standard Cityscapes weights. We
refer to this in the remainder of this paper as ”without agricultural pre-
training”.

  2. 2.

Active learning sampling starting with the weights of a model pre-trained on
the PhenoBench (train) dataset [14]. We refer to this in the remainder of this
paper as ”with agricultural pre-training”.

For a fair comparison between the two scenarios, the adjustable parameters in
active learning were kept the same. A sample size of 10 images was chosen,
with number of iterations set to 9; this means that training and sampling were
stopped upon reaching 90 images. The dropout probability was set to 0.5. In
this experiment, three replicate trials were performed for each acquisition
function to account for the randomness in the Monte-Carlo dropout method. A
one-way ANOVA with a significance level of 5% was employed to test whether
there were significant performance differences between BALD, PowerBALD and
Random. Furthermore, an in-depth analysis was made by making a histogram of
the uncertainty values from BALD on the models without and with agricultural
pre-training for both the Sugarbeet and Phenobench-validation dataset. If BALD
is calculating the uncertainty values accordingly, then images on the
Sugarbeet dataset should have a higher uncertainty than the Phenobench-
validation images when using agricultural pre-training. The in-depth analysis
also contains a visualization of pixel-level uncertainty values for a sample
image. Finally, the active learning performance was compared to the benchmark
performance by training the segmentation model on the entire training pool
(9287 images). The resulting mIoU value was considered as the maximum mIoU
that could be achieved on our datasets.

###  4.2 Effect of model pre-training on the active learning performance on
Corn-Weed

In this experiment, active learning performance was evaluated in an industrial
application. Often in agriculture, a model is trained on a set of images
obtained from multiple fields. When deploying the model in practice, new
fields will be encountered, meaning that the model is exposed to different
conditions than the ones learned during model training. This may lead to
suboptimal model performance in practice. To overcome this, it is necessary to
select and annotate images from new fields, but for humans it is unclear which
images should be selected. In this experiment, the goal is to verify the
effectiveness of active learning when the segmentation model has already been
pre-trained on a dataset with the same camera setup but with different field
conditions. To test the added value of the pre-training, the following
scenarios were tested:

  1. 1.

Active learning sampling starting with the standard Cityscapes weights. Again,
we refer to this as ”without agricultural pre-training”.

  2. 2.

Active learning sampling starting with the weights of a model pre-trained on
the Corn-Weed pre-train dataset of Table 2. Again, we refer to this as ”with
agricultural pre-training”.

In this experiment, the training pool was composed of 331 images from the
Corn-Weed dataset (Table 1). The models were tested on the Corn-Weed
validation dataset, which consisted of 117 images (Table 1). All active
learning settings were similar to Section 4.1. The benchmark test was done by
training the segmentation model on all 331 images.

##  5 Results

###  5.1 Effect of model pre-training on the active learning performance on
SugarBeet

The experimental results of testing the active learning framework on the
Sugarbeet dataset are shown in Figure 4(a) (without agricultural pre-training)
and Figure 4(b) (with agricultural pre-training). The overall performance of
the segmentation model with agricultural pre-training was better compared to
the performance without agricultural pre-training. Without agricultural pre-
training, the Random acquisition function significantly outperformed
(p$\leq$0.05) BALD and PowerBALD at 50 sampled images (see the two black dots
in Figure 4(a)). For the other iterations, there were no significant
differences between the three acquisition functions. With agricultural pre-
training, there were no significant differences observed (Figure 4(b)).
Despite the lack of significant differences, PowerBALD had a consistently
better performance than Random sampling. Furthermore, it should be noted that
by only sampling 40 images, the agriculturally pre-trained model had a
performance of approximately 0.80 mIoU, which was only 0.07 mIoU lower than
the maximum performance of the model trained on all images (black dashed
line).

(a) Sugarbeet without agricultural pre-training

(b) Sugarbeet pre-trained on PhenoBench

Figure 4: Test performances (mIoU) as a function of the number of training
images for BALD, PowerBALD, and Random on the Sugarbeet dataset. Subfigure (a)
highlights the results without agricultural pre-training, and subfigure (b)
with agricultural pre-training. In both subfigures, the solid colored lines
are the mean performance values over the three repetitions. The colored areas
around the lines represent the 95% confidence intervals around the means. The
black dashed lines indicate the performance when FCHarDNet was trained on the
entire training pool.

The two blue histograms in Figures 5a and 5b show the uncertainty values of
all images of the training pool when deploying BALD on the segmentation models
trained without and with agricultural pre-training. The histogram without pre-
training (Figure 5a) had lower and less scattered uncertainty values than the
one with pre-training (Figure 5b), indicating that BALD was less
discriminative in finding uncertain images. This explains the lower
performances of both BALD and PowerBALD in the without agricultural pre-
training scenario (Figure 4(a)). The orange histograms in Figures 5a and 5b
visualize the uncertainty values on the Phenobench-validation dataset. In
Figure 5a, the blue and orange histograms are centered around 0.02, in other
words, the images in both validation sets have roughly the same uncertainty.
After agricultural pre-training (Figure 5b), the blue histogram of the
SugarBeet dataset has shifted to the right, indicating that images from the
SugarBeet dataset got a higher uncertainty value assigned. This is according
to our expectation that BALD assigns higher uncertainty values to images that
are relatively unknown to the pre-trained model. Figure 5c shows the
uncertainty values visualized on a pixel-level for an example image of the
SugarBeet dataset. The green- and yellow-colored pixels have a higher
uncertainty than the blue-colored pixels. Figure 5c clearly indicates that the
model uncertainty was high on the plant pixels.

(a)

(b)

(c)

Figure 5: Histogram of the uncertainty values when deploying BALD on a model
trained (a) without agricultural pre-training and (b) with agricultural pre-
training. In both (a) and (b), the blue histogram depicts the uncertainty
values calculated on the 9287 images of the SugarBeet dataset. The orange
histogram depicts the uncertainty values calculated on the 772 images of the
Phenobench validation dataset. (c) Uncertainty values visualized on a pixel-
level in an example image.

###  5.2 Effect of model pre-training on the active learning performance on
Corn-Weed

On the Corn-Weed dataset, we observed marginal differences between the three
acquisition functions. Without agricultural pre-training (Figure 6(a)), the
performances were very similar and not significantly different. With
agricultural pre-training (Figure 6(b)), there were also no significant
differences between PowerBALD, BALD, and Random. This was mainly due to the
high standard deviations between the repetitions of the Random acquisition
(see the blue areas in Figure 6(b)). Interestingly, PowerBALD had the lowest
variance between the three repetitions (Figure 6(b)). This indicates that
PowerBALD had a more consistent performance than BALD and Random on the Corn-
Weed dataset. This result is promising when it comes to applying active
learning in a production setup. BALD was on average underperforming compared
to Random and PowerBALD, meaning that selecting images with only the highest
uncertainty did not lead to better model performance on the unseen corn field.

(a) Corn-Weed without agricultural pre-training

(b) Corn-Weed with agricultural pre-training

Figure 6: Test performances (mIoU) as a function of the number of training
images for BALD, PowerBALD, and Random on the Corn-Weed dataset. Subfigure (a)
highlights the results without agricultural pre-training, and subfigure (b)
with agricultural pre-training. In both subfigures, the solid colored lines
are the mean performance values over the three repetitions. The colored areas
around the lines represent the 95% confidence intervals around the means. The
black dashed lines indicate the performance when FCHarDNet was trained on the
entire training pool.

##  6 Discussion

In this research, the active learning framework was tested on two agricultural
datasets with three acquisition functions: BALD, PowerBALD and Random. On the
Sugarbeet dataset, both BALD and PowerBALD did not show a significantly better
performance than Random acquisition. There are two possible explanations for
the minimal difference between the acquisition functions. First, there was a
high image redundancy in the Sugarbeet dataset. Training the model on the
complete training pool (9287 images) resulted in a mIoU of 0.87. In Figure
4(b), it was shown that after 4 iterations (40 sampled images), the model had
a mIoU of approximately 0.80. This indicates that by training on only 40
images, a performance can be obtained that is close to the performance of
training on 9287 images. This implies that there was not much image diversity
in the dataset.

The second reason for the minimal difference is that BALD’s epistemic
uncertainty was calculated across all semantic classes. The fact that the
background class represented more than 98% of the pixels means that BALD
probably placed more emphasis on optimizing the background class than on the
crop and weed classes. This is also why active learning has been successfully
applied to instance segmentation in agriculture [6], because there the
uncertainty calculation was conducted on the instance level, ignoring the
background pixels. To solve the problem of over-present but less relevant
classes in semantic segmentation, it may be worthwhile to calculate the sum of
uncertainties for the top-K most uncertain pixels. By focusing on the
uncertainty sum of a small part of the image, the influence of the background
class can be reduced. This would work for the image shown in Figure 5(c),
because all pixels related to the plant had a high uncertainty. Another
possibility is to test other acquisition functions, such as Entropy. In a
recently published article, it was shown that Entropy was effective for
quantifying uncertainty on the Sugarbeet dataset [15]. In their work, high
pixel uncertainties were observed around the edges of the plants. However, it
should be noted that in their research the algorithm was pre-trained on 75% of
the dataset. It would be interesting to see if uncertain pixels are observed
when testing on a larger set of unseen images.

The results on the Corn-Weed dataset showed that PowerBALD was better than
Random sampling (Section 5.2). The performance of BALD was worse than Random
acquisition. To verify this outcome, additional experiments were done on
Cityscapes. In these experiments both PowerBALD and BALD were outperforming
Random, indicating that these acquisition functions can be used for active
learning. The high redundancy in our agricultural datasets probably made it
more difficult for the active learning to outperform the Random acquisition.
In our experiments, selection was performed on images that differed from those
used in pre-training the model. For example, images in the Corn-Weed dataset
had fewer shadows and a more uniform sandy soil than the images in Corn-Weed
pre-trained dataset (Figure 3a versus 3c). Due to the difference in images
between the two datasets, each randomly sampled image is still relatively
high-informative for the segmentation network. In the situation that the
training pool also contained images similar to the pre-trained dataset, then
there would be a higher probability of selecting non-informative images when
using Random acquisition. In that scenario, we would expect larger performance
differences between the Random acquisition and the uncertainty-based
acquisition functions. Figure 5(b) highlights this scenario because the blue
histogram had higher uncertainties than the orange histogram. Unfortunately,
in our experiments with the relatively small sample sizes and iterations,
there was no clear performance difference between the active learning and the
Random acquisition. To investigate the added value of active learning in a
future research, it is recommended to do additional testing with different
numbers of sample sizes, iterations and by including more similar or redundant
images in the training pool. The latter would be a better representation of
real agricultural datasets, especially when those are obtained with slow-
moving platforms.

In future work, our active learning framework could also be improved. Our
proposed framework now uses Monte-Carlo dropout, which requires the same image
to be analyzed multiple times to be able to calculate the model uncertainty.
The disadvantage of Monte-Carlo dropout is that it takes significantly more
time to conduct the image analysis compared to a standard image analysis.
Therefore, Monte-Carlo dropout is not preferred for larger datasets, because
of its limited scalability [16]. In future work, other potentially more time-
efficient sampling methods can be tested, including the learning loss method
[17]. With learning loss, an additional regression module is added to the
neural network architecture to predict the loss of unlabeled images. The
advantage of this method is that it can predict the image uncertainty in one
forward pass, thus drastically reducing the sampling time compared to Monte-
Carlo dropout sampling. Next to that, active learning was evaluated by
comparing it to Random sampling. In industrial applications, pre-selection of
fields or locations is often done by humans. In the research of [18], human
annotators decided which images should be annotated. Since humans are often
involved in selection tasks, it may be interesting to compare the active
learning with human sampling as well.

##  7 Conclusions

Active learning focuses on optimizing model performance with fewer image
annotations. In this work, PowerBALD was shown to outperform Random
acquisition when the segmentation model was pre-trained on an agricultural
dataset. PowerBALD also showed less variation in performance across the three
repetitions than Random acquisition, indicating better sampling stability.
Although the better trend, there were no significant differences observed
(p$\leq$0.05). The lack of other significant differences can be explained by
the fact that the datasets were highly unbalanced with more than 89% of the
pixels belonging to the less relevant background class. Future research should
focus on sampling the top-K most uncertain pixels to address the class
imbalance at image level.

## Acknowledgments

We would like to thank the following people from Wageningen University and
Research, who helped us in our research: Ard Nieuwenhuizen and Trim Bresilla.

## References

  * [1] Andreas Kamilaris and Francesc X. Prenafeta-Boldú  “Deep learning in agriculture: A survey”  In _Computers and Electronics in Agriculture_ 147, 2018, pp. 70–90  DOI: https://doi.org/10.1016/j.compag.2018.02.016
  * [2] Zifei Luo et al.  “Semantic segmentation of agricultural images: A survey”  In _Information Processing in Agriculture_ , 2023  DOI: https://doi.org/10.1016/j.inpa.2023.02.001
  * [3] Jiajia Li et al.  “Label-Efficient Learning in Agriculture: A Comprehensive Review”, 2023  arXiv:2305.14691 [cs.CV]
  * [4] Yarin Gal, Riashat Islam and Zoubin Ghahramani  “Deep Bayesian Active Learning with Image Data”  In _Proceedings of the 34th International Conference on Machine Learning - Volume 70_ , ICML’17  Sydney, NSW, Australia: JMLR.org, 2017, pp. 1183–1192
  * [5] Usman A. Zahidi and Grzegorz Cielniak  “Active Learning for Crop-Weed Discrimination by Image Classification from Convolutional Neural Network’s Feature Pyramid Levels”  In _Computer Vision Systems: 13th International Conference, ICVS 2021, Virtual Event, September 22-24, 2021, Proceedings_ Berlin, Heidelberg: Springer-Verlag, 2021, pp. 245–257  DOI: 10.1007/978-3-030-87156-7˙20
  * [6] Pieter M Blok et al.  “Active learning with MaskAL reduces annotation effort for training Mask R-CNN on a broccoli dataset with visually similar classes”  In _Computers and Electronics in Agriculture_ 197 Elsevier, 2022, pp. 106917
  * [7] Laurent Valentin Jospin et al.  “Hands-On Bayesian Neural Networks—A Tutorial for Deep Learning Users”  In _IEEE Computational Intelligence Magazine_ 17.2 Institute of ElectricalElectronics Engineers (IEEE), 2022, pp. 29–48  DOI: 10.1109/mci.2022.3155327
  * [8] Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani and Máté Lengyel  “Bayesian active learning for classification and preference learning”  In _arXiv preprint arXiv:1112.5745_ , 2011
  * [9] Andreas Kirsch et al.  “Stochastic Batch Acquisition for Deep Active Learning”, 2022  arXiv:2106.12059 [cs.LG]
  * [10] Nived Chebrolu et al.  “Agricultural robot dataset for plant classification, localization and mapping on sugar beet fields”  In _The International Journal of Robotics Research_ , 2017  DOI: 10.1177/0278364917720510
  * [11] Yarin Gal and Zoubin Ghahramani  “Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning”  In _Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48_ , ICML’16  New York, NY, USA: JMLR.org, 2016, pp. 1050–1059
  * [12] Parmida Atighehchian et al.  “Baal, a bayesian active learning library”, https://github.com/baal-org/baal/, 2022
  * [13] Ping Chao et al.  “HarDNet: A Low Memory Traffic Network”  In _Proceedings of the IEEE International Conference on Computer Vision (ICCV)_ , 2019, pp. 3551–3560  DOI: 10.1109/ICCV.2019.00365
  * [14] Jan Weyler et al.  “PhenoBench–A Large Dataset and Benchmarks for Semantic Image Interpretation in the Agricultural Domain”  In _arXiv preprint arXiv:2306.04557_ , 2023
  * [15] Ekin Celikkan, Mohammadmehdi Saberioon, Martin Herold and Nadja Klein  “Semantic Segmentation of Crops and Weeds with Probabilistic Modeling and Uncertainty Quantification”  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2023, pp. 582–592
  * [16] Ozan Sener and Silvio Savarese  “Active learning for convolutional neural networks: A core-set approach”  In _arXiv preprint arXiv:1708.00489_ , 2017
  * [17] Donggeun Yoo and In So Kweon  “Learning Loss for Active Learning”, 2019  arXiv:1905.03677 [cs.CV]
  * [18] Akshay L. Chandra et al.  “Active learning with point supervision for cost-effective panicle detection in cereal crops”  In _Plant Methods_ 16.1, 2020, pp. 34  DOI: 10.1186/s13007-020-00575-8


# Efficient Maximal Coding Rate Reduction by Variational Forms

Christina Baek∗
Carnegie Mellon University
<EMAIL_ADDRESS>*equal contribution
Ziyang Wu∗
International Digital Economy Academy
<EMAIL_ADDRESS>*equal contribution Kwan Ho Ryan Chan
Johns Hopkins University
<EMAIL_ADDRESS>Tianjiao Ding
Johns Hopkins University
<EMAIL_ADDRESS>Yi Ma
University of California, Berkeley
<EMAIL_ADDRESS>Benjamin D. Haeffele
Johns Hopkins University
<EMAIL_ADDRESS>
###### Abstract

The principle of Maximal Coding Rate Reduction (MCR2) has recently been
proposed as a training objective for learning discriminative low-dimensional
structures intrinsic to high-dimensional data to allow for more robust
training than standard approaches, such as cross-entropy minimization.
However, despite the advantages that have been shown for MCR2 training, MCR2
suffers from a significant computational cost due to the need to evaluate and
differentiate a significant number of log-determinant terms that grows
linearly with the number of classes. By taking advantage of variational forms
of spectral functions of a matrix, we reformulate the MCR2 objective to a form
that can scale significantly without compromising training accuracy.
Experiments in image classification demonstrate that our proposed formulation
results in a significant speed up over optimizing the original MCR2 objective
directly and often results in higher quality learned representations. Further,
our approach may be of independent interest in other models that require
computation of log-determinant forms, such as in system identification or
normalizing flow models.

##  1 Introduction

Given a classification task, deep networks aim to learn a nonlinear mapping,
consisting of a series of linear and nonlinear functions, that can map data to
their correct labels. The overall deep network can often be interpreted as a
composition of a nonlinear “featurizer” $f_{\theta}$ and a linear classifier
$g(\mathbf{z})=\mathbf{W}\mathbf{z}$ for some matrix $\mathbf{W}$. The hidden
layers or the featurizer is designated with learning a latent representation
$\mathbf{z}_{\theta}=f_{\theta}(\mathbf{x})\in\mathbb{R}^{d}$ that best
facilitates the final layer or classifier for the downstream task.

The canonical way to train a deep learning model for a classification task is
empirical risk minimization using cross-entropy (CE) loss. While CE measures
the difference between the model’s prediction and the true labels, it does not
explicitly enforce any structure over the representation. In fact, Papyan,
Han, and Donoho [16, 7] show that this direct label fitting implicitly leads
to neural collapse in deep networks. That is, as CE loss converges to $0$, the
representations of each class at the last hidden layer collapse to a single
point, suppressing within-class variability.

Beyond neural collapse and failing to represent within-class variation,
several works [19, 4, 18] have empirically shown that training neural networks
using stochastic gradient descent (SGD) on CE loss often leads the network to
utilize the simplest, often spurious, feature in the image for classification.
This hypothesis is theoretically supported by [1] which verified that when
multiple explanations can describe a class, models trained with CE often pick
a subset of features that can classify a majority of the points well and then
classifies the remaining points from noise in the data.

To alleviate this issue, Yu et al. [21] proposed a framework for learning
geometrically meaningful representations, via a featurizer $f_{\theta}$, by
maximizing the coding rate reduction (MCR2). In brief, the MCR2 objective
encourages the latent representation of the entire training set to expand or
occupy as much volume as possible, while simultaneously pushing each class to
compress or occupy as little space as possible. Empirically and theoretically,
it is shown that this objective drives the latent representations of each
class to lie on a low dimensional linear subspace, with the subspaces
orthogonal to each other, which empirically provides robustness against label
noise, a notable advantage of MCR2 compared to CE [21].

However, despite these inherent advantages, the MCR2 metric can be costly from
a computational perspective. In particular, the loss involves calculating the
$\log\det$ of the Gram matrix of the representations of each class. Not only
does the number of $\log\det$ terms grow linearly with the number of classes,
but computing (and back-propagating) the $\log\det$ of a $d\times d$ matrix
incurs a computational cost of $O(d^{3})$. For this reason, MCR2 methods to
date have been limited to datasets with a relatively small number of classes
such as MNIST and CIFAR-10, where the loss is computationally feasible. In
order to make MCR2 scalable, there is a significant need to improve the
computational efficiency, particularly as it pertains to computing $\log\det$
terms, to allow for large numbers (hundreds or thousands) of classes in high
dimensional spaces.

Contributions. In this paper we make significant progress towards this goal,
with the following contributions:

  1. 1.

We provide an alternative formulation of the MCR2 objective based on a
variational form of the $\log\det$ function which scales much more gracefully
with the number of classes and the problem dimension.

  2. 2.

We show experimentally that the alternative formulation requires approximately
the same number of epochs to converge as the original MCR2 formulation, but
achieves a significant speedup in the training cost per epoch, particularly as
the number of classes in the dataset grows.

  3. 3.

As an additional benefit of our formulation, we also observe empirically that
training over our proposed variational formulation often results in higher
quality learned latent representations and better test accuracy than the
original MCR2 objective.

Finally, we note that our approach for optimization with variational forms may
be of independent interest for other models which require computing $\log\det$
terms, such as in system identification [3] or normalizing flow models [8].

##  2 Preliminaries

Here we first describe the original MCR2 formulation as well as introduce
relevant background material.

###  2.1 MCR2 Objective

The original MCR2 objective [21] 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$ classes111Here 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
probability222Note 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 [21]. 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})$:

$\displaystyle\max_{\theta}\Delta R(\mathbf{Z}_{\theta})$ $\displaystyle\equiv
R(\mathbf{Z}_{\theta})-R_{c}(\mathbf{Z}_{\theta},\mathbf{\Pi})\quad\text{s.t.
}\mathbf{Z}_{\theta}\in\mathcal{S},\quad\text{where}$ (1) $\displaystyle
R(\mathbf{Z}_{\theta})$
$\displaystyle=\frac{1}{2}\log\det\left(\mathbf{I}+\alpha\mathbf{Z}_{\theta}\mathbf{Z}_{\theta}^{\top}\right),\quad\text{and}$
$\displaystyle R_{c}(\mathbf{Z}_{\theta},\mathbf{\Pi})$
$\displaystyle=\sum_{j=1}^{k}\frac{\gamma_{j}}{2}\log\det\left(\mathbf{I}+\alpha_{j}\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}\right)$

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}$ norm333Note 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 [12]. 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 MCR2 objective (1), one can already observe a potential
drawback of MCR2 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 MCR2 for datasets with even moderate
numbers, say hundreds, of classes as the objective becomes computationally
infeasible on common machines.

###  2.2 Variational Forms of Spectral Functions

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

$\log\det(\mathbf{I}+c\mathbf{M})=\sum_{i=1}^{r}\log(1+c\sigma_{i}(\mathbf{M})),$
(2)

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 functions[5, 15]. In particular note the
following result in [15].

######  Theorem 2.1 (Adapted from [15])

For any matrix $\mathbf{X}$, let r denote the rank of $\mathbf{X}$, let
$\sigma_{i}(X)$ denote the $i^{th}$ singular value of $\mathbf{X}$, and define

$H(\mathbf{X})=\sum_{i=1}^{r}h(\sigma_{i}(\mathbf{X})).$

for some function $h$. If $h$ is a concave, non-decreasing function on
$[0,\infty)$ with $h(0)=0$, then the following holds

$H(\mathbf{X})=\min_{\mathbf{U},\mathbf{V}:\mathbf{U}\mathbf{V}^{\top}=\mathbf{X}}\sum_{i}h\big{(}\left\|\mathbf{U}_{i}\right\|_{2}\left\|\mathbf{V}_{i}\right\|_{2}\big{)},$

where $(\mathbf{U}_{i},\mathbf{V}_{i})$ denotes the $i^{th}$ columns of
$(\mathbf{U},\mathbf{V})$. Note also that $(\mathbf{U},\mathbf{V})$ can have
an arbitrary number of columns $(\geq r)$ provided
$\mathbf{U}\mathbf{V}^{\top}=\mathbf{X}$.

##  3 Proposed Formulation

Having introduced the above background material, we now describe our proposed
approach. In particular, note that Theorem 2.1 immediately gives the following
result as a proposition.

###### Proposition 3.1

Let $\mathbf{M}$ be any real positive semi-definite matrix and let $c\geq 0$
be any non-negative scalar. Then the following holds:

${-\log\det(\mathbf{I}+c\mathbf{M})=\max_{\begin{subarray}{c}\mathbf{U}:\
\mathbf{U}\mathbf{U}^{\top}=\mathbf{M}\end{subarray}}\\!\\!-\\!\\!\sum_{i}\log\big{(}1+c\left\|\mathbf{U}_{i}\right\|_{2}^{2}\big{)}.}$
(3)

Further, if $\bar{\mathbf{U}}\mathbf{S}\bar{\mathbf{U}}^{\top}=\mathbf{M}$ is
a SVD of $\mathbf{M}$ then $\mathbf{U}^{*}=\bar{\mathbf{U}}\mathbf{S}^{1/2}$
is a solution to the above problem.

Proof. First recall the basic fact that for any function $\psi(x)$ one has
$-\min_{x}\psi(x)=\max_{x}-\psi(x)$. Additionally, recall (2) and note that
the function $h(x)=\log(1+cx)$ satisfies the conditions required for $h$ in
Theorem 2.1. These facts give the following:

$\displaystyle-$ $\displaystyle\log\det(\mathbf{I}+c\mathbf{M})=$ (4)
$\displaystyle\max_{\begin{subarray}{c}\mathbf{U},\mathbf{V}:\mathbf{U}\mathbf{V}^{\top}=\mathbf{M}\end{subarray}}\\!\\!-\\!\\!\sum_{i}\log(1+c\left\|\mathbf{U}_{i}\right\|_{2}\left\|\mathbf{V}_{i}\right\|_{2}).$

Further, note (4) implies that
$-\log\det(\mathbf{I}+c\mathbf{M})\geq-\sum_{i}\log(1+c\left\|\mathbf{U}_{i}\right\|_{2}^{2})$
for all $\mathbf{U}$ such that $\mathbf{U}\mathbf{U}^{\top}=\mathbf{M}$ since
we have simply added the constraint $\mathbf{U}=\mathbf{V}$. The result is
completed by noting that for the choice of $\mathbf{U}=\mathbf{U}^{*}$ the
maximum can be attained since
$\log(1+c\left\|\mathbf{U}^{*}_{i}\right\|_{2}^{2})=\log(1+c\sigma_{i}(\mathbf{M})),\forall
i\in[r]$.

###  3.1 Variational Formulation of MCR2

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 MCR2
objective in (1) is equivalent to the following constrained variational form:

$\displaystyle\max_{\theta}\Delta R(\mathbf{Z}_{\theta})=$ (5)
$\displaystyle\max_{\theta,\\{\mathbf{U}^{(j)}\\}_{j=1}^{k}}\frac{1}{2}\log\det\left(\mathbf{I}+\alpha\sum_{j=1}^{k}\mathbf{U}^{(j)}(\mathbf{U}^{(j)})^{\top}\right)$
$\displaystyle\quad\quad\quad\quad-\sum_{j=1}^{k}\frac{\gamma_{j}}{2}\sum_{i}\log\left(1+\alpha_{j}\left\|\mathbf{U}_{i}^{(j)}\right\|_{2}^{2}\right)$
$\displaystyle\quad\text{s.t. }\forall
j,\mathbf{U}^{(j)}(\mathbf{U}^{(j)})^{\top}=\mathbf{Z}_{\theta}\textrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}\
\text{and}\ \mathbf{Z}_{\theta}\in\mathcal{S}.$

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 MCR2 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):

$\displaystyle\max_{\theta}\Delta R(\mathbf{Z}_{\theta})=$ (6)
$\displaystyle\max_{\theta,\Gamma\in\mathbb{R}^{d\times
q}\cap\mathcal{S},\mathbf{A}\in\mathbb{R}_{+}^{q\times
k}}\frac{1}{2}\log\det\left(\mathbf{I}+\alpha\sum_{j=1}^{k}\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\right)$
$\displaystyle\quad\quad\quad-\sum_{j=1}^{k}\frac{\gamma_{j}}{2}\sum_{l=1}^{q}\log\left(1+\alpha_{j}\mathbf{A}_{l,j}\right)$
$\displaystyle\text{s.t. }\forall
j,\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}=\mathbf{Z}_{\theta}\textrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}\
\text{and}\ \mathbf{Z}_{\theta}\in\mathcal{S}.$

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-MCR2:

$\displaystyle\max_{\theta,\Gamma\in\mathbb{R}^{d\times
q}\cap\mathcal{S},\mathbf{A}\in\mathbb{R}_{+}^{q\times
k}}R^{v}(\Gamma,\mathbf{A})-R^{v}_{c}(\mathbf{A})-\frac{\mu}{2m}M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$
(7) $\displaystyle\text{ where}\quad
R^{v}(\Gamma,\mathbf{A})=\frac{1}{2}\log\det\left(\mathbf{I}+\alpha\sum_{j=1}^{k}\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\right),$
$\displaystyle
R^{v}_{c}(\mathbf{A})=\sum_{j=1}^{k}\frac{\gamma_{j}}{2}\sum_{l=1}^{q}\log\left(1+\alpha_{j}\mathbf{A}_{l,j}\right),$
$\displaystyle
M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})=\sum_{j=1}^{k}\frac{1}{\gamma_{j}}\left\|\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}\\!-\\!\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\right\|_{F}^{2},$

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})$.

###  3.2 Interpretation of the Variational Form

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

Sparsifying dictionary learning interpretation. Notice that the above
variational reformulation takes on a natural interpretation as learning a
sparsifying dictionary: it essentially “parameterizes” the subspaces spanned
by each class with a common shared dictionary $\Gamma$. Every class then
selects a “sparse” number of eigenbases, with $\mathbf{A}_{j}$, from this
dictionary and forms its estimate of the sample covariance within the
subspace. Notice that the scalar $\log$ terms in $R^{v}_{c}(\mathbf{A})$ are
precisely nonconvex sparsity promoting measures adopted in early studies of
sparse representation [13, 14]. The sparsity in the (spectral) bases in terms
of $\mathbf{A}_{j}$ precisely corresponds to the subspace spanned by each
class being low-dimensional or low-rank.

Penalty function method and other options. Notice that in our formulation (7),
the equality constraint in (6) is enforced through a penalty function $M$. As
the penalty weight $\mu$ increases to infinity and the dictionary $\Gamma$ is
sufficiently large444In the worst case the model become equivalent when
$\Gamma$ is large enough to contain a concatenation of the singular vectors of
each class
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$.,
the formulation becomes exactly equivalent to the original formulations (6)
and (1). Of course, to deal with the equality constraint in (6) more
precisely, one may also consider adopting more advanced methods such as the
augmented Lagrangian multiplier method to incorporate the equality constraint
[6, 17], which we leave for future work. However, as we discuss next, by
relaxing the strict equality constraint we also gain a potential advantage
when the latent representation contains noise.

Low-rank LASSO interpretation. Notice that the sparse/low-rank promoting term
$R^{v}_{c}(\mathbf{A})$ and the quadratic penalty term $M$ together resemble
the classic LASSO method for recovering a sparse solution from noisy
measurements [20]. The only difference here is that we are seeking a sparse
solution in the spectrum of a covariance matrix – hence seeking a low-rank
solution for the covariance. So to some extent, one may consider the
variational form as a “low-rank LASSO.” We have noticed a nice side benefit of
this LASSO-type formulation: empirically it seems to lead to better solutions
than solving the original MCR2 objective (see experimental results in Section
5). Part of the reason is likely because the LASSO type relaxation introduced
by the variational form finds a solution that is more stable to small noise in
the data or deviation from an ideal low-dimensional linear subspace.

###  3.3 Optimization Strategy

####  3.3.1 Alternating Maximization

To optimize (7), we adopt an alternating maximization strategy [2] between the
variation parameters $(\Gamma,\mathbf{A})$ and the network parameters
$(\theta)$. At each iteration, we first optimize $\Gamma$ and $\mathbf{A}$ by
taking one step of a proximal gradient ascent update which consists of taking
a gradient ascent step on the relevant part of variational loss $\Delta
R^{v}-\frac{\mu}{2m}M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$, followed by
normalizing the columns of $\Gamma$ to have unit $\ell_{2}$ norm and
thresholding the negative entries of the updated $\mathbf{A}$ matrix to $0$
(i.e., applying the ReLU function). To ensure stability of the gradient-based
method, we inversely scale the learning rate of $\Gamma$ and $\mathbf{A}$ by
upper-bounds of the Lipschitz constants of the gradients,
$\frac{1}{L_{\Gamma}}$ and $\frac{1}{L_{\mathbf{A}}}$, respectively. (See
Appendix for our derivation for the bounds.) Next, the matrix approximation
term $M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$ is recomputed using the
updated $\Gamma$ and $\mathbf{A}$, and we then update the network parameters,
$\theta$, by taking a gradient step on the relevant part of the variation
loss, $\nabla_{\theta}M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$.

In addition, note that from the variational form in Proposition 3.1 we know
that the optimal variational parameters $(\Gamma,\mathbf{A})$ should be
closely related to the singular values and vectors of
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$
for each class $j$ (the relationship becomes exact for large values of $\mu$).
We exploit this fact to initialize the variational parameters and to make
periodic ‘approximately closed-form’ updates to the variational parameters. We
call this procedure latching, which we describe in more detail in the next
Section 3.3.2. We summarize our overall training process in Algorithm
1.555Note that for clarity we describe the full procedure for gradient ascent,
but in our experiments stochastic gradient ascent is implemented.

Algorithm 1 Variational MCR2 Training

1:Input: data $\mathbf{X}$, labels $\mathbf{Y}$, featurizer
$f_{\theta}(\cdot)$, latch-freq, step sizes ($\nu_{\theta}$, $\nu_{\Gamma}$,
$\nu_{\mathbf{A}}$)

2:Initialize
$\mathbf{A},\Gamma\leftarrow\mathrm{latching}(\mathbf{X},\mathbf{Y},f_{\theta})$

3:for $\texttt{iter}=0,1,...,n-1$ do

4: Get $\mathbf{Z}_{\theta}=f_{\theta}(\mathbf{X})$ and membership matrices
$\mathbf{\Pi}$

5: Get $\ell_{\text{V-MCR}^{2}}(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})=$ (7)

6: Compute $L_{\mathbf{A}},L_{\Gamma}$ (see Appendix)

7:
$\Gamma\leftarrow\Gamma+\frac{\nu_{\Gamma}}{L_{\Gamma}}\nabla_{\Gamma}\ell_{\text{V-MCR}^{2}}(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$

8:
$\mathbf{A}\leftarrow\mathbf{A}+\frac{\nu_{\mathbf{A}}}{L_{\mathbf{A}}}\nabla_{\mathbf{A}}\ell_{\text{V-MCR}^{2}}(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$

9: Project $\mathbf{A}\leftarrow\mathrm{ReLU}(\mathbf{A})$

10: Project
$\Gamma_{l}\leftarrow\frac{1}{\left\|\Gamma_{l}\right\|_{2}}\Gamma_{l}\quad\forall
l\in[q]$

11: Recompute $M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$

12:
$\theta\leftarrow\theta-\nu_{\theta}\nabla_{\theta}(M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A}))$

13: if $\texttt{iter}\,\mbox{mod}\,\texttt{latch-freq}=0$ then

14:
$\mathbf{A},\Gamma\leftarrow\mathrm{latching}(\mathbf{X},\mathbf{Y},f_{\theta})$

15: end if

16:end for

17:return $f_{\theta}$

Algorithm 2 Latching

Input: data $\mathbf{X}$, labels $\mathbf{Y}$, featurizer $f_{\theta}(\cdot)$

Get $\mathbf{Z}_{\theta}=f_{\theta}(\mathbf{X})\in\mathbb{R}^{d\times m}$ and
membership $\mathbf{\Pi}\in\mathbb{R}^{m\times k}$

$\mathbf{A}\leftarrow\mathbf{0}\in\mathbb{R}^{q\times k}$ (assume $q$ is
divisible by $k$)

$\Gamma\leftarrow\mathbf{0}\in\mathbb{R}^{d\times q}$

for $j=1,...,k$ do

Get
$\mathbf{U}\mathrm{Diag}(\mathbf{\sigma})\mathbf{V}^{\top}=\mathrm{SVD}(\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top})$

$s\leftarrow q/k$

$\Gamma[:,(j-1)*s:j*s]=\mathbf{U}[:,0:s]$ % python indexing

$\mathbf{A}[(j-1)*s:j*s,j]=\mathbf{\sigma}[0:s]$ % python indexing

end for

return $\mathbf{A},\Gamma$

####  3.3.2 Latching

In order to optimize the variational MCR2 objective, the dictionary $\Gamma$
and $\mathbf{A}$ must maximize $\Delta
R(\mathbf{Z}_{\theta})=R^{v}(\Gamma,\mathbf{A})-R^{v}_{c}(\mathbf{A})$ whilst
minimizing the $\ell_{2}$ regularization term
$M(\mathbf{Z}_{\theta},\Gamma,\mathbf{A})$. This trade-off is controlled by
the regularization constant $\mu$. In practice, when $\mu$ is too large, each
gradient step does not allow for
$\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}$ to stray too far away from
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}$, which
can result in slow convergence. We observe that the following procedure
improves convergence in practice. Note that the variational form is
maximized666The maximization is exact as $\mu$ becomes large, but a good
approximation otherwise. when $\Gamma$ and $\mathbf{A}$ are derived from the
SVDs of
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$,
as given in Proposition 3.1. This gives a means to periodically reinitialize
the variational parameters $(\Gamma,\mathbf{A})$ based on the SVDs of
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$,
which we refer to as latching as described in detail in Algorithm 2. This
latching step can be viewed as taking an (approximate) full-maximization step
w.r.t. the variational parameters (as opposed to a proximal gradient descent
step) based on the closed-form solution provided in Proposition 3.1. This will
be an exact maximization step as $\mu$ becomes large. In short, given a
dictionary with $q$ columns, we initialize the dictionary as the concatenation
of the top $q/k$ singular vectors of
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$
for each class $j\in[k]$. Similarly, the columns of $\mathbf{A}$ are
initialized as the corresponding singular values. Though latching is in itself
an expensive procedure, requiring one to compute the SVD of a matrix $k$
times, it is optional (though we notice a benefit in practice) and can be done
relatively infrequently throughout training with a proper choice of the hyper-
parameter latch-freq. As a result, the amortized cost of latching becomes
insignificant.

##  4 Experimental Setup

We compare $\Delta R$, wall-clock time, and accuracy of models trained with
the original MCR2 objective and the variational MCR2 objective on MNIST [11],
CIFAR-10 [9], CIFAR-100 [9], and Tiny ImageNet [10] (with 200 classes)
datasets. We also compare the performance to cross-entropy (CE) training as a
benchmark of the correctness of the learned representations (for
classification). The high-level goal of these experiments is to show that 1)
the variational MCR2 objective is feasible for datasets where the original
MCR2 objective is computationally expensive (or impossible) to train (such as
CIFAR-100 and Tiny ImageNet), 2) show that training on the variational MCR2
maximizes the true $\Delta R$ objective and obtains the desired subspace-like
representations.

###  4.1 Hyperparameters

For fair comparison across training objectives (original MCR2, variational
MCR2), 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
MCR2 objectives, we use $\epsilon^{2}=0.5$ for all datasets. $\epsilon$ and
batch size are consistent with the experimental settings in the original MCR2
work [21] for MNIST and CIFAR-10. For Variational MCR2, 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$.

###  4.2 Nearest Subspace Classifier

MCR2 is a loss over the featurizer $f_{\theta}$. To classify the test data, we
use the nearest-subspace classifier similar to the original $\text{MCR}^{2}$
work [21]. As shown by [21], at the global optima of MCR2, the representations
of each class lie on low-dimensional subspaces that are orthogonal to each
other. Yu, et al. [21] also empirically observe this property for networks
trained by SGD. Assuming that the learned representations satisfy this
property, given a test datapoint, we can simply identify the closest subspace
for the final classification. Formally, given a test sample
$\mathbf{z}_{test}=f_{\theta}(\mathbf{x}_{test})$, the predicted label is
given as

$y=\operatorname*{arg\,min}_{j\in{1,...,k}}\left\|(\mathbf{I}-\mathbf{V}^{(j)}(\mathbf{V}^{(j)})^{\top})\mathbf{z}_{test}\right\|^{2}_{2},$
(8)

where $\mathbf{V}^{(j)}$ is a matrix of the top $\lfloor{\frac{d}{k}}\rfloor$
principal components of
$\mathbf{Z}_{\theta}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}_{\theta}^{\top}$
with $\mathbf{Z}_{\theta}=f_{\theta}(\mathbf{X})$ being the embedding of the
training data $\mathbf{X}$.

##  5 Experimental Results

We discuss the performance of variational MCR2 below. Performance is measured
by 1) training speed, 2) the true $\Delta R$ value the model reaches over the
training data, and 3) the classification accuracy over the test set. Note that
we are more interested in comparing the accuracy and training efficiency
between MCR2 and V-MCR2, over trying to reach state of the art results in this
paper.

(a) MNIST

(b) CIFAR-100 $d=100$ (top), $d=500$ (bottom)

(c) CIFAR-10

(d) Tiny ImageNet $d=200$ (top), $d=500$ (bottom)

Figure 1: Convergence of training $\mathbf{\Delta R}$. We compare the $\Delta
R(\mathbf{Z}_{\theta})$ of the training data over epochs for V-MCR2 and MCR2.
For both V-MCR2 and MCR2, the network is optimized by stochastic gradient
descent with a learning rate of $10^{-3}$. All training runs are 2000 epochs,
excluding the MCR2 runs for Tiny ImageNet, which we stop at 500 and 200 for
$d=200$ and $d=500$ respectively, due to significant computational cost from
the high number of classes. Also, note that for V-MCR2, we often observe small
undulations in the training loss due to the regular reinitialization by
latching.

###  5.1 Computational Efficiency

Across datasets, we compare the wall-clock time to train one epoch using the
MCR2 (1) and V-MCR2 (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-MCR2 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 MCR2
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 MCR2 becomes
nearly impractical, while it is easily handled by V-MCR2.

Dataset | MCR2 | V-MCR2
---|---|---
MNIST | 11.56 | 6.29
CIFAR-10 | 33.06 | 20.71
CIFAR-100 | 157.45 | 31.14
Tiny ImageNet | 527.85 | 44.23
Table 1: Wall-Clock Time per Epoch (in seconds). We compare the wall-clock
time to complete one epoch of training. MNIST and CIFAR-10 use a batch size of
1000. For CIFAR-100 and Tiny ImageNet, we use a batch size of 2000.

Additionally, we compare the true $\Delta R(\mathbf{Z}_{\theta})$ (i.e.,
computing the original MCR2 objective with the iterates of the V-MCR2 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-MCR2 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-MCR2 and MCR2 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
MCR2 and V-MCR2. With a feature dimension of $d=500$, we observe across 5
seeds that training with the original MCR2 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-MCR2 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.

(a) MNIST

(b) CIFAR-10

(c) CIFAR-100

(d) Tiny ImageNet

Figure 2: Inner product of representations. We plot the heatmap of
$|\mathbf{Z}_{\theta}^{\top}\mathbf{Z}_{\theta}|$ where $\mathbf{Z}_{\theta}$
are the representations of the training data, ordered by the class they belong
to. For CIFAR-100 and Tiny ImageNet, 10 classes are randomly chosen. If the
classes lie on low-dimensional orthogonal subspaces, we expect to see a block
diagonal structure.

###  5.2 V-MCR2 Obtains Better Representations

There are two properties of the representation that we aim to attain by
optimizing the MCR2 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 MCR2 and V-MCR2 (Figure 1). For
MNIST and CIFAR-10, note that both MCR2 and V-MCR2 obtain a block diagonal
structure. However, on CIFAR-100 and Tiny Imagenet, we observe no block
diagonal structure after MCR2 training, whereas we see a clear block diagonal
structure for V-MCR2. These findings suggest that V-MCR2 training is more
robust to avoiding poor local minima than training on the original MCR2 model,
particularly as the number of classes increase. We leave a rigorous study of
these phenomena for a future work.

###  5.3 Performance on Classification Tasks

In Table 2, we present test accuracies on the four aforementioned datasets
when trained under original MCR2 and V-MCR2 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 purpose777We 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 [22] 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 hyperparameters888For experiments using
CE, we use a higher learning rate of $10^{-2}$ to improve convergence for CE
training. as in V-MCR2 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 MCR2 and V-MCR2 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 MCR2 objective
results in a poor test accuracy, while the V-MCR2 model achieves comparable
performance to training the same network with CE. Due to limited resources and
heavy computational requirements of MCR2 training with a large number of
classes, we can only report the result on Tiny ImageNet by optimizing MCR2
after 200 epochs while the other two objectives complete the full training
session of 2000 epochs.

Dataset | Objective |  Training $\mathbf{\Delta R}$ | Test Accuracy
---|---|---|---
|  MCR2 | 44.6429 | 0.9785
MNIST |  V-MCR2 | 44.2117 | 0.9788
| CE | - | 0.9738
|  MCR2 | 49.40 | 0.8956
CIFAR-10 |  V-MCR2 | 48.43 | 0.8997
| CE | - | 0.8665
|  MCR2 | 226.0519 | 0.2421
CIFAR-100 |  V-MCR2 | 218.0185 | 0.5872
| CE | - | 0.5840
Tiny |  MCR2 | 227.6468 | 0.1319
ImageNet |  V-MCR2 | 231.1538 | 0.2665
200 | CE | - | 0.1907
Table 2: Comparison of classification performance. We evaluate the training
$\Delta R$ and test accuracy of CE, MCR2, and V-MCR2 after 2000 training
epochs for MNIST, CIFAR-10, and CIFAR-100. For Tiny ImageNet, we report the
results after 200 epochs for MCR2 due to very slow training, though note that
$\Delta R$ had already converged (Fig. 1(d)). For V-MCR2 and CE, results are
after 2000 epochs. For CIFAR-100 and Tiny ImageNet, we report results for
$d=500$. See Appendix for results for other choices of $d$.

##  6 Conclusion

Building on Yu et al. [21], we propose an alternative cost-efficient
formulation of the MCR2 objective that is scalable to datasets with a large
number of classes. Namely, for CIFAR-100 and Tiny ImageNet, we observed a
$5\times$ and $12\times$ speedup per training epoch, respectively. The gain
would be even more significant as the number of classes increases.
Additionally, we show that we do not make any compromises when it comes to
performance by using this approximate, variational formulation. In fact, in
all datasets we tested, not only does V-MCR2 reach similar $\Delta R$ values
as MCR2, the learned representations of V-MCR2 were just as good, oftentimes
better, than those learned by MCR2 for the same number of training epochs. On
a related note, we observe an interesting phenomenon that directly optimizing
the MCR2 objective is not only slow, but often completely fails to learn the
desired orthogonal subspace structures when the number of classes increases
even after $\Delta R$ has successfully converged. In our experiments, we found
that V-MCR2 is surprisingly not prone to this problem, suggesting that V-MCR2
allows for better control over learning the desired representations. By these
means, we find that V-MCR2 is a very promising adaptation of MCR2 and we leave
a rigorous comparison of the representations learned by MCR2 and V-MCR2 for
future work.

Acknowledgements This work was partially supported by NSF Fellowship
DGE2139757, NSF grants 1704458 and 2031985, and the Northrop Grumman Mission
Systems Research in Applications for Learning Machines (REALM) initiative. Yi
Ma was funded by ONR grants N00014-20-1-2002 and N00014-22-1-2102, the joint
Simons Foundation-NSF DMS grant $\\#2031899$.

## References

  * [1] Zeyuan Allen-Zhu and Yuanzhi Li.  Towards understanding ensemble, knowledge distillation, and self-distillation in deep learning.  https://arxiv.org/abs/2012.09816, 2021.
  * [2] J. Eckstein.  Augmented Lagrangian and alternating direction methods for convex optimization: A tutorial and some illustrative computational results.  RUTCOR Technical Report, 2012.
  * [3] M. Fazel, H. Hindi, and S.P. Boyd.  Log-det heuristic for matrix rank minimization with applications to hankel and euclidean distance matrices.  In Proceedings of the 2003 American Control Conference, 2003., volume 3, pages 2156–2162 vol.3, 2003.
  * [4] R Geirhos, J.-H. Jacobsen, C Michaelis, R Zemel, W Brendel, Bethge M, and F. A. Wichmann.  Shortcut learning in deep neural networks.  Nature Machine Intelligence, 2020.
  * [5] Paris Giampouras, René Vidal, Athanasios Rontogiannis, and Benjamin Haeffele.  A novel variational form of the schatten-$p$ quasi-norm.  Advances in Neural Information Processing Systems, 33, 2020.
  * [6] Magnus R Hestenes.  Multiplier and gradient methods.  Journal of optimization theory and applications, 4(5):303–320, 1969\.
  * [7] Like Hui, Mikhail Belkin, and Preetum Nakkiran.  Limitations of neural collapse for understanding generalization in deep learning.  arXiv preprint arXiv:2202.08384, 2022.
  * [8] Ivan Kobyzev, Simon Prince, and Marcus Brubaker.  Normalizing flows: An introduction and review of current methods.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020\.
  * [9] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [10] Ya Le and Xuan Yang.  Tiny imagenet visual recognition challenge.  Stanford cs231, 2015.
  * [11] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [12] Yi Ma, Harm Derksen, Wei Hong, and John Wright.  Segmentation of multivariate mixed data via lossy data coding and compression.  IEEE transactions on pattern analysis and machine intelligence, 29(9):1546–1562, 2007.
  * [13] Bruno A Olshausen and David J Field.  Emergence of simple-cell receptive field properties by learning a sparse code for natural images.  Nature, 381(6583):607, 1996.
  * [14] Bruno A Olshausen and David J Field.  Sparse coding with an overcomplete basis set: A strategy employed by v1?  Vision research, 37(23):3311–3325, 1997.
  * [15] Marcus Valtonen Ornhag, Carl Olsson, and Anders Heyden.  Bilinear parameterization for differentiable rank-regularization.  Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 2020.
  * [16] Vardan Papyan, X.Y. Han, and David Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  Proceedings of the National Academy of Sciences of the United States of America, 2020.
  * [17] R Tyrell Rockafellar.  The multiplier method of Hestenes and Powell applied to convex programming.  Journal of Optimization Theory and Applications, 12(6):555–562, 1973.
  * [18] Luca Scimeca, Seong Joon Oh, Sanghyuk Chun, Michael Poli, and Sangdoo Yun.  Which shortcut cues will dnns choose? a study from the parameter-space perspective.  arXiv preprint arXiv:2110.03095, 2021.
  * [19] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli.  The pitfalls of simplicity bias in neural networks.  Advances in Neural Information Processing Systems, 2020.
  * [20] R. Tibshirani.  Regression shrinkage and selection via the LASSO.  Journal of the Royal Statistical Society B, 58(1):267–288, 1996\.
  * [21] Yaodong Yu, Ryan Chan, Chong You, Chaobing Song, and Yi Ma.  Learning diverse and discriminative representations via the principle of maximal coding rate reduction.  NeurIPS, 2020.
  * [22] Sukmin Yun, Jongjin Park, Kimin Lee, and Jinwoo Shin.  Regularizing class-wise predictions via self-knowledge distillation.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 13876–13885, 2020.

##  7 Appendix

###  7.1 Derivation of the Lipschitz Constants

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:

$\begin{split}f(\Gamma,\mathbf{A})=&\frac{\mu
k}{2m}\sum_{j=1}^{k}\|\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}-\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\|_{F}^{2}\\\
=&\frac{\mu
k}{2m}\langle\mathbf{\Pi}\mathbf{\Pi}^{\top},(\mathbf{Z}^{\top}\mathbf{Z})^{\odot
2}\rangle-\frac{\mu
k}{m}\langle\mathbf{A}\mathbf{\Pi}^{\top},(\Gamma^{\top}\mathbf{Z})^{\odot
2}\rangle+\frac{\mu
k}{2m}\langle\mathbf{A}\mathbf{A}^{\top},(\Gamma^{\top}\Gamma)^{\odot
2}\rangle\end{split}$ (9)

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

From this, the relevant gradients are

$\displaystyle\nabla_{\Gamma}f$ $\displaystyle=-\frac{2\mu
k}{m}\sum_{j=1}^{k}\left(\underbrace{\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}\Gamma\mathrm{Diag}(\mathbf{A}_{j})}_{:=g_{j}(\Gamma)}-\underbrace{\Gamma\mathrm{Diag}(\mathbf{A}_{j})\Gamma^{\top}\Gamma\mathrm{Diag}(\mathbf{A}_{j})}_{:=h_{j}(\Gamma)}\right)$
(10) $\displaystyle\nabla_{\mathbf{A}}f$ $\displaystyle=\frac{\mu
k}{m}\left[(\Gamma^{\top}\Gamma)^{\odot
2}\mathbf{A}-(\Gamma^{\top}\mathbf{Z})^{\odot 2}\mathbf{\Pi}\right]$ (11)

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

$\displaystyle\frac{\|g_{j}(\Gamma+\Delta\Gamma)-g_{j}(\Gamma)\|_{F}}{\|\Delta\Gamma\|_{F}}$
$\displaystyle=\frac{\|\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}\Delta\Gamma\mathrm{Diag}(\mathbf{A}_{j})\|_{F}}{\|\Delta\Gamma\|_{F}}$
(12)
$\displaystyle\leq\|\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}\|_{F}\|\mathbf{A}_{j}\|_{\infty},$
(13)

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 [21]. Now, we have the following
approximation

$\bar{h}_{j}(\Gamma)\approx\Gamma\mathrm{Diag}(\mathbf{A}_{j})^{2}.$ (14)

which implies,

$\displaystyle\frac{\|\bar{h}_{j}(\Gamma+\Delta\Gamma)-\bar{h}_{j}(\Gamma)\|_{F}}{\|\Delta\Gamma\|_{F}}$
$\displaystyle\approx\frac{\|\Delta\Gamma\mathrm{Diag}(\mathbf{A}_{j})^{2}\|_{F}}{\|\Delta\Gamma\|_{F}}$
(15) $\displaystyle\leq\|\mathbf{A}_{j}\|_{\infty}^{2}$ (16)

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

$L_{\Gamma}=\frac{2\mu
k}{m}\sum_{j=1}^{k}\|\mathbf{Z}\mathrm{Diag}(\mathbf{\Pi}_{j})\mathbf{Z}^{\top}\|_{F}\|\mathbf{A}_{j}\|_{\infty}+\|\mathbf{A}_{j}\|_{\infty}^{2}$
(17)

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

$\displaystyle\frac{\left\|\nabla_{\mathbf{A}}f\big{|}_{\mathbf{A}+\Delta\mathbf{A}}-\nabla_{\mathbf{A}}f\big{|}_{\mathbf{A}}\right\|_{F}}{\|\Delta\mathbf{A}\|_{F}}$
$\displaystyle=\frac{\|\frac{\mu k}{m}(\Gamma^{\top}\Gamma)^{\odot
2}\Delta\mathbf{A}\|_{F}}{\|\Delta\mathbf{A}\|_{F}}$ (18)
$\displaystyle\leq\frac{\mu k}{m}\|(\Gamma^{\top}\Gamma)^{\odot
2}\|_{F}=L_{A}$ (19)

###  7.2 Architecture

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.

Architecture 1 Neural Network Architecture for MNIST

1:Conv2d(in_channel=1, out_channel=32, kernel=3, stride=1)

2:ReLU()

3:Conv2d(in_channel=32, out_channel=64, kernel=3, stride=1)

4:ReLU()

5:MaxPool2d(kernel=2, stride=None)

6:Dropout(p=0.25)

7:Flatten()

8:Linear(12544, d)

9:ReLU()

10:Dropout(p=0.5)

11:Linear(d, d)

12:Normalize()

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 MCR2 objective.

Architecture 2 Reshaping Layers for ResNet18

1:Linear(512, 512, bias=False)

2:BatchNorm1d()

3:ReLU()

4:Linear(512, d, bias=True)

5:Normalize()

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

###  7.3 Data Augmentation

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

Transformations 1 Transformations for MNIST, CIFAR-10, CIFAR-100

1:import torchvision.transforms as transforms

2:TRANSFORM = transforms.Compose([

3: transforms.RandomCrop(32, padding=8),

4: transforms.RandomHorizontalFlip(),

5: transforms.ToTensor()])

Transformations 2 Transformations for Tiny ImageNet

1:import torchvision.transforms as transforms

2:TRANSFORM = transforms.Compose([

3: transforms.Resize(32)

4: transforms.RandomHorizontalFlip()

5: transforms.ToTensor()])

###  7.4 Additional Experiments

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 [22].

Here we report experimental results on CIFAR-100 and Tiny ImageNet by training
with a similar training recipe adopted in [22]. Specifically, as in [22], we
use PreAct ResNet-18 as the backbone and train with both cross-entropy loss
and V-MCR2. All networks are trained by SGD with momentum 0.9, and weight
decay of $10^{-4}$. Similar to [22], 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 [22] 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-MCR2. In addition
to the training strategy, we set $\mu=10^{-1}$ for V-MCR2 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.

Dataset | Objective |  Training $\mathbf{\Delta R}$ | Test Accuracy
---|---|---|---
CIFAR-100 |  V-MCR2 | 130.2177 | 0.6951
| CE | - | 0.7146
Tiny ImageNet |  V-MCR2 | 134.9504 | 0.4189
200 | CE | - | 0.4843
Table 3: Comparison of classification performance. We evaluate the training
$\Delta R$ and test accuracy of V-MCR2 after 2000 training epochs for
CIFAR-100 and Tiny ImageNet. For CE, we evaluate the test accuracy at each
epoch and report the highest test accuracy achieved across the 2000 epochs.
[22] similarly achieved an accuracy 0.7529 for CIFAR-100 and 0.5647 for Tiny
ImageNet for CE training. With V-MCR2, we achieve a lower test accuracy. Note
that the hyperparameters for V-MCR2 such as latch frequency, dictionary size,
and $\mu$ were not rigorously tuned. Additionally, due to the computational
cost of the nearest subspace classification algorithm, we only evaluate the
test performance at the 2000 epoch mark.