Massively Scaling Heteroscedastic Classifiers
Authors: Mark Collier
Rodolphe Jenatton
Basil Mustafa
Neil Houlsby
Jesse Berent
Effrosyni Kokiopoulou
Venue: ICLR 2023 poster (2023)

Abstract: Heteroscedastic classifiers, which learn a multivariate Gaussian distribution over prediction logits, have been shown to perform well on image classification problems with hundreds to thousands of classes. However, compared to standard classifiers, they introduce extra parameters that scale linearly with the number of classes. This makes them infeasible to apply to larger-scale problems. In addition heteroscedastic classifiers introduce a critical temperature hyperparameter which must be tuned. We propose HET-XL, a heteroscedastic classifier whose parameter count when compared to a standard classifier scales independently of the number of classes. In our large-scale settings, we show that we can remove the need to tune the temperature hyperparameter, by directly learning it on the training data. On large image classification datasets with up to 4B images and 30k classes our method requires 14X fewer additional parameters, does not require tuning the temperature on a held-out set and performs consistently better than the baseline heteroscedastic classifier. HET-XL improves ImageNet 0-shot classification in a multimodal contrastive learning setup which can be viewed as a 3.5 billion class classification problem.

Published as a conference paper ICLR 2023
MASSIVELY SCALING HETEROSCEDASTIC CLASSIFIERS
Mark Collier*, Rodolphe Jenatton*, Basil Mustafa, Neil Houlsby, Jesse Berent & Effrrosyni Kokiopoulou Google Al
(markcollier, rjenatton, bas 1m, nei nei leeilhhulsby jberent, kokiopou com
ABSTRACT
Heteroscedastic classifiers, which learn multivariate Gaussian distribution over predic- tion logits, have been shown to perform well on image classification problems with hun- dreds to thousands of classes. However, compared to standard classifiers, they introduce extra parameters that scale linearly with the number classes. This makes them infeasible to apply to larger-scale problems. In addition heteroscedastic classifiers introduce critical temperature hyperparameter which must be tuned. We propose HET-XL, a heteroscedastic classifier whose parameter count when compared to standard classifier scales indepen- dently of the number of classes. In our large-scale settings, we show that we can remove the need to tune the temperature hyperparameter, by directly learning it the training data. On large image classification datasets with up to 4B images and 30k classes our method requires 14x fewer additional parameters, does not require tuning the tempera- ture on a held-out set and performs consistently better than the baseline heteroscedastic classifier. HET-XL improves ImageNet O-shot classification multimodal contrastive learning setup which can be viewed as 3.5 billion class classification problem.
INTRODUCTION
Heteroscedastic models learn input-dependent noise term to capture uncertainty in their predictions. In deep learning, they have been used successfully in large-scale image classification (Collier et al., 2021), image segmentation (Kendall & Gal, 2017: Collier et al., 2020), regression (Lakshminarayanan et al., 2017),
uncertainty quantification (Tran et al., 2022; Nado et al., 2021) and in bandit problems (Osband et al., 2021).
It is known from the economics literature that heteroscedastic classifiers particularly suited modelling classification problems with many classes (Train, 2009) and this has been further observed in deep learning (Collier et al., 2021). However, heteroscedastic classifiers add additional parameters to standard "determin- istic" classifiers (DET) to define their K x K covariance matrix, with K the number of classes. Even with low-rank approximations, the number of additional parameters scales linearly in K, thus imposing a signif- icant cost in large-scale settings. Also, these additional parameters must be stored in long-term storage and loaded in memory which can pose problems for both storage and memory bound applications. For example, on JFT-4B, a dataset with 29,593 classes, the state-of-the tott----ttt-------e and most scalable, to the best of our knowledge, heteroscedastic classification method HET (Collier et al., 2021), 1), does not fit in memory on a large TPU slice (64 TPU v3 cells with 128 cores) when using modest-sized ViT-L/32 base architecture.
this paper, we propose HET XL whose extra parameter count over DET scales independently of the number of classes. In addition, HET requires tuning temperature hyperparameter which hinders the adoption of heteroscedastic classifiers in large-scale settings where hyperparameter sweeps are either very costly not feasible at all. HET in contrast, learns directly on the training set. We argue and demonstrate empirically that this is feasible precisely in this very large-scale setting. Despite the improved
*Equal contribution.

---

Published as conference paper ICLR 2023
<((x))
W
W
Pre Pre-logiit
Preflggits
Pre-iogits + noise
Logits noise
Logits
Logits noise
(a2; 0)
Network
Neevork
Figure 1: HET-XL (left) injects the noise e((x) ~N(, E' (a)) in the pre-logits W (o(x;; + (a))), while HET (right) adds the noise in the logits W ((x;0) <(x) and e(x) ER E' does not scale with K.
efficiency and ease adoption, HET XL performs consistently better across large-scale image classification tasks compared to DET and HET and improves upon a DET baseline in the contrastive learning setting.
Contributions. summary, our contributions are:
(1) We develop the HET-XI. heteroscedastic classifier, whose cost of deployment is significantly reduced as compared to НET, the prior state-of-the-art. In large-scale settings, HET-XL does not require tuning a tem- perature hyperparameter and has massively reduced parameter count compared to НET. Moreover, HET- XL allows for plug-in compatibility with existing large-scale classifiers, which not the case for HET. (2) On three image classification Lenchmarks-Jep 300M, ImageNet-21k and JFT-4B and for two differ- ent popular model classes, ResNet152 and ViT-L/32, HET-XL consistently outperforms HET and НET-H, hashing-based baseline introduce. For example, with ResNet152 on JFT-300M, we increase precision@ @ by 2.3% compared to HET, while adding about times fewer parameters. (3) We extend HET -X contrastive learning where the method improves ImageNet 0-shot classification accuracy from 85.29% for DET model to 85.56% on a LiT setup (Zhai et al., 2022).
2 BACKGROUND ON HETEROSCEDASTIC CLASSIFIERS
We now review the core prior work that our method builds on, wider review of related work is provided in Appendix A. We focus on classification tasks where we learn classifiers of the form
softmax [W 0)) and s sigmoid(W) (xx;0)
based training data D (k) Mn) A pair (æn., Vn.) corresponds to input In. e.g., image, together with its label yn. E {0, 1] belonging to one, or multiple, of the K classes, in the multi-class and multi-label settings, respectively. The model is parametrized by W RR and the D-dimensional representation (f::0) ()) output by neural network with parameters O. We have omitted the bias term to ease the presentation. Throughtout the paper, we refer to W o( E Æ as the logits, while we will use the term pre-logits for @(x; 0) E R.D. We denote the elementwise product between tensors by o.
Heteroscedastic classifiers learn an additional input-dependent noise distribution placed the logits capture uncertainty in the predictions of the model (Kendall & Gal, 2017; Collier et al., 2021; Train, 2009). We consider the setting where this noise is modelled by Gaussian, leading to the class predictions
(2) where a can be either the softmax or the s sigmoid. transformation (see Fig. 1, right). Above, we have introduced the covariance matrix Exx; fewr) of that is parametrized by Awwv will describe Σ in more details in Section 2.2. The resulting conditional probability p(y) æ;; [W, A, f in Eq. (2) used to train the model D by maximum likelihood and make predictions evaluation time.
Ee [a (WT 6(x;6) e(x))] R with <(x) R (O
2.1 ESTIMATING THE HETEROSCEDASTL PREDICTIONS
We will focus on the нET method from Collier et al. (2021) as obtains state-of-the-ae performance on several benchmarks and operates at the largest scale. As shown in Eq. (2), heteroscedastic modelling requires
2

---

Published as a conference paper ICLR 2023
marginalizing over the noise E(x). The corresponding expectation cannot be solved analytically (Lu al., 2020) and Collier et al. (2021) use Monte Carlo (MC) estimate by sampling from (o,, (x; 9 )), Fig. 1.
By relating Eq. (2) to generative process where o approximates discrete choice (Train, 2009), Collier al. (2020; 2021) further consider temperature parameter: 7 > 0 to control this approximation. More precisely, a Eq. (2) replaced by o, so that or (u( o(1/ - u). We present more formally the construction of the HET classifier in Appendix A.I. Importantly, Collier et al. (2020) show that is crucial to regulating a cas-variance trade-off between the bias with respect to the generative process and the variance of the MC estimate. is tuned on a held-out set and the test performance is sensitive to its choice.
2.2 PARAMETRIZING THE COVARIANCE MATRIX
A central component of heteroscedastic classifiers is the covariance matrix (ax; Bene) E RRKK It enables the model to learn which regions of the input space have noisy labels and what the correlations across classes are in those regions (Collier et al., 2021). When clear from the context, we will omit the parameters &co--
Making E(x) dependent on the inputs is challenging, Let us assume that ((x) = L(x) L()) is positive definite with the full-rank matrix L(x) k that vectorize vec (Lx)) R= Even simply using a linear transformation of the D-dimensional pre-logits to parameterize vec(L(x)) as
vec(L(x)) Cá(x;; R2 with CERKAXD
is not feasible due to the large number of elements in Oan {C} (for instance, K = 29, 593 and D = 2048 in some of our settings in Section 6). We could restrict E(x) to diagonal matrix (Kendall & Gal, 2017), scaling down C' to R x but this comes with drop in performance (Collier et al., 2021). Instead, low-rank parametrization, with R < K, of the form
E(x) = V (x) V (x) + diag(d(x:) with V(x) E IR Rxx and d(x) E R
(4)
offers a good trade-off between memory footprint and performance of the classifier (Collier et al., 2021). In that case, using a linear transformation of the pre-logits, as above, leads to Oen of size O(DKR= DK).
Collier et al. (2021) consider a further optimized parametrization whereby V(x) = J (1æv(z)T) where J E IR Rx K and v(x) R K with 6eov thus scaling in O(DK + KR) (in Appendix L. we discuss slightly more expressive variants of Eq. (4) that do not increase the parameter count). Despite these advances, O(DK + K R) remains restrictive for modern large models and problems with large number of classes
3 HET-XL: SCALING HETEROSCEDASTIC CLASSIFIERS TO VERY LARGE K
Heteroscedastic classifiers have resulted in impressive performance gains across many tasks, see Section 1. Empirically those improvements have been the biggest in large-scale classification tasks with many classes (Tran et al., 2022; Collier et al., 2021). However, their additional parameter count scales linearly in the number of classes K, see Section 2. The increase in parameter count and corresponding memory use can be large relative to the performance gains. For example, DET ResNet152 applied to ImageNet-21K has 102.9M parameters, while a НET ResNet152 has 88% more parameters than DET. НEТ provides a 1.4% increase in precision@1 for this task, however the additional cost may be considered too high for this gain.
In this section we will introduce our proposed solution, the HET-XL classifier. HET-XL has significantly reduced deployment cost when K is large and still provides performance gains over HET, whose large parameter count can lead to overfitting, see Section On the above ImageNet-21k example, HET-XL provides 2.4% performance gains with only 8% more parameters compared to DET.
3

---

Published as a conference paper ICLR 2023
Table Summary of the properties of HET and HET-XL.
Extra parameter count
Automated tuning of Plug-in compatibility with existing classifiers
HET HET-XL
ODD R
+
3.1 HET-XL
The additional parameter count of HET compared to DET scales in O(DK KR) because the noise term 6 is directly parameterized in the logit space with dimension K. We make the simple yet important obser- vation: If we directly inject the noise term the pre-logits d( E RD, and then reuse the existing linear transformation W E DxX from pre-logits to logits, we can remove this dependency on the number of classes K. This core idea illustrated Fig. We refer to the resulting extension of HET as HET-XL. For HET-XL, we thus replace Eq. (2) with
Ee [or (WT (6(=;)) + c(()))) with c((x) NO
E' (x; ))
(5)
where crucially e' E R2 and the covariance matrix Ex(x; DxD apply to the pre-logits. Therefore, when we parameterize E' as per the HET parametrizations recalled in Section 2.2, the additional parameter count of lwv compared to DET scales in O(D² + DR) and D(D2R8 respectively. In large-scale settings D is often small relative K. For example, for JFT-4B, K 29, 593 while
1024 for ViT-L and D 2048 for 52s two representative large-scale image classification models.
Because of the properties of Gaussian distributions under linear transformations, W e' still defines a Gaussian distribution in logit space: N(W oxx;; 0). W W)) However, the covariance matrix, or some decomposition thereof, is never explicitly computed in this space. Interestingly, (Appendix I), we show that the choice of sharing the W transformation between the pre-logits and the noise samples does not sacrifice performance compared to separate transformations W
Table 15
<(x;; + (W')Te')
Plug-in compatibility with existing large-scale classifiers. In extreme classification tasks, where K may be in the millions, the standard matrix- multiplication W o(x; 8) may have to be replaced by more scalable logic, involving for instance distributed lookup of the active classes (Zhang et al., 2018; Song et al., 2020). Let us assume this logic is abstracted by an API of the form get logi segiitsodd (6(x; 0)). HET-XL, directly piggyback that functionality by using get log i t (6(1;; 0) + e'(x)) while HET cannot be easily adapted to this case. Table summarises the differentiating properties of HET and HET-XL.
4 LEARNING THE TEMPERATURE
As described in Section 2.1, the temperature of HET controls bias-variance trade-off. The resulting performance is sensitive to the value r whose choice is often dataset-dependent (Collier et al., 2020; 2021). In practice, is included in the hyperparameter sweep which may become prohibitively expensive to run at scales considered here. At the same time, bypassing this step or making the sweep too coarse may degrade performance (Collier et al., 2020). We thus investigate strategies to automatically tune
4.1 A REPRESENTATIVE BENCHMARK
As working study, we consider the tuning of in HET-XL for ResNet50x2 model ImageNet-21k trained for 7 epochs (the dataset will be presented in more details in Section 6). "x2" refers to the factor multiplying the width of each layer (Szegedy al., 2017). This setting is representative of the large-scale tasks we deal with in our subsequent experiments, although slightly scaled down, both in terms of model size and training regime. We do so to be able to compare with wider spectrum of methods and assess the
4

---

Published as a conference paper ICLR 2023
Table 2: Comparisons (means + standard errors) of methods to tune T for a ResNet50x2 on ImageNect22kk
Method
.NNLLL
PPre 1
Single training?
No validation set needed?
Grid search
6.08 0:04
0.433
0,002
X
X
Bayesian optimization
6.03 0.06
0.428
0,000
X
X
Luketina et al. (2016)
6.08
0:04
0.425
0,000
Simple training (as in HET-XL)
5.99
0.04
0.435
0.001
validity of the results with multiple replications (we will report the average and the standard error over three replications). We provide further details of the benchmark in Appendix G.
Approaches with multiple trainings. To obtain gold-stane gold-standard references, we consider methods that tune r based on multiple successive trainings. In particular, we focus on grid search (GS) (assuming a grid of values for t informed by an experienced user; see Appendix G) and Bayesian optimisation (BO) (Snoek et al., 2012) which we run for 30 full trainings; we use the implementation of Song et al. (2022); Golovin et al. (2017). None of those approaches would scale up to our largest experiments in Section 6.
Online gradient-based tuning. Given that T is a one-dimensional continuous hyperparameter, it is mean- ingful to consider approaches that optimize validation objective by gradient descent, typically following a bi-level formulation (Maclaurin et al., 2015; Luketina et al., 2016; Pedregosa, 2016; Franceschi et al., 2018; Lorraine et al., 2020). Those approaches need to approximate costly high-order derivatives that account for the dependency of the hyperparameters in the validation objective. In our study, we will take Luketina representative gradient-based method that considers the hyperparameter dependency only through the current training step gradient (we formally describe their approach in Appendix G). Moreover, because of the particular structure of e-explicilly appearing both at training and validation time, unlike coniiiisttto--eellteee or regularisation-reltte hyperparameters that surface only at training time- it is feasi- ble to consider and evaluate simpler alternative gradient estimator, see Appendix G.
al. (2016)
Do we need a validation set? Simple training of т. The experiments we consider in Section 6 bring into play datasets with a large number of data points (from 12.8M up to 4B) over which training limited to only few epochs. As result, it is natural to wonder to what extent metrics collected on the training set are faithful to metrics computed on a proper held-out validation set, since overfitting seems unlikely to happen. We thus consider an approach wherein is treated as simple training variable, optimized via the training objective like { (W, 0,00
the contrastive learning setting, Jia et al. (2021); Radford et al. (2021); Mustafa et al. (2022) already successfully learned temperature during training, although in different context where T was not responsible for a specific Sarrinai trade-off like in HET-XL.
4.2 RESULTS AND DISCUSSION
We summarise the results in Table 2. The main conclusion is that there is no opportunity cost in using the simplest approach-tunne tuning as a standard training parameter-eve even when compared to more advanced, c.g., Luketina et al. (2016), and more expensive methods like GS and BO requiring multiple trainings. Remarkably, we still observe the benefit of "Simple training" in our larger-scale experiments, see Table 8.
In the perspective of having a plug-in deployment of HET HET-XL, "Simple training" has the advantage over Luketina et al. (2016) not to require an extra validation set to iterate over during training, which may be constraining in production pipelines. In Appendix G. we try to replace the validation set of Luketina et al. (2016) by some held-out subset, say 25% of the (training) batches, but we observe some drop in performance (past the first epoch, this held-out subset of the batches does not behave actual held-out data anymore).
Let us consider the split of the entire training set D into Dirai D We could hypothesise that GS and BO perform worse because their corresponding models are trained on less data, i.e., just D Dumin (with D being
5

---

Published as a conference paper ICLR 2023
used guide the tuning), as opposed to Drain U Daal for "Simple training". Appendix G, we check the effect of retraining on vD the best models obtained by GS and BO, with no change in our conclusion.
Hence, from now on, we assume that the temperature T in HET-XL is trained like any other model parameter.
5 CONTRASTIVE LEARNING
Contrastive learning is growing area of research (Radford et al., 2021; Jia et al., 2021; Zhai et al., 2022; Yu et al., 2022; Mustafa et al., 2022) where pairs of data, possibly of different modalities, e.g., image and text, are learned to be aligned. Contrastive models have shown strong performance in zero-shot image classification and retrieval, as well as when learning representations for downstream tasks.
5.1 HETEROSCEDASTIC CONTRASTIVE LEARNING
Contrastive learning as classification. We can view contrastive learning as massive classification prob- lem. Assume a dataset of N image-text pairs D {(an: bn) We are trying to learn well-aligned representations of (a,, b,.) that we denote by (an, Bn), with on Bn R output by some neural network.
We want to push (axz B..) to be similar while {a.., Br Jise are encouraged to be dissimilar. In this sense, we can see this setting as having two classification tasks, one from image to text-given On, identify the most similar B,, and one from text to mmage-given Bn, identify the most similar at, at,-with K = N classes in both cases. Since N > those two "global" classification tasks are approximated by "local", batch-level classification tasks. Indeed, during training, given a batch B C {l,..., .N7 of size B = |B] << N, we focus the contrastive loss L(B) 2B
E (C (iimgge-text (B) Ln plext-image (B)) where we have defined
TTT Am
rimage text (B) = lolo log
and Clext ilextiimage (B) = = -olog
(6)
72
Σ Eieg
EieB ea. / Pa
Heteroscedastic classification for contrastive learning, Contrastive learning datasets have been primarily constructed automatically by web scraping (Radford al., 2021; Jia et al., 2021: Zhai al., 2022; Yu et al., 2022). This process leads to noisy image-text pairs (Jia et al., 2021). This noise and the above view of contrastive learning as a massive classification task suggests naturally the use of heteroscedastic models. naive application of Eq. (2) would define two heteroscedastic classifiers with two covariance matrices
Simage-text (a) E RRN N and teettimmage (b) E E N (remember N K)
and for each batch B, the appropriate sub-matrices (with rows and columns indexed by B) would be used for the sampling of e, i.e., Dinage-rext (a) Bx and Σ Ieext-image(() (b) E TR. BxB Unfortunately, the scale of Simage-text( (a) and Etext-image (b) is not manageable, e-g., N 4B for some standard contrastive learning datasets (Zhai et al., 2022). We next develop HET-XX XL style contrastive learner where the noise is injected in the representation space (i.e., On and Bn), enabling a scaling independent of the number of classes K N.
5.2 CONTRASTIVE LEARNING WITH HET-XL
Let us focus just on text-to-image classification, as the image-to-text case is symmetric. Looking at Eq, (6). it is natural to think of the vector an E RD as being equivalent to the n-th column of W E R (remembering again that K N) in standard classifier like Eq. (1). Under this viewpoint, the image network producing the representations {ax) acts like a lookup function into the columns of W. This perspective is especially natural in the case where the image tower is fixed and not learned (Zhai et al., 2022).
In summary, we can thus mirror the standard classification setup as follows. Given a batch B of size B, we can collect the vectors {a, hes into a matrix Ag E IR DxB leading to the following (homoscedastic)
6

---

Published as conference paper ICLR 2023
Method comparison on datasets, JFT-300M trained for epochs (top), ImageNet-2Ik trained for 90 epochs (middle) and JFT-4B trained for 1 epoch (bottom). Two architectures used, 125xl (left) and a ViT-I/32 (right). NL.L stands for negative log-likelihood. For HET, r" found by grid search.
Table 3:
ResNet152x1
ViT-L/32
Method
JNLL
Prre
P#Paams
JNLL
Prec@1
4#Parass
DET
8.57
0.429
95.6M
7.83
0,468
325.3M
JFT-300M
HET T
8.39
0.440
171.5M
7.79
0.469
363.7M
HET T
8.45
0.444
171.5M
7.89
0.484
363.7M
НЕT-H
8.38
0.452
104.
7.68
0.493
327.5M
HET-XL
8.11
0.467
104.IM
7.65
0.498
327.5M
JFT-4B ImageNet-211
DET
5.79
0.461
102.9M
5.72
0.472
328,9M
НET T
5.73
0.469
193.5M
5.89
0.468
374.8M
HET 7*
5.72
0.475
193.5M
5.80
0.471
374.8M
НET-H
5.73
0.469
111.4M
5.72
0.474
331.1M
HET-XL
5.71
0,486
111.4M
5.78
0.474
331.1M
DET
5,44
0.508
118.8M
5.17
0.539
M
HET T*
5.38
0,512
241.6M
5.19
0.533
399. 399.IM IM
HET-XL
5.34
0.517
127.2M
5.10
0.544
339M
classifier for the input text b,, p(a = softmax (A ( 3.). echoing the form of Eq. (1). We can now extend this reasoning apply HET-XL to contrastive learning
p(a(ba) = sof comt (A (Bn, + <(b,))) with <(b..) ~ NO inext-image (b,)))
(8)
where, desired, Lext-imugee (b,) remains of size DxD independent of K. Table Appendix A.2 shows the shoo results. HET-X XL improves ImageNet zero-shot performance from 85.29% to 85.56%.
RESULTS
6.1 A COMPETITIVE BASELINE HET-H: HASHING APPROACH
The description of HET-XI (Section 3.1) I) has highlighted that we can improve HET by efficiently mapping a covariance defined in a lover-dimensional space the logit space. Equipped with this insight, and stress test our approach competitively as possible, we design baseline that retains the rich heteroscedastic modeling НET without making the parameters Ocov scale with respect to K.
To this end, we take inspiration from the "hashing trick" (Weinberger et al., 2009; Chen et al., 2015; Eban al., 2020), notably popularised by application to large-scale advertising problems (Chapelle et al., 2014). The interested readers can find more background about the hashing trick in Appendix A.3. Let us consider a maximum covariance dimension H, with H < K. Moreover, analogously to the "hashing trick", let h be hashing function going from f1...., K) {l...., H). Representing the hashing map via the matrix H E
with the (j. k) entry being non-zero when hi(k) = j. we obtain
(9) where, like in Section 3.1. the covariance Ex 0cow) H is in a diver-dimessionll space. The map- ping induced by the hashing function can lead to collisions, i.e., some different classes will be modeled by the same entry in the covariance matrix. In advertising problems, the collisions in features have a regular- izing effect (Chapelle al., 2014). Since H is defined via the hashing function i, only manipulate H
Be" (o (WT (( 0) + H ((())) with 6"(x) E I N(0, E" Ocov)),
7

---

Published as conference paper ICLR 2023
DET
1.4
DET HET
X
Sogh(Toin mas
13.00
НЕTHI
loga(Evvl moimeg
HET-XL
log- (Barch
logs (Barch
12
0.8
2,50
0.6
11
11
200
HET
HET -XL
0.0
10
10
logg(MC samples)
(MC samples)
IM samples)
logg(MC samples)
Figure 2: ImageNet-21k ResNet152x1 latency (left) and memory usage (right). Black areas mean that the method goes out of memory for a given number MC samples and batch size (the maximum batch size 6,144 and maximum MC samples is 1,000). Grey areas mean that HET-XL fits in memory but HET does not. HET-XL has lower memory usage than HET. The relative latency between the methods is sensitive to the number of MC samples. At lower MC samples, HET-XL is faster than HET and vice versa. is possible to reduce HET-XX XL's 's MC samples with minimal effect on predictive performance, see Appendix H.
through simple indexing operations. To make a fair comparison with HET-XL, we will take H = D, noting that if H W, we fall back HET-XL. We name this method HET-H for heteroscedastic hashing.
6.2 LARGE-SCALE IMAGE CLASSIFICATION
We evaluate HET-X XL on three image classification benchmarks: (i) Imagenet-2 which is an expanded version of the ILSVRC2012 ImageNet benchmark (Deng et al., 2009; Beyer et al., 2020) with 12.8M training examples and 21,843 classes, (ii) JFT-300M (Hinton et al., 2015; Chollet, 2017; Sun et al., 2017; Kolesnikov al., 2019) with over 300M training images and 18,291 classes and (iii) JFT-4B, an expanded version of the JFT dataset with over 4B training examples and 29,593 classes. All these datasets have previously been used to benchmark the HET method (Collier et al., 2021; Tran et al., 2022).
We evaluate on two performant large-scale image classification architectures; ResNet152 (Szegedy et al., 2017) and ViT-L/32 (Dosovitskiy al., 2020). Further experimental details and hyperparameterss are given in Appendix N. We build upon open-source implementations (Nado et al., 2021; Tran et al., 2022). The code implement HET-XL
drop-in classifier last-layer, and the scripts to replicate our ImageNet-2 results are publicly available on GitHub (https: / / thub, com/ google wwnnertainty-eaaeen The JFT and contrastive learning code not publicly released as makes use proprietary datasets.
line)..
Analysis and discussion. For both architectures on JFT-300M and ImageNet-21k, there is consistent ranking the methods, see Table 3. HET outperforms DET. HET-H outperforms HET, with one exception (ImageNet-21k ResNet152) and HET-XI outperforms is par with all other baselines.
is also important to quantify the various costs of each method. Each table shows the parameter count of the different methods. Of course, the DET baseline always has the lowest parameter count. We also choose the H hyperparameter of НЕТ-H to equalize the number of parameters with the HET-XL method. These two methods have significantly reduced parameter count compared to the HET method. This is an important consideration for real-world deployment which often has tight storage constraints (Song et al., 2020).
Parameters must be loaded into memory, impacting memory usage. While it is more difficult quantify memory usage and latency considering interactions with XLA compilation (Sabne, 2020), we follow the best practices from the literature (Zhai et al., 2022). To measure memory usage, evaluate which models
8

---

Published as a conference paper ICLR 2023
Table 4: Effect of data scaling. ViT-L/32 JFT 300M comparison of methods. See Table for epoch results.
Table 5: Effect of model scaling: ResNet1 on JFT-300M. See Table 3 for ResNet152x results.
14 epochs
28 epochs
Method
JNLL
1ree@@1
Method
JNLL
1Prec@1
JNLL
*Prec@1
DET
7.83
0,490
DET
7.61
0.487
7.51
0.496
HET
7.62
0,484
7.41
0.505
HET-XL
7.42
0.526
HET-XL
7.41
0.518
7.33
0.523
fit into memory at varying batch sizes for a series of MC samples. As is demonstrated in Fig. 2 and Fig. 5 (Appendix F), HET-XI has lower memory usage than HET and slightly lower memory usage НET-H.
The conclusions about HET-XI'' latency are more mixed. The relative latency of the methods depends the number of MC samples. There is trade-off between the cost of computing the samples from the [(x) matrix for HET and the cost of the W matrix multiplication in HET-XL. For smaller numbers of MC samples, HET-XL is faster than HET and vice versa. We provide a detailed analysis of this phenomenon in Appendix F.1. When latency is a concern, the number of MC samples in the HET-XL method can be significantly reduced with only a small impact on predictive performance, see Table (Appendix H) for a sensitivity analysis. In this regime, HET-XL is faster, uses less memory, has substantially reduced parameter count and still provides significant performance gains over HET.
Model and data scaling ablations. In Table 5, we see the results when we increase the model scale from a 15sNet152x1
ResNer152x2. While both methods benefit from model scaling, the relative gains between the DET and HET -XL methods remain roughly the same. Similarly, we scale the number of training epochs for JFT-300M Table 4 and JFT-4B Table 10 (Appendix D). Again all methods benefit from more training examples and the relative differences between methods remain roughly the same.
Learning ablation. In Table 8 (Appendix B), we run an ablation to separate the effect of learning 7 from other features of the HET-XL method. We see that learning accounts for part but not all of the gains over HET. For example when we do not learn т, HET-XL, on JFT-300M still outperforms HET by 1.6% in precision@1 1, further increasing 2.3% by learning Not only does removing the temperature hyperpa- rameter ease the adoption of HET-XL, the performance gains also empirically confirm the insight that we can learn in the large-scale setting (see Section 4). Outside of that large-scale setting, i.e., where models are more prone to overfitting by training for further epochs, we demonstrate that learning T becomes indeed less useful, see Table 9 in Appendix B. How precisely we parameterize T and the HET-XL'S insensitivity choices in this parameterization is set out in Appendix K.
7 CONCLUSIONS AND FUTURE WORK
We have developed HET-XL that consistently outperforms the DET and HET approaches in image classifi- cation and zero-shot classification from a contrastively trained model. In our large-scale settings, not only does HET-XI substantially reduce the parameter count compared to HET, it also removes the need to tune the extra temperature hyperparameter on a held-out set. As result, HET-XB XL significantly eases the deploy- ment of heteroscedastic classifiers in problems with many classes. In future research, would like to study approaches that can bypass Monte Carlo sampling, which can negatively impact the latency and memory of HET-XL,
could exploit approaches based some Taylor expansion of Eq. (2) (Collier et al., 2021) or based on mean-field approximations (Lu et al., 2020). One central question is how the simultaneous increase in bias introduced by those approximations together with reduction in variance in the estimation of the integral of Eq. (2) affect the training and end performance of the classifiers.
Figure 2. For example,
9

---

Published as conference paper ICLR 2023
Acknowledgements. measurements.
We thanks Lucas Beyer for his help with the FLOPs, latency and memory usage
REFERENCES
Lucas Beyer, Olivier J Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. Are we done with imagenet? arXiv preprint 1xXiv:2006.07159, 2020.
Christopher Bishop. Pattern recognition and machine learning. Springer, 2006.
Christopher M Bishop and Cazhaow S Quazaz. Regression with -eput-dependent noise: A bayesian treat- ment. In Advances in neural information processing systems, pp. 347-353, 1997.
Ake Björck and Gene H Golub. Numerical methods for computing angles between linear subspaces. Math- ematics of computation, 27(123):579-3 -59,494, 1973.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty neural network. In International conference muchine learning, pp. 1613-1622. PMLR, 2015.
Guillaume Bouchard.
Efficient bounds for the softmax function, applications to inference in hybrid models. In Presentation at the Workshop for Approximate Bayesian Inference in Continuous/Hybrid Systems NIPS-07. Citeseer, 2007.
Olivier Chapelle, Eren Manavoglu, and Romer Rosales. Simple and scalable response prediction for display advertising. ACM Transactions on Intelligent Systems and Technology (TIST), 5(4):1-34, 2014.
Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinherger, and Yixin Chen. Compressing neural net- works with the hashing trick. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 2285- 2294, Lille, France, 07-09 Jul 2015. PMLR. URL https: https:: chenc15 html.
/ http:s/s/ppcccednnga... mlr pre.ss.es./........e............ v37/
François Chollet. Xception: Deep learning with depthwise separable convolutions. In The IEEE Conference on Computer Vision and Puttern Recognition (CVPR), July 2017.
Mark Collier, Basil Mustafa, Ef Kokiopoulou, Rodolphe Jenatton, and Jesse Berent. A simple probabilistic method for deep classification under input-dependent label noise. arXiv preprint arXiv:2003.06778 2020.
Mark Collier, Basil Mustafa, Ef Kokiopoulou, Rodolphe Jenatton, and Jesse Berent. Correlated input- dependent label noise in large-scale image classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1551-1560, 2021.
Jean Daunizeau. Semi-analytical approximations statistical moments of sigmoid and softmax mappings 2017.
of normal variables. arXiv preprint
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 2009.
leee,
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 2001v /9999, 2020.
10

---

Published as conference paper ICLR 2023
Elad Eban, Yair Movshovitz-Attias, Hao Wu, Mark Sandler, Andrew Poon, Yerlan Idelbayev, and Miguel Carreira-Perpinán. Structured multi-hashing for model compression. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11903-11912, 2020.
Vincent Fortuin, Mark Collier, Florian Wenzel, James Urquhart Allingham, Jeremiah Zhe Liu, Dustin Tran, Balaji Lakshminarayanan, Jesse Berent, Rodolphe Jenatton, and Effrosyni Kokiopoulou. Deep classifiers with label noise modeling and distance awareness. Transactions Machine Learning Research, 2022.
URL htsps::////.
/pppneevieeeeeeeeeeeeeeeeeeeeeee ew. wwt.eeeeee/eeee................e www..neetoormiiiccmmmm
Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel pro- gramming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning, pp. 1568-1577. PMIR, 2018.
Yarin Gal and Zoubin Ghahramani. Dropout a bayesian approximation: Representing model uncertainty in deep learning. international conference on machine learning, pp. 1050-1059. PMLR, 2016.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and Sculley. Google vizier: A service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, Halifax, NS, Canada, August 13 17, 2017, pp. 1487-1495, ACM, 2017. doi: 0111443309998330988800 URL https::/ www.dcii.oooooooooooooooooooooo. / 10. 1145/ 3097983. 3098043.
Edouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, and Hervé Jégou. Efficient softmax approximation for GPUs. In International Conference on Machine Learning, volume 70, 1302-1310.
PMLR, 2017.
Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. In NIPS Deep Learning and Representation Learning Workshop, 2015. URL http: //ap::////xv....///aa/// org/ abs/ 1503. 02531.
Chao Jia, Yinfci Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision laiguale representation learning with noisy text supervision. In International Conference on Machine Learning, pp. 4904 4916. PMLR, 2021.
Alex Kendall and Yarin Gal. What uncertaintics do we need bayesian deep learning for computer vision? In Proceedings of the st International Conference Neural Information Processing Systems, pp. 5580- 5590, 2017.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. arXiv preprint 19Xiv: 13700, 2019.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncer- tainty estimation using deep ensembles. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 6405-6416, 2017.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pp. 1540-1552. PMLR, 2020.
Zhiyun Lu, Eugene Ie, and Fei Sha. Mean-field approximation to gaussian-softmax integral with application to uncertainty estimation. arXiv preprint 20xivv2006.07584,, 2020.

---

Published as a conference paper ICLR 2023
Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable gradient-based tuning of con- tinuous regularization hyperparameters. In International conference machine learning, pp. PMLR, 2016.
Dougal Maclaurin, David Duvenaud, and Ryan Adams Gradient-based hyperparameter optimization through reversible learning. In International conference on machine learning, pp. 2113-2122. PMLR, 2015.
Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, 32, 2019.
Miguel Monteiro, Loic Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Kon- stantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation networks: Modelling spatially correlated aleatoric uncertainty. arXiv preprint arXiv:2006.06015 2020.
Kevin P Murphy. Machine learning: probabilisticc perspective. MIT press, 2012.
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multimodal con- trastive learning with limoe: the language-image mixture of experts. arXiv preprint arXiv:2206.027700 027770 2022.
Zachary Nado, Neil Band, Mark Collier, Josip Djolonga, Michael W Dusenberry, Sebastian Farquhar, Qix- uan Feng, Angelos Filos, Marton Havasi, Rodolphe Jenatton, al. Uncertainty baselines: Benchmarks for uncertainty & robustness in deep learning, arXiv preprint arXiv:2106. 040055, 2021.
lan Osband, Zheng Wen, Mohammad Asghari, Morteza Ibrahimi, Xiyuan Lu, and Benjamin Van Roy. Epis- temic neural networks. arXiv preprint 2 200.x:::0000000 08924, 107.089444 2021.
Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In International conference on machine learning, pp. 737-746. PMLR, 2016.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International Conference on Machine Learning, pp. 7748-8766 PMLR, 2021.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto. Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583-8595, 2021.
Amit Sabne. Xla: Compiling machine learning for peak performance. 2020.
Aravind Sankaran, Navid Akbari Alashti, Christos Psarras, and Paolo Bientinesi. Benchmarking the linear algebra awareness of tensorflow and pytorch. arXiv preprint 1xXi::2202.09888 2022.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances neural information processing systems, 25, 2012.
Liuyihan Song, Pan Pan, Kang Zhao, Hao Yang, Yiming Chen, Yingya Zhang, Yinghui Xu, and Rong Jin. Large-Scale Training System for 100-Million Classification at Alibaba. In 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD). 2020.
Xingyou Song, Sagi Perel, Chansoo Lee, Greg Kochanski, and Daniel Golovin. Open source vizier: Dis- tributed infrastructure and api for reliable and flexible blackbox optimization. In Automated Machine Learning Conference, Systems Track (AutoML-Com Conf Systems), 2022.
12

---

Published as conference paper ICLR 2023
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, pp- 843-852, 2017.
Christian Szegedy, Sergey Ioffe. Vincent Vanhoucke, and Alexander A. Alemi. Inception-v4, inception- resnet and the impact of residual connections learning. In Proceedings of the Thirty-First AAAI Con- ference on Artificial Intelligence, AAAI'1,, pp. 4278-4284 AAAI Press, 2017.
Kenneth E Train. Discrete choice methods with simulation. Cambridge university press, 2009.
Dustin Tran, Jeremiah Liu, Michael W Dusenberry. Du Phan, Mark Collier, Jie Ren, Kehang Han, Zi Wang, Zelda Mariet, Huiyi Hu, et al. Plex: Towards reliability using pretrained large model extensions. arXiv preprint arXiv:2207.074// 2022.
Eric A Wan and Rudolph Van Der Merwe. The unscented kalman filter for nonlinear estimation. In Proceed- ings of the IEEE 2000 Adaptive Systems for Signal Processing, Communications, and Control Symposium (Cat. No. 00EX373), pp. 153-158. Icee, 2000.
Kilian Weinberger, Anirban Dasgupta, John Langford, Alex Smola, and Josh Attenberg, Feature hashing for large scale multitask learning. In Proceedings of the 26th annual international conference on machine learning, pp. 1113-1 120, 2009.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume MIT press Cambridge, MA, 2006.
Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Scyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. arXiv preprint arXiv:2205.00 20t.:2995.119 2022.
Zhuoning Yuan, Zhishuai Guo, Xiaotian Yu, Xiaoyu Wang, and Tianbao Yang. Accelerating deep learning with millions of classes. In European Conference on Computer Vision (ECCV), pp. 711-726, 2020.
Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, et al. The visual task adaptation benchmark. 2019.
Xiaohua Zhai, Xiao Wang. Basil Mustafa, Andreas Steiner, Daniel Keysers, Alexander Kolesnikov, and Lucas Beyer. Lit: Zero-shot transfer with locked-image text tuning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 18123-18133 2022.
Xingcheng Zhang, Lei Yang, Junjie Yan, and Dahua Lin. Accelerated training for massive classification via dynamic class selection, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
13

---

Published as a conference paper ICLR 2023
A
RELATED WORK
Our methodology connects with several lines of research, as described in the following paragraphs.
Uncertainty quantification. Total uncertainty can be divided into epistemic/model uncertainty and aleatoric/data uncertainty (Kendall Gal, 2017). The neural network uncertainty quantification literature has primarily focused on epistemic uncertainty (Blundell et al., 2015; Gal & Ghahramani, 2016; Lakshmi- narayanan et al., 2017; Maddox et al., 2019). Aleatoric uncertainty has the interesting property that is irreducible, meaning that even in the large-scale settings such those considered in this paper there may be considerable aleatoric uncertainty. Heteroscedastic models are a popular method to model aleatoric un- certainty in regression (Bishop & Quazaz, 1997; Williams & Rasmussen, 2006), classification (Kendall & Gal, 2017: Collier et al., 2020; 2021; Osband et al., 2021; Fortuin et al., 2022; Tran et al., 2022) and image segmentation (Monteiro et al., 2020). Prior work on heteroscedastic classification has not been scaled up problems of the size of JFT-4B and would require restricting the covariance matrix parameterization, for ex- ample to the diagonal, to enable scaling. Additionally prior work, either ignores the temperature parameter leading to reduced predictive performance or requires an expensive grid search over the temperature.
Gaussian-sigmoid and Gaussian-softtax integrals. A central component of our methodology lies in the computation of either Gaussian-sigmoll or Gaussian-softmal integrals (Daunizeau, 2017). In both cases, analytical forms exist for those integrals and we need to resort to approximations. Those integrals play a central role in Bayesian inference (Bishop, 2006) and probabilistic machine learning (Murphy, 2012). For heteroscedastic classifiers, Kendall & Gal (2017); Collier et al. (2020; 2021) consider estimation based on Monte Carlo sampling, which is feasible and practical for (say, with several hundreds to thousands samples) since the operation takes place at the final layer, requiring only a single forward pass. In a different context, for the modelling of ensembles from a single model, Lu et al. (2020) propose to use mean-field approximations that they benchmark against numerical approaches such as the unscented Kalman Filter (Wan & Van Der Merwe, 2000). Variational bounds have also been proposed in the context of variational inference, see Bouchard (2007) and references therein. Daunizeau (2017); Collier et al. (2021) further consider approximations resulting from Taylor expansion.
To the best of our knowledge, only the approaches based on Monte-Carlo sampling and Taylor expansion have been successfully exploited for heteroscedastic classifiers where the large-scale covariance is learned during training, unlike, c.g., Lu et al. (2020).
Extreme classification. Previous work has looked into classification problems with millions of classes (often referred to extreme classification or massive classification). Zhang et al. (2018) found empirically that the class probabilities concentrate on a few classes only, which are called "active classes". Motivated by this observation the authors propose an efficient hashing-based implementation of "selective softmax", which omits the non-active classes from the softmax computation. Yuan et al. (2020) proposed an efficient training framework handle extreme classification tasks based on random projections. The idea consists in training first random projected softmax classifier and then recover the original classifier. Alonog the lines of approximating full softmax, Song et al. (2020) introduce a large-scale training system for 100M classification at Alibaba. The proposed approach for accurately approximating the full softmax is pretty involved, including building a k-NN graph on W for fast retrieval of active classes, a specialized com- munication strategy (including gradient sparsification) and fast continuous convergence strategy for the end-to-end system. Finally, there exist also frequency-based methods (see e.g., Grave et al. (2017) and references therein) assuming that the class distribution is imbalanced and focusing the most frequently occurring classes. Although this may be the case in language modeling context, these methods are not suit- able for problems where all classes important (e.g., in face recognition). We emphasize that HET-XL
14

---

Published as conference paper ICLR 2023
can be easily combined with the methods above, unlike НET, which requires custom solutions enable the matrices defining (xx) to fit in memory.
A.1 MORE FORMAL BACKGROUND ABOUT HET
this section, we derive the form of the HET classifier, in particular, justifying why we introduce a temper- ature parameter and why which we reproduce here for completeness.
MC sampling. Our exposition follows that of Collier et al. (2020; 2021),
Following the generative process typically employed in the econometric (Train, 2009), Gaussian-process (Williams & Rasmussen, 2006) and noisy-label (Kendall Gal, 2017) literature, we consider a utilitty vector (xx) E Rx and (zero-mean) stochastic term e(x: R I (z)) that the probability p(c/z) for the class c EE EI, {I....,K] K} is defined via
arg max (u(xx) e(x)) } sE(I....K)
1 ce==
p(c/x) = (((
with e, E 10 denoting the one-hot vector with one at the j-th entry and zeroes elsewhere. In words, the probability for corresponds how often the utility, in expectation with respect to the stochastic perturbation, is maximal for c.
classical, inpu--nndependent instantiation for (()) is i..d Gumbel distribution, thus recovering the standard so soft ftmax with logits given by uxx)) (Train, 2009). HET considers instead x(x) No, [(r)), particular extending the "identical" and "independent" properties of the Gumbel distribution. We have
(c)r()
= ((.)))N(0.E((a))
1 c arg max ej (u(x) + c(())
je(l...,kk
1 (u(2)) + +((x))
= (0.E())
lim e softmax r-0
(((u())+ +c(()))
2 (((3))N((03(()))
e eoftmaa
for some small >>> 0
(u(x) + +<((a)) for S samples <'(x) M/0, E(x)).
& 5 e softmax
The last two lines show how the temperature controls the approximation of the actual arg max while the Monte Carlo sampling approximates the woffmaa--Gausiii integral with S samples.
A.2 CONTRASTIVE LEARNING
We evaluate HET-XL added to LiT style contrastive learner (Zhai al., 2022). During training, HET-XL XLL takes just one MC sample. At test time, the deterministic logits (using no MC samples) are used for zero-shot
Table 6: HET-XL LiT vs. DET LiT zero-shot accuracy + 1 std. error ImageNet, CIFAR-100 and Oxford- III Pets datasets. Average over 3 random seeds. Ip < 0.01 in a two-sample unequal variance t-test.
Method
ImageNet
CCIFAR-100
10ford-IIIT Pet
DET HET-XL
85.29t (1:0.0242)
83.39 (=0.117)
97.87 (+:0.722)
85,561 (=0.027)
83.41 (=0.101)
97.75 (+00.71)
15

---

Published as a conference paper ICLR 2023
Table 7: Comparison of different placements of the normalization operation in contrastive learning with
НET -LL. MmngeNee/CIFARR-000Ooforr IIIT Pet zero-shot accuracies are shown.
Method
Post-noise normalization=Fall Post-noise normalization=TTuu
Pre-noise normalization-Faale
82.22/81 148/95.099
84.07/82.55/95.55 84.65/82.56/97.63
Pre-noise normalization-TTue 84.37/82.81/977.2
classification. This test procedure combined with the single MC sample during training means the cost of adding HET-XL head is negligible compared to the overall network. In Appendix M, which justifies this choice by demonstrating no change in the zero-shot accuracy of HET-XL contrastive learner when using MC sampling versus deterministic representation at test time.
ablation
For our experiments, we use a ViT-g text encoder network and ViT-G image encoder (Dosovitskiy et al., 2020). The setup is trained on 256 v3 TPU cells (512 cores) for 18 billion examples with batch size of 16,384, on the same data as the original LiT model. Further experimental details are given in Appendix N. In Table 6, we show the zero-shot ImageNet, CIFAR-100 and Oxford-IIIT Pets classification accuracy for HET-XL versus a baseline DET model. We see a small but significant gain on the primary ImageNet accuracy metric and no significant difference in the CIFAR 100 and Oxford- -IIT Pets accuracies. To contextualize the magnitude of the gains, we note that when the same DET model is trained for 4 billion rather than 18 billion examples the zero-shot accuracy 84.98/83. 14/97.57 for ImageNet, CIFAR-100 and Oxford-IIIT Pet respectively. The gains we see from HET-X on ImageNet are similar in magnitude to increasing the number of training examples shown to DET by a factor of 4.5.
To clarify the computation of the p-values: we compute 3 separate p-values using 3 separate two-sample un- equal variance t-tests. We compute a p-value comparing the DET vs. HET-XL Imagenet zero-shot accuracy, a separate p-value for the CIFAR-100 results and similarly for Oxford-III Pet. Only the ImageNet results are significant at the 0.01 level. Note however that the other two computed p-values are well above this level of significance. For elarity we report here the CIFAR-101 100 p-value: 0.839 and Oxford-IIIT Pet p-value: 0.596.
In Table 7 compare the effect of different placements of the normalization operations in HET-XL for contrastive learning. Pre-nornalization-TTne, means that the normalization operation is applied to the de- terministic representation outputted by the image text tower. Post-normalizatioo==Tnee means that the normalization operation is applied after the sampled noise vector has been added. We train with a ViT-L/32 text tower for 2 billion examples. We observe that it is crucial to add the noise term after normalization, however re-applying normalization after the noise has been added does not have a consistent effect. In the full scale experiments, Table 6 we apply normalization before adding the noise but not after i.e. pre- normalization=TTue and post-normalization-Falee.
A.3 BACKGROUND ABOUT THE "HASHING TRICK"
It can be cumbersome to deal with the one-hot, or bag-of-word. encoding of features that have many cate- gorical attributes. A case in point is large-scale advertising problems (Chapelle et al., 2014) with categorical features describing the large sets of publishers and advertisers. In such settings, we need to carefully book- keep the different categories based on some associative data structures and it is difficult to control the space taken by the resulting feature representations (i.e., Σ C2 where Cy stands for the number of categories for the feature of and F is the set of all the features).
16

---

Published as a conference paper ICLR 2023
The "hashing trick" (Weinberger et al., 2009; Chen et al., 2015; Eban al., 2020) was developed to provide a simple-to-ig tommleeeeen,, robust', fast and space-efficient technique to deal with the above setting. The "hashing trick" considers
A user-defined parameter H explicitly controlling the space of the representation, with possibly
H Erer Cj. Here, H will correspond to the number of hash buckets.
A hashing function hash_f fn mapping a pair of feature-value (y, to (1,.... H).
The H -i-diennsional feature representation & E IR" of the raw data {(f.u)))ser is obtained by the follow- ing simple procedure (see Section 4.3 in Chapelle et al. (2014)):
Initialize +b- = O.
For f E F, update 2 =b + with j = hash.fn((f.
Depending on the choice of H. the hashing trick operates as a dimensionnllty-redutii technique. It leads collisions wherein two different feature-value pairs (f, uy) and (f", uy.) end up with the same entry in the feature representation. In the example of large-scale advertising problems, the collisions generally act as regularizer (Chapelle et al., 2014). This technique works well in practice and has been widely adopted.
"Hashing trick" for HET. Equipped with the previous insights, we can use the hashing trick to have a fast and space-efficient mapping from the space of labels R to RR We can then define the covariance matrix into the resulting lower-dimensional space R", which is precisely what HET in Section 6.1 does. Interestingly, in the context of HET H, collisions have different meaning: If two classes and c' are mapped to the same entry in I H then they will share their covariance structure. Our experiments with H < K showed that those collisions did not seem to be harmful in our settings.
B LEARNING 7
We ablate the effect of learning T in the HET-XL method and also evaluate whether adding a learned tem- perature parameter to DET and HET makes up the gap with HET-XI, see Table We observe that on JFT-300M and JFT-4B all methods benefit from having a learned temperature parameter, however the gains from HET-XL can only be partially explained by learning т. On ImageNet-2Ik, where we train for 90 epochs, only validation-set performance.
sees a benefit from learning
opposed to applying grid search based on the
In the main paper, argue that 7 can be learned in the large-scale settings considered in the experiments. By "large-scale" settings, we mean that we are in a regime whereby (i) the data scale is large relative to the model capacity and (ii) the training schedule, typical at the scale, entails few epochs. Overfitting is therefore unlikely to happen.
Above, we observed ImageNet-22 21k, with models trained for 90 epochs, that learning only has positive effect for HET-XL and not the other methods. We hypothesise that this may indicate that this setup is already the limit of the large-scale setting we described.
We evaluate our hypothesis that learning r helps more when we can only train for small number of epochs by comparing training ImageNet-21k for 7 and 90 epochs, Table 9. We observe, expected, that we train for longer the gains from learning 7 reduce.
For example, the hashing trick can easily deal with new features and it can handle any feature type that is hashable. Sometimes, signed variants of the hashing trick are used (Weinberger al., 2009).
17

---

Published as conference paper ICLR 2023
Table 8: Ablation on the effect of learning T. all results for 52sNet152x1 models.
JFT-300M
ImageNet-21k
JFT-4B
Method
JNLL
Pree
JNLL
Prec@1
JNLL
1Prec@1
DET
8,57
0.429
5.79
0,461
5,44
0.508
w/ learn
8.44
0.449
5.82
0.458
5.39
0.514
HET
8.45
0.444
5.72
0.475
HET w/ learn
8.28
0.461
5.72
0.475
HET-XL, w/o learn
8.15
0.460
5.72
0.475
5.37 5.34
0.513
HET-XL
8.11
0.467
5.71
0.486
0.517
Table 9: Evaluating the effect of learning т as we reduce the number of training epochs on ImageNet-21k with ViT-L/32 HET-XL model.
Method
JNLL
Pree
90 epochs grid search r
5,80
0.471
90 epochs w/ learn 5.78
0.474
epochs grid search
7.988
0.232
epochs w/ learn
7.507
0.274
C MODEL SCALING ABLATION
We evaluate whether gains from the HET-XL method persist when the base model scale doubled. We see in Table that when the base model scale is increased from ResNet152x1 to ResNet152x2 the relative gains from the HET-X. XXL method on JFT-300M are similar to those observed earlier with a ResNet152x1 Table 3.
DATA SCALING ABLATION
In Table 4 (main paper) and Table 10, we see that increasing the number of training epochs has little effect on the relative gains from the HET-XL method, while increasing the performance of all methods. We scale training from 7 to 14 and 28 epochs JFT-300M with the ViT-L/32 architecture, Table 4. Similarly scale the training epochs from to and epochs on the JFT-4B dataset with ResNet152x1 architecture, Table 10.
Table 10: Effect of data scaling: ResNet152x1 JFT-4B comparison of methods. We report test set negative log-likelihood and precision@ I.
epoch
epochs
epochs
Method
JNLL
1Prec@1
JNLL
1
JNLL
tPrec@)
DET
5,44
0.508
5.38
0.518 0.523
0.515
5.35
HET-XL
5.34
0.517
5.28
5.25
0.526
18

---

Published as conference paper ICLR 2023
8.0
log:(Train (Train GFL (PLOPs/igg)
HET
(oga(Eval GFLOPs/img)
66.25
HET
7.75
HETHH
HET-H
7.50
HET--L
6.00
HET-XL
7.25
$.50
7.00
5.25
6.75
6:50
(Mg samples)
IMC2 samples)
Figure 3: JFT-4B ViT-L/32 latency in terms of FLOPs.
8.5
DET
DET
(TTuin GPL GFLOPw/men
8.0
HET X
(Eval GFLOPs/img)
X
HET
HET-H HET-XL
HET-H
7.5
6.0
HET-XL
5.5
6.5
4.5
6.0
8
(MC samples)
(MC samples)
Figure 4: ImageNet-21k ResNet152x1 latency in terms of FLOPs.
FLOPS METRIC
In the main paper, we report training and evaluation ms/img as the latency metric, Fig. 2. FLOPs is another popular metric evaluating the computational cost of neural network models. In Fig. 3 and Fig. 4, we report the training and evaluation FLOPs for a ViT-L/32 on JFT-4B and a ResNet152x1 on ImageNet-2 respec- tively. We see a similar trend as for ms/img. At high numbers of MC samples, HET-XL higher FLOOPs versus the key baseline, HET. However the trend is reversed at lower numbers of MC samples. Table of Appendix H, we will see that predictive performance HET-X0 XXLL not highly sensitive to the number of MC samples.
a significantly
FURTHER LATENCY AND MEMORY RESULTS
the main paper, we only report memory usage and latency metrics in the ResNet152x1 ImageNet-21k setting, Fig. 2. In Fig. 5 and Fig. 6, additionally report memory usage and latency in the ViT-L/32 JFT-4B setting. We see similar trends as for the ImageNet-21k setting, with HET-XL memory usage less that HET and the ordering of the methods dependent the number of MC samples.
19

---

Published as conference paper ICLR 2023
logy(Bushh slze)
logn/Bach slne)
logs/Bush size)
HET-H
HET-XL
12
12
11
log:MMC samples)
log:(MC samples)
(Mg2(MC
Figure 5: JFT-4B ViT memory usage. A black rectangle means that for the given number of MC samples and batch size the method goes out of memory. Plotted in grey are the cells where HET-H and HET-XL fit in memory but HET does
DET
3.00
HET
HET
log:( (Train me/mg!
HET-H
(oga(Eval ma/img)
0.8
HET-H
2.75
HET-XL
HET-XL
2.50
2.25
0.2
2.00
1.75
0.
(Mg samples)
IMC samples)
Figure 6: JFT-4B ViT-L/32 latency.
20

---

Published as a conference paper ICLR 2023
F. AN ANALYSIS OF THE COMPUTATIONAL DEPENDENCY ON THE NUMBER OF MC SAMPLES
As previously observed, whether HET-XL has better latency than HET depends on the number of MC samples used. We next study this observation by deriving the time complexity for both HET XL and HET, showing that the comparison depends on trade-off with respect to K, D, R (as defined in Section 2.2) and S, which stands for the number of MC samples.
Let define the dimension Q equal K in the case of HET and equal D, the dimension of the pre-logits, otherwise. With the notation from Section 2.2, we recall that o(x) E RP and E IR RxQ where stands for the number of factors of the low-rank parametrization. For both HET-XL and HET and number S of MC samples, the noise matrices (x) E RR5K and <s(2)) E RSxD are generated according to (we omit all the bias terms for simplicity)
(1s v)x))) (ZJ) + (Is ((x)) T K1) (z 12) with v(xr)= = K2 ((z) E R
and
z ERSS I),
Z E RSSR ~ N(0. J).
Ki E DXQ,
K2 DXQ,
Once the noise matrices (x) and ca)(a:) have been generated, the logits for the S MC samples are obtained by computing
(5)(2) W + 6s(x)) E R and (IsØ(æ)) T + <>( W e RSK
(10)
for respectively HET and HET
In Table 11, we derive step by step the resulting time complexity for both HET and HET and show that we recover the experimental observation that HET-XI XL is more efficient than HET for small numbers of MC samples.
Operations
Complexity HET
Complexity HET-XL
v(z)
O(DK) O(KS + KRS)
O(D2) O(DS DRS)
HET
(1s n(x) )o (ZJ) (1s o(x) (x
1.0
HET-XL
O(DK + KS)
O(D2 + DS)
Logits Eq. (10)
O(DK + KS)
O(DS DKS)
Cmmplexity
0.8
Time complesity
Total
O(DK + KRS) O(DRS + D2 + DKS)
0.6
Table 11: Time complexity analysis of HET and HET-XL. For the total complexity in the last row, only keep dominating terms, without constants if repeated multiple times (assuming K > 1, S 1 and R > 1). We additionally plot on the right figure the non-simplified time complexities of HET and HΣT-XL. instantiated with the values of the ResNet1 52x1 setting on ImageNet-21k. We do observe the phenomenon wherein HET-XL is more efficient for fewer MC samples. The crossing point is not exactly the same as that in the real experiments because of the big-O analysis and some difficull-to-mont modellll linear-algebra optimization due to XLA (Sankaran el al., 2022).
0.4
0.2
0.0
IMg sumplex)
We note that this time complexity analysis agrees with the experimentally observed training times computed our profiling. In particular Table 12 shows that on ImageNet-21K with ResNet152x1 the training time for HET-XL is faster than HET up to 16 MC samples, with HET being faster from 64 MC samples on. Note that there are several points at which HET-XX -X dominates HET c.g. HET-XL has better Prec@1 I and NLL at I MC sample than HET at 1000 MC samples despite also being substantially faster train. Training times computed from by multiplying the profiler train ms/img by the number of examples seen during training, as commonly done to benchmark large-scale vision models (Riquelme et al., 2021).
21

---

Published as conference paper ICLR 2023
Table 12: TPU training time vs. performance for varied MC samples for HET and HET-XL on ResNet152x1 ImageNet-2Ik.
MC Samples
4
16
64
256
512
1000
HET Prec@1
0.464
0.433
0.462
0.466
0.469
0.472
0.475
HET NLL
5.80
6.00
5.77
5.75
5.73
5.73
5.72
HET TPU Training Hours
1217
1221
1230
1263
1547
1666
1786
HET-XI, Prec@1
0.476
0.478
0.476
0.480
0.485
0.484
0.486
HET-XI. NLL
5.69
5.68
6.12
5.99
5.79
5.70
5.71
HET-XI TPU Training Hours
1171
1181
1205
1316
1745
2244
3278
G
TUNING OF
In this section, we discuss further details and results of Section 4.
G.1 DETAILS ABOUT THE BENCHMARK
We consider the tuning of HET-XL for ResNet50x2 model over ImageNet-21k with a training duration of 7 epochs. The hyperparameters of the model are those from Table in Dosovitskiy et al. (2020).
To define our 3-fold split of the dataset, we take the standard validation set as the points), and extract from the training set a validation set (also with 102,400 points). All the methods in the benchmark tuned minimize the validation negative log likelihood."
(containing 102,400
The grid search uses the following list of temperatures E {0.05, 0.1, 0.2, 0.4, 0.8, 1.5, 3.0,5.0) 5.0) which cover typical perfooerfor values mentioned in previous work (Collier al., 2021).
The Bayesian optimization approach follows the default configuration from Song et al. (2022). We run a total of 30 training jobs, divided into 6 sequential steps of 5 parallel training jobs.
As far as the -ET-EE---EEE XL-reelt-dd hyperparameters are concerned, we set the range of as [0.05, 5.0]. consider 50 factors and use 1000 MC samples, as done in the rest of the paper.
G.2 THE SPECIFIC STRUCTURE OF THE HYPERPARAMETER T
The temperature has the property to explicitly appear both at training and validation time, unlike optimisation-relatee or regularisation-reltte hyperparameters that surface only training time. As we next see, this observation has implications about the way we can compute gradient estimator to tune
To ease the exposition, 1/T [æz θ) IR with one-dimensional output, e.g., in regression task. The parameters of the model are summarised in E I... Moreover, let ..) and (x.)) be two loss functions, defined respectively for training data (x, y) Purin and validation data (x, y) Pval.
us momentarily simplify our setup. We consider a temperature-scaled model
We typically seek a temperature parameter that minimises some validation objective a bi-level optimisa- tion formulation (Franceschi et al., 2018)
((o;
min (T)(T) with Fva(T) = ((m..))~ppa +50
(11)
'There is exception: We also check in Table 13 the effect of applying grid search and Bayesian optimization directly to optimize for the test metric. Those variants referred "2-fold".
22

---

Published as a conference paper ICLR 2023
for some parameters (7) defined as a solution of the training problem
Curin ( 0),
(T) E arg min Fuain (T, θ) with Frain (T, 3) = ((x,))) Prrin
(12)
GERP
By applying the chain rule, we can compute the gradient of the objective in Eq. (11) for given pair (æ, of the expectation, leading to:
@Cvaly's y)
(+) { - O())) (Ve 00 (, (.O
(13)
Oy
In the expression above, the most complex term to handle is V V7) since it requires to differentiate through the solution of another optmization problem, notably involving higher-order derivatives (Maclaurin et al., 2015; Luketina et al., 2016; Pedregosa, 2016; Franceschi et al., 2018: Lorraine et al., 2020).
Several approaches have been proposed to approximate this term. For example, Luketina et al. (2016) only consider the hyperparameter influence through the current gradient step solving Eq. (12). This formally translates into (T) x - V Ve 2x)] with st and θ being respectively the learning rate and iterate produced step / by iterative algorithm solving Eq. (12), e.g., SGD. In our comparison, we will take Luketina et al. (2016) as a representative approach for efficient gradient-based hyperparameter
tuning.
A simple gradient estimator. A remarkable property of the temperature is that Fval in Eq. (11) depends on not only through (r( it would be for example the case for weight decay- also directly via I/T. While in absence of this property, the gradient of Fval would solely rely on (T)" in the case of T, we can assume that (r) = (i.e., the optimal O (r) mildly varies with 7) and focus on - f(; (T))) for the gradient. Interestingly, further approximating O (T) ~ O₁ the resulting estimator extremely simple and cheap to compute, e.g., no higher-order derivatives to compute (in Table 13, we refer to this approach as "Simple gradient"). Moreover, we can notice that, provided that Larain Lval (in our comparison, both are set to the negative log likelihood) and the training and validation sets are identical, 3 (x; θ,) falls back to the gradient we would take by treating standard training parameter by directly solving min, Firain(T. θ).
G.3 ADDITIONAL OBSERVATIONS
In addition to the results reported and discussed in Section 4, make the following observations from Table 13:
In the training regime focusing on, it appears that applying grid search and Bayesian op- timization to directly optimize the test likelihood ("2 fold") does not lead to significantly better performances.
Retraining the best models obtained by grid search and Bayesian optimization ("+retrain") typi- cally comes with drop in performance, which can be explained by the mismatch of the selected temperature when the training set is enlarged.
The coarser approximation made by "Simple gradient" compared with Luketina et al. (2016) leads to worse likelihood.
The approach of Luketina et al. (2016) suffers more than "Simple gradient" when we approximate the true validation set by fraction of the training batch (here, 12.5% and 25%). may be explained by the fact that "Simple gradient" reduces to "Simple training" in the case where the validation and training sets coincide. Note that, when 12.5% and 25% of the training batch are held-out as validation data, the model effectively sees fewer training points. We account for this effect by running for respectively 12.5% and 25% more epochs.
23

---

Published as conference paper ICLR 2023
Table 13: Comparisons (means + standard errors) of methods to tune for a ResNet50x2 on ImageNet-2Ik.
Method
JNLL
PPrec @
Single training?
No validation set needed?
Grid search
6.08
0,433
Grid search (2 folds)
6.11 a: 0.08
0,434
0000
Grid search (+ retrain)
6. 0.05
0.434
0001
X
Bayesian optimization
6.03 #: 0.06
0.428
0.000
Bayesian optimization (2 folds)
6.05 0:07
0.435
0.001
X
X
Bayesian optimization (+ retrain)
6.09 0.05
0.429
0.001
X
Luketina et al. (2016)
6.08 0.04
0,425
00000
Luketina al. (2016) (12.5%)
6.19 0.07
0,423
Luketina (2016) (25.0%)
6,20 0.00
0,420
0.001
Simple gradient (val)
6.22 #:
0,432
Simple gradient (12.5%)
6.18 1: 0.04
0.432
0.001
Simple gradient (25.0%)
6.16 0.06
0.429
0.001
Simple training (as NET-XL)
5.99 0.04
0.435
0.001
Table 14: Sensitivity analysis of HET-XL to the number of MC samples on ImageNet-21k with ResNet 15x1 architecture.
MC Samples
16
64
256
512
1000
Prec@)
0.476
0.478
0.476
0,480
0.485
0.484
0.486
NLL
5.69
5.68
6.12
5.99
5.79
5.70
5.71
H HET-XL SENSITIVITY TO THE NUMBER OF MC SAMPLES
We have seen above that the latency of the HET-XL method is sensitive to the number of MC samples used at training and eval time. In Table 14, we conduct a sensitivity analysis of the method's predictive performance to the number of MC samples. We see that while, expected, more MC samples results in broadly better predictive performance, the number of training and evaluation samples can even be reduced for -ET and still outperform HET with 1,000 MC samples.
I EFFECT OF SHARING W PROJECTION WITH NOISE SAMPLES
Three possible explanations for HET XL's outperformance of HET are (i) learning (ii) knowledge sharing due to the use W to project both the pre-logits and the noise to logit space and (iii) the regularizing effect of vastly reducing the parameter count. We have already seen, Table Appendix B, that learning the temperature only accounts for part of the gains. Below, Table 15, we run an ablation to test (ii) whether the W sharing has a significant effect HET-XL's performance. When do not share the W, instantiate a separate independent randomly initialized matrix W which is used solely to transform the noise samples logit space; formally.
p(ylar) = E. (a (W ((x;0) (W')) with é'(x) E M and W e MXK
(14)
and M # D.
24

---

Published as a conference paper ICLR 2023
Table 15: Effect of sharing W projection between the pre-logits ((z; 0) and noise samples c(x) Evaluated on ImageNet-2Ik with ResNet152x1 architecture.
Method
JNL
1Prec@1
HET XL w/ W sharing HET-XL XL w/o W sharing
5.71
0.486
5.62
0,487
Table 16: Sensitivity to M of 1esNet152x1 ImageNet-21k performance when not sharing W.
M 64
128
256
512
1024
2048
Prec@1 1
0,465
0.468
0.472
0.477
0,484
0.487
NLL
5.74
5.68
5.66
5.66
5.74
5.62
the setting with a ResNet1 on ImageNet-21k,, we can see that sharing W has a small, anything slightly negative effect on performance. Hence, we conclude that the primary drivers of HET-XL's im- proved performance are learning т and the regularization due to HET XL's s reduced parameter count. For the HET-X XL method presented in the other experiments throughout this paper, we opt to share W due to the minimal performance difference compared to not sharing W as well as the added simplicity and efficiency of the method as result.
HET-XL ties the dimensionality of e'(x) to the dimensionality of the last layer of the base network, D. Following Eq. (14), this is not strictly necessary and we could allow the dimension of e' (xx to be independent D. In Table 16 we vary precisely M. Note that M is different from R, the rank of E' (ax; fewv). We observe that indeed M by a factor of 8 from D down to M = 256, and still have performance better than HET.
increase M to D we improved Prec@1 I and NLL. Note however that
decrease
L.1 ANOTHER MOTIVATION FOR THE SHARING OF W
We provide another observation that further motivates why sharing the same matrix W for both p( 0) and é' (æ), as done by HET in Eq. (5), is natural design decision
Let us write the rank-R covariance matrices for both HET and HET-XL
Σ Ber (x) = Q(x) Q(x) with (
and
ΣEet (x) W P(x) P(x)W with P(x) E RXD W E Rxx A sufficient condition for HET-XL to be able to express a covariance similar to that of HET is to have
Q(x) W TPx))
(15)
other words, the span of the rows of W-a subspace of æ also referred to as row space W or rowsp(W) (W) for short-is well aligned with rowsp(Q (.),
find matrix P ((x) to satisfy the above
relationship.
We test the hypothesis that the W learned by HET naturally aligns with rowsp(Q(z)), which would justify our design choice for HET-XL. To this end, we consider the HET ResNet152x1 model trained on ImageNet- 21k. To quantitatively measure the alignment of rowsp(W) (W) and (owsp(Q(=r)) we the smallest principal angle (SPA) between those two subspaces (Björck & Golub, 1973). We will manipulate the SPA via cosine, i.c., cos(aspa (z)). Computing the SPA requires compute the SVD of Bq(z)) with B. denoting orthonormal basis for rowsp(W) and rowsp (Björck & Golub, 1973).
25

---

Published as conference paper ICLR 2023
Table 17: Measurement of the alignment between rowsp(W W) and rowsp(Q(z)) for HET. Evaluated on ImageNet-2Ik with 1esNet152x1 architecture. The table reports the average and the standard deviation of the (cosine the) smallest principal angle (SPA) over the validation set. A higher cosine implies a smaller angle and better alignment of rowsp (W) and (owsp(Q(x))) Here, have arccos(0.829)) 0.593 arccos(0.350) = 1.213.
Setting
Ex (c08( (aspx (x))]
With W learned by HET
0.829 + 0.042
With random W
0.350 0.002
Table 18: Ablation on the effect of the dynamics of T throughout training on predictive performance. We first train HET model and then train new model from scratch fixing the temperature to be the final value of the learned temperature from the first training run. Evaluated on ImageNet-21k with a R5sNet152x1 architecture.
Method
JNLL
Prec 1
HET-XL
5.71
0.486
HET XL w/ 7 fixed to final value of learned temperature
5.74
0.486
We compute cos( (os(Asp (a:) for each image of the validation and report in Table 17 the resulting mean and standard deviation. To put in perspective the value obtained in the case of the learned W by HET, we also compute the same quantities in the case randomly generated Gaussian W -O(00, 1). We can that in the case of HET, the SPA is significantly smaller (i.e., its cosine significantly larger) than in the case of random W, with arccos(O. (0.829) 0.593 (0.305)) 1.213. Compared with the random case, the learned W is thus such that zowsp (W) and rowsp (x)) are well aligned, as required by Eq. (15).
J T DYNAMICS
Fig. 7 shows the evolution of T during training on ImageNet-21k with a ResNet152x1. is interesting to note the non-monotonic dynamics of r. In Table 18, we take the temperature at the end of this training run, - = 0.50940 and re-run training from scratch with this fixed temperature value. The performance of the methods are similar, suggesting that the benefits from learning come primarily from finding a very fne-grained value of to train with and not that can take different values over time. Finding such specific values T would only be possible with an expensive fine-grained hyperparameter search, which is infeasible in the settings considered in this paper.
K T PARAMETERIZATION AND SENSITIVITY TO STARTING POINT AND RANGE
When learning the temperature parameter, we parameterize 7 as
T()) (Tmax Tiin) s sigmoi d()) + Trin
(16)
where is an unbounded trainable parameter to be learned and Tmin and Tmax are hyperparameters defining the valid range of minimumn and maximum possible temperatures. For all experiments in the main paper, we use Tmin = 0.05, Tiaa 5.0 and initialize t such that the initial 7o is set to the midpoint of [0.05, 5.0).
i.e. 2.525.
26

---

Published as conference paper ICLR 2023
2,50
2.25
Temperature
2.00
1.75
1.25
1,00
0.75
0.50
0.0
0.4
0.6
0.8
1.0
Training Steps
Figure 7: T dynamics during training of ResNet152x1 on ImageNet-2Ik.
Table 19: Sensitivity to r starting point. Evaluated on ImageNet-2Ik with ResNet152x1 architecture.
TD
Tmin
Tmax
NNLL
PPrec I
1.25
0.05
5.0
5.59
0.485
2.525
0.05
5.0
5.71
0.486
3.75
0.05
5.0
5.73
0.485
4.9
0.05
5.0
5.72
0.478
In Table 19 and Table 20 we evaluate the effect of the choice of Tmin and Tmax and the initialization point along this range. We see that the final performance in the ResNer152x1 ImageNet-21k setting is not sensi- tive to these choices and that, if anything, further small improvements the results in the paper would be possible by tuning these hyperparameters. Since the removal of the need to tune the temperature hyperpa- rameter in the НET method is the primary motivation for learning т, we are satisfied with the insensitivity in performance to the temperature range and starting point and simply use this good default of Tnin = 0.05, Tiax 5.0 for all experiments.
L DIAGONAL COMPONENT PARAMETERIZATION
We recall that Collier et al. (2021) parameterize the covariance matrix
(x)) V(x) + diag(d(x)) with V(x) E RxX and d(x:) E R
(17)
Table 20: Sensitivity to 5 range. Evaluated on ImageNet-2lk with ResNet152x1 architecture.
70
Tmin
Tmax
JNL
P
2.525
0.05
5.0
5.71
0.486
2.525
0.001
5.0
5.50
0.488
2.525
0.01
5.0
5.73
0.485
2.525
0.05
4.0
5,95
0,485
2.525
0.05
3.0
5.61
0.485
27

---

Published as a conference paper ICLR 2023
Table 21: Comparing two parametrizations of the covariance matrix with equal parameter counts, namely V(x) V(x) + diag(d(x)) versus V(z) V(x)) + d(x)d(()) Evaluated ImageNet-2Ik with ResNet152x architecture.
Method
JNLL
1Pre@@1 1
HET-XI, w/ (()) HET-X- XL w/ diag (d(x))
5.71
0.486
5.67
0,486
HET-XL w/ d(x)d(()) leam
5.72
0.474
HET-XP w/ diag(d(()) learn т
5.72
0.473
However, we observe that, for the same parameter count, we can increase the expressiveness (xx) while also reducing (marginally) the cost of MC sampling <(x) by replacing the diagonal component with rank-1 component
E(x) VV d(x))d(=))
(18)
The marginally reduced cost of MC sampling is consequence of the shape of underlying standard Normal samples (0, I) required to sample from diag d())) versus d(x))((x)) Let us denote by S the number MC samples.
the diagonal case, we must sample tensor of standard Normal values shape [S, DJ. whereas similar to Monteiro et al. (2020), to induce the non-diagonal covariance terms only a tensor of shape [1, D] is required. This represents a saving of (S 1) x D standard Normal samples per example. In Table 21 empirically evaluate these choices. We that the choice of this rank-1 vs. diagonal compo- nent parameterization has little effect on the predictive performance. Hence, due to the marginal efficiency savings from the rank-1 parameterization, comparison, we also modify the Collier et al. (2021) method parameterization to use this rank-1 approach. This is the HET method reported throughout the paper.
choose parametertize HET-XL, accordingly. To ensure fair
M DETERMINISTIC VERSUS STOCHASTIC HETEROSCEDASTIC CONTRASTIVE LEARNING AT TEST TIME
In our contrastive learning setup, terministic representations to perform zero-shot classification, despite having added heteroscedastic noise at training time. This is similar the original HET paper, (Collier al., 2021) which fine-tunes deterministic representations from the HET model downstream on the VTAB (Zhai et al., 2019) benchmark. Therefore, test time the computational cost HET-XL. equal to that of DET. However we could also perform MC sampling at test time. Here ablate our choice by performing test time MC sampling with 22 MC samples (the maximum that fits in memory) and comparing the results to using the non-stochastic represen- tations. We use smaller scale setup with ViT-L/32 text encoder and billion training examples (all other hyperparameters as per the contrastive experiment in the main paper). We observe exactly the same zero- shot ImageNet accuracy of 84.35% regardless of whether perform MC sampling or the deterministic representation at test time.
only MC sample training time. At time use the de-
EXPERIMENTAL DETAILS AND HYPERPARAMETERS
All hyperparameters follows unless otherwise stated.
28

---

Published as a conference paper ICLR 2023
N.1 IMAGE CLASSIFICATION
All image classification experiments are trained on 64 TPU v3 cells with 128 cores. Heteroscedastic methods use 1,000 MC samples and the parameter efficient covariance parameterization with I = 50 is used. The temperature parameter for HET XL is parameterized as per Appendix K.
JFT-300M. All methods and models are trained for epochs with the Adam optimizer with 31 = 0.9, 82 = 0.999 and weight decay of 0.1 and otherwise default JAX hyperparameters. The learning rate dergoes a 10,000 linear warm-up phase starting at 10-6 and reaching x 10-1. A batch size of 4096 used.
ImageNet-22ks All methods and models are trained for 90 epochs with the Adam optimizer with 31 = 0.9, 32 0.999 and weight decay of 0.03 for ResNet152 and 0.1 for ViT-L/32 and otherwise default JAX hyperparameters. The learning rate undergoes a 10,000 linear warm-up phase starting at 10-5 and reaching 10 A batch size of 1024 is used for ResNet152 and 4096 for VIT-L/32.
JFT-4B. All methods and models are trained for epoch with the Adam optimizer with 81 0.9, B2 0.999 and weight decay of 0.1 and otherwise default JAX hyperparameters. The learning rate undergoes a 10,000 linear warm-up phase starting at 10 and reaching 6 x 10 4. A batch size of 4096 is used.
We note that we reuse, without modification, all the hyperparameters tuned in prior work (Tran et al., 2022). The only hyperparameters introduced by the HET-XL method are the boundaries of the learned temperature parameter which we set to be min=0.05 and max=5.0 for all experiments. We find that performance is insensitive to these boundaries provided the boundaries contain the final temperature learned by HET-XL, see Appendix K. Therefore we conclude that HET-XI. is relatively insensitive to hyperparameter choice, given that all existing hyperparameters can be reused and we remove the need for tuning the main sensitive hyperparameter from the HET method (T).
the HET method which have been
N.2 CONTRASTIVE LEARNING
All hyperparameters are taken from the original LiT paper Zhai et al. (2022) unless otherwise specified. We a JFT pre-trained ViT-G image encoder and ViT-g text encoder with vocabulary size of 32,000. We train with a batch size of 16,384 for 18B examples on 256 v3 TPU cells (512 cores). A cosine learning rate schedule is used with base learning For HET-XL, the learned temperature initialized to 3.0 and for DET initialized to the non parameter-efficient covariance parameterization in HET-XL, which is trained with MC sample but evaluated using the deterministic representations.
of 10 and 10,000 warm-up steps. Weight decay of 10 4 is used. 10.0. R = 15 for
FEW-SHOT PERFORMANCE
We follow Tran et al. (2022) in evaluating the representations learned by HET-XL with a linear evaluation protocol 9 few-shot learning datasets; birds, caltech, cars, cifar100, col hist, dtd, imagenet, pets, and uc merced; see Tran et al. (2022) for further details about the datasets. We evaluate I-shot, 5-shot, 10-shot and 25-shot performance. We evaluate 9 upstream settings
ViT-L/32 trained on JFT-4B for 1 epoch,
ViT-L/32 trained on JFT-300M for 7, 14 and 28 epochs,
ResNet152x1 trained on JFT-4B for 1, 2 and epochs,
522x trained JFT-300M for epochs and
29

---

Published as a conference paper ICLR 2023
Table 22: Average few-shot accuracy over 9 downstream datasets and different upstream settings std.
An upstream setting is defined by upstream dataset, base architecture and number of epochs. Please note that the standard errors are only equal after rounding and the reported standard errors are correct.
Method
1-shot ACC.
5-shot ACC.
10-shot ACC.
25-shot ACC.
DET
59.00 (+0.054)
81.27 (+0.025)
81.66 (+.0.033
85.12 03.120000025 0.029)
HET-XL
59.02 (=0.054)
81.92 (=0.025)
82.70 (#0.033)
85.87 (=0.029)
ResNet152x2 trained on JFT-300M for epochs.
Table 22 shows the performance for each number of shot(s), aggregated over the downstream datasets and the upstream settings. We can observe that in all settings HET-XL yields representations that transfer to downstream tasks better than DET.
P EQUALIZE DET PARAMS WITH HET-XL
Similar to (Collier et al., 2021) we implement an ablation in which the DET model's capacity is increased to the same number of parameters HET-XL and the performance is compared. In particular, for JFT-300M trained for 7 epochs with ViT-L/32 we increase the width of the last Dense layer in the network from 1,024 1,138. As result the DET network has 327.5M parameters, the same number the HET network. The DET model with increased capacity achieves Prec@) I: 0.470 and NLL: 7.82, a small improvement on the results in Table 3 which show the lower capacity DET model with Prec 0.468 and NLL: 7.83, but still far behind the HET-XL model which achieves Prec@) I: 0.498 and NLL: 7.65 at the same parameter count. This demonstrates that HET-XL's performance gains cannot simply be attributed increased model capacity.
30


Reviews:
[{"review_id": "uB9J-MneQsh", "reviewer": "Reviewer_EUZN", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigates the use of heteroscedastic classifiers for large scale classification problems. Heteroscedastic classifiers consider uncertainty in the decision function by learning a multivariate gaussian distribution over classifier logits. This paper proposes two simple tricks to scale such approach forclassification problems with large number of classes:\n 1) add gaussian noise at the pre-logit level of the networks, to remove the parameter dependencies on the number of classes\n 2) learn the softmax temperature instead of more involve strategies to stabilize the MC  sampling processing.\n\nExtensive experimentation on large scale datasets (ImageNet21K, JFT-300M/4B) shows the good scalability of their approaches and good performances over deterministic baseline and a strong hashing-based baseline. \n\nAuthors also apply their approach to a contrastive learning setting and investigate the 0-shot classification setup.\n", "summary_of_the_review": "The paper proposes a simple approach with good scalability property and extensive evaluation to demonstrate the value of their proposal. I therefore support acceptance. \n\n", "strength_and_weaknesses": "Strength:\n-\tSimplicity of the approach.\n-\tExtensive empirical study on large scale datasets.\n-\tStrong results that show both scalability of the approach and good performance.\n\nWeaknesses:\n- Lack of evaluation on downstream transfer tasks: Usually, representations learned from large-scale supervised datasets are then evaluated in a transfer setting through finetuning or linear probing. Would you expect HET-XL to show a gain in that setting, beyond the zero-shot case.\n- HET extra cost due to Monte-Carlo Sampling: HET requires to perform a MC sampling to compute the network output, which increase the computational cost of the forward prop compared to a deterministic model. How would a deterministic model with the same computational budget for the forward prop than HEC perform?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "u7jZMFUr6h", "reviewer": "Reviewer_r78j", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This manuscript proposes HET-XL, which is a simple yet effective solution for handling heteroscedastic classifiers for extreme classification. Specifically, the proposed method has advanced the scaling issues in the existing work. The extensive experiments and ablation studies showed the effectiveness of the proposed method on large-scale image classification benchmarks such as JFT-300M and ImageNet-21k.", "summary_of_the_review": "Overall, I would recommend the acceptance, as it has solid contributions and is worth sharing. Although the proposed solution is straightforward, it clearly advanced the previous works and has shown significant improvements on large-scale benchmarks. ", "strength_and_weaknesses": "Strengths\n- The manuscript is well-written and easy to understand.\n- The proposed idea is simple yet effective.\n- The proposed method has advanced training efficiency and scalability of learning heteroscedastic classifiers for extreme classification tasks using a simple idea.\n- Extensive experimental results demonstrate the effectiveness of the proposed method.\n\nWeaknesses\n- In my understanding, the proposed method takes benefits from breaking the full rank assumption of the covariance matrix (from $K$ to $D$). The authors provide empirical evidence for this, but it is still unclear why it works. Could you provide any theoretical analysis for this? or How big does $D$ work well? Specifically, does the proposed method effectively work on ViT models having small dimensionality (e.g., ViT-small and ViT-tiny)?\n- I think applying the proposed method to contrastive learning is an interesting idea. However, the reported improvements are marginal. However, under the perspective of instance discrimination view, there exist very strong class-correlations; for example, CIFAR-100 has actual 100 categories but 50,000 classes in this case. I am curious about whether we can expect heteroscedastic classifiers to work well under such a strong class correlation. Moreover, in contrastive learning, I think adding noise before and after $l_2$ normalization would have different behaviors.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eqVQcVQL-QF", "reviewer": "Reviewer_HrTZ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "As indicated by the title, this paper considers the problem of scaling heteroscedastic classifiers to problems where the number of classes is large.  The main contribution is a simple trick which eliminates the dependence on the number of classes on the total number of parameters.  The authors also show that the temperature parameter commonly used for these networks does not need to be treated as a hyperparameter; rather, it can be optimized as a trainable parameter.  The authors discuss how these insights can be applied to contrastive learning, and then they provide thorough experiments showing that their method scales more efficiently than existing heteroschedastic classifies and that it outperforms a deterministic baseline.", "summary_of_the_review": "Overall, I thought this was a solid paper.  The empirical results are impressive for the heteroscedastic setting.  The method is simple and clearly presented.  And the experiments are thorough.  On the negative side, aside from a few minor points of confusion, I would argue that the contrastive results are not particularly impactful and that the time-complexity of HET-XL is concerning.  I think that this paper has some empirical novelty, but it does not introduce tools that may be more broadly applicable to the learning community.  All of this being said, I think that this paper does exactly what it sets out to do, which is to improve the SOTA in heteroscedastic classification.  Therefore, I recommend that this paper be accepted.", "strength_and_weaknesses": "### Strengths\n\n**Writing quality.**  This is the best-written paper in my batch of ICLR papers this year.  The arguments are clearly articulated and it's easy to understand the main ideas.\n\n**Simplicity.**  I found the simplicity of the main ideas to be appealing.  The trick used to change the layer at which the noise is added is simple, and it's encouraging (as shown in an appendix) that trick comes at no loss of performance.  The main ideas seem like they would easy to implement for those who wanted to reproduce this work, and authors also seem committed to making their code available, which is a positive and encouraging sign.\n\n**Experiments.**  The experiments are thorough and they show that the method seems to offer a performance improvement over both deterministic classifiers (henceforth DET) and heteroscedastic classifiers (henceforth DET).  Furthermore, relative to HET, the proposed architecture HET-XL has a smaller number of parameters; indeed, the authors show their architecture's parameter count does not scale with the number of classes $K$.  These two results -- both highlighted in Table 3 -- are exactly what one would hope to see.\n\nOverall, the rest of the experiments also seem to be quite thorough.  In the main text, we get a couple of ablations concerning the latency and memory usage of the relevant algorithms.  There are also quite a few extra results in the appendix, including ablation studies on learning $\\tau$, model scaling, data scaling, and the number of floating point operations.  In eacfh case, it seems that HET-XL compares favorably to the baselines.\n\n## Weaknesses\n\n**Minor points of confusion.**  Here are a couple of things that confused me while I was reading the paper:\n\n* I don't understand why this model requires *marginalizing* over the noise distribution.  Granted, I'm not an expert on heteroscedastic classifiers, so perhaps I have missed something obvious here.  However, from my understanding, the noise is modelled by a Gaussian, which we assume has zero mean.  Furthermore, the covariance matrix of this Gaussian is learned from data.  One one has access to this covariance matrix, why does one need to run MCMC to sample from this noise distribution?  Can one not directly sample, given that the distribution is normal?  Perhaps it has something to do with needed to differentiate through the covariance estimation?  If the authors could elaborate on this point, I think it would make the paper clearer.\n\n* What is the so-called \"hashing trick?\"  It's relatively difficult to understand the innovation in Section 6.1 without knowing this, and I think that it's possible that many readers of this paper may not have heard of it before.  In my opinion, it would be worth adding a few lines to the paper to describe the main idea behind this trick, and why it's useful here.\n\n**Contrastive learning.**  Having read the paper, I'm not sure what the contribution of the contrastive learning section is.  Indeed, it seems clear that contrastive learning can be formulated as a massive classification problem.  However, the question is: Does this perspective result in better algorithms or architectures for contrastive learning.  And based on the results in Table 4, it's not clear that it does.  The table indicates that on ImageNet, the heteroscadastic classification view of contrastive learning yields a marginal (albeit, statistically significant) improvement.  However, on the other two datasets, HET-XL performs worse than the baseline deterministic classifier.  Therefore, it's somewhat unclear whether it's worth using a classifier with potentially many more parameters relative to DET in the setting of contrastive learning.  For these reasons, I do not see the contrastive results as an impactful contribution.  If I have missed something, I would welcome the thoughts of the authors on this point.\n\n**Efficiency.**  After reading the paper, it seemed unintuitive that when taking many MC steps, HET-XL underperformed relative to HET in terms of latency.  Could the authors shed more light on this?  Why is sampling less efficient for HET-XL?  And if this is the case, in terms of time-complexity, would one then tend to prefer HET over HET-XL.  (Of course, in space-complexity, it does seem that HET-XL is to be preferred in these large-scale settings.)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "RaVJuyMRPRD", "reviewer": "Reviewer_GaFB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a technique for heteroschedastic classification, HET-XL, which does a simple trick to maintain scalability to settings with large amounts of classes.\n\nThe key idea is deceptively simple: the authors elect to model the covariance over noise not over class-space in the logits, but over the empedding space pre-logits, leading to a covariance matrix of much smaller dimension int he large class scenario which maintains tractability at ease.\n\nExperiments show this approach also confers some regularization benefits and remains indeed tractable.", "summary_of_the_review": "The authors propose a simple trick to make heteroschedastic classification scalable to large class-scenarios.\nTheir key idea is simple, easily applicable, and although it appears to lose out on some modeling oompf theoretically performs quite well in practice.\n\nI think this could be useful for practitioners going after raw performance, though I would caution not to overinterpret the learned covariance matrix semantics as equivalent to one learned in class-space.", "strength_and_weaknesses": "Strengths:\n- this is a simple technique to model heteroschedastic label noise and makes it tractable to large dimensions. \n- the technique is enormously simple to utilize and plug in to existing classifiers.\n\nWeakness:\n- in theory, the expressivity of this parametrization should be lower than in the \"full rank\" scenario. However, the authors show graceful performance in practice. I feel it is noteworthy to keep this in mind in practice.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "G0y8hhK-UO", "reviewer": "Reviewer_mUAK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a method for performing Heteroscedastic classification that can scale independently of the number of classes.  The method simply defines noise not over the logits of a network as is commonly done, but over the layer previous to the logins.  In this way the learned instance-specific covariance of the noise distribution is a DxD matrix (where D is the dimensionality of the pre-logit representation) instead of KxK (where K is the number of classes).  The authors also provide some empirical justification for treating the commonly validated temperature Hyperparameter as a free parameter in the training procedure.  In addition to those two contributions, authors define a contrastive learning model that uses their proposed noise model.  These three contributions are aimed at scaling heteroscedastic classification to cases where there are a considerable (> 10,000) classes.  Empirically, their method out-performs both heteroscedastic and deterministic baselines in terms of negative log likelihood and precision while producing models with fewer parameters than other heteroscedastic baselines.", "summary_of_the_review": "While I feel the contributions of this paper likely improve the practice of learning heteroscedastic classifiers, I feel they are incremental and not significant enough to warrant acceptance.\n\n(See response to rebuttal for change of final score)", "strength_and_weaknesses": "Strengths\n1. Overall the method is presented in a clear manner.  It is both easy to see why the method would result in models with fewer parameters and likely easy to implement.\n2. The empirical results seem to justify the changes to standard heteroscedastic learning as proposed in the paper.\n\nWeaknesses \n1. The biggest weakness is that the the work lacks novelty.  Each of the three contributions is simply applying well known techniques with slight variations to reduce practical run time of training and inference.  The paper provides no principled or theoretical contribution, which makes the paper seem like more incremental than substantial.\n2. The empirical results are not very compelling.  First, the number of epochs that each model is trained for does not seem representative of complete training procedures.  The JFT-4B (which are from a proprietary data set, and thus is impossible to replicate) results are from training from a single epoch, for instance.  A more compelling result would be to plot the NLL and/or precision of the model as a function of wall run time, as this would be illustrate the claimed main benefit of HET-XL: With fewer parameters, HET-XL is able to scale better than HET methods.  Further, the ablation results need to be in the main paper and more clearly presented in the context of the results using the individual components.  The memory and train timing figures are a good idea in principle, but it is unclear how the number of MC samples practically effects training.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "FDBTLZZAv8h", "reviewer": "Reviewer_FYog", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies heteroscedastic models and extends the previous studies that require an explicit modelling of the noise structure into a simple modelling by reusing the logits from the mean (non-noise part) term. And then the idea is further extended to contrastive learning setting. ", "summary_of_the_review": "reasonably good work, but some clarification on why the idea can bypass the identification issue might need to be deeply discussed. ", "strength_and_weaknesses": "- strength\n    - the idea is straightfoward and a natural extension of the existing methods\n    - the discussion and summary to previous methods are nice\n- weakness\n    - it seems unclear to me that where the strength of the methods come from \n    - the empirical strength of the methods does not seem clear to me", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster