Quantile Activation: Correcting a failure mode of ML models

## Abstract

Abstract A classifier is, in its essence, a function which takes an input and returns the class of the input and implicitly assumes an underlying distribution. We argue in this article that one has to move away from this basic tenet to obtain generalization across distributions. Specifically, the class of the sample should depend on the points from its “context distribution” for better generalization across distributions. How does one achieve this? – The key idea is to “adapt” the outputs of each neuron of the network to its context distribution. We propose quantile activation, QAct , which, in simple terms, outputs the relative quantile of the sample in its context distribution, instead of the actual values in traditional networks. The scope of this article is to validate the proposed activation across several experimental settings, and compare it with conventional techniques. For this, we use the datasets developed to test robustness against distortions – CIFAR10C, CIFAR100C, MNISTC, TinyImagenetC, and show that we achieve a significantly higher generalization across distortions than the conventional classifiers, across different architectures. Although this paper is only a proof of concept, we surprisingly find that this approach outperforms DINOv2(small) at large distortions, even though DINOv2 is trained with a far bigger network on a considerably larger dataset.

## 1 Introduction

Deep learning approaches have significantly influenced image classification tasks on the machine learning pipelines over the past decade. They can easily beat human performance on such tasks by non-trivial margins by using innovative ideas such as Batch Normalization [] and other normalization techniques [, ], novel rectifiers such as ReLU/PReLU [, , ] and by using large datasets and large models.

However, these classification systems do not generalize across distributions [, ], which leads to instability when used in practice. [] shows that deep networks with ReLU activation degrades in performance under distortions. [] observes that there exists a feature collapse which inhibits the networks to be reliable.

### Fundamental Problem of Classification:

We trace the source of the problem to the fact that – *Existing classification pipelines necessitates single point prediction*, i.e., they should allow classification of a single sample given in isolation. We argue that, for good generalization, one should move away from this basic tenet and instead allow the class to be dependent on the *“context distribution”* of the sample. That is, when performing classification, one needs to know both the sample and the context of the sample for classification.

While this is novel in the context of classification systems, it is widely prevalent in the specific domain of Natural Language Processing (NLP) – The meaning of a word is dependent on the context of the word. However, to our knowledge, this has not been considered for general classification systems. Even when using dominant NLP architectures such as Transformers for vision [], the technique has been to split the image into patches and then obtain the embedding for an individual image.

Figure: Figure 1: Comparing TSNE plots of QAct and ReLU activation on CIFAR10C with Gaussian distortions. Observe that QAct maintains the class structure extremely well across distortions, while the usual ReLU activations loses the class structure as severity increases.
Refer to caption: extracted/5602265/img/article/Fig1/quant_test.png

### Obtaining a classification framework for incorporating context distribution:

We suspect that the main reason why the context distribution is not incorporated into the classification system is – The naive approach of considering a lot of samples in the pipeline to classify a single sample is computationally expensive. We solve this problem by considering the context distribution of each neuron specifically. We introduce *Quantile Activation* (QAct), which outputs a probability depending upon the context of *all* outputs. This, however, gives rise to new challenges in training, which we address in section 2.

Figure 1 illustrates the differences of the proposed framework with the existing framework. As severity increases (w.r.t Gaussian Noise), we observe that ReLU activation loses the class structure. This behaviour can be attributed to the fact that, as the input distribution changes, the activations either increase/decrease, and due to the multiplicative effect of numerous layers, this leads to very different features. On the other hand, the proposed QAct framework does not suffer from this, since if all the pre-activations(^1^11We use the following convention – “Pre-activations” denote the inputs to the activation functions and “Activations” denote the outputs of the activation function.) change in a systematic way, the quantiles *adjust* automatically to ensure that the inputs for the next layer does not change by much. This is reflected in the fact that class structure is preserved with QAct.

### Remark:

Quantile activation is different from existing quantile neural network based approaches, such as regression [], binary quantile classification [], Anomaly Detection [, ]. Our approach is achieving best in-class performance by incorporating context distribution in the classification paradigm. Our approach is also markedly different from Machine unlearning which is based on selective forgetting of certain data points or retraining from scratch [].

### Contributions:

A decent amount of literature on neuronal activation is available. However, to the best of our knowledge, none matches the central idea proposed in this work.

In [], the authors propose an approach to calibrate a pre-trained classifier $f_{\theta}({\bm{x}})$ by extending it to learn a *quantile function*, $Q({\bm{x}},\theta,\tau)$ ($\tau$ denotes the quantile), and then estimate the probabilities using $\int_{\tau}I[Q({\bm{x}},\theta,\tau)\geq 0.5]d\tau$(^2^22$I[]$ denotes the indicator function). They show that this results in probabilities which are robust to distortions.

- 1.
In this article, we extend this approach to the level of a neuron, by suitably deriving the forward and backward propagation equations required for learning (section 2).
- 2.
We then show that a suitable incorporation of our extension produces context dependent outputs at the level of each neuron of the neural network.
- 3.
Our approach contributes to achieving better generalization across distributions and is more robust to distortions, across architectures. We evaluate our method using different architectures and datasets, and compare with the current state-of-the-art – DINOv2. We show that QAct proposed here is more robust to distortions than DINOv2, even if we have considerably less number of parameters (22M for DINOv2 vs 11M for Resnet18). Additionally, DINOv2 is trained on 20 odd datasets, before being applied on CIFAR10C; in contrast, our framework is trained on CIFAR10, and produces more robust outcome (see figures 3,5).
- 4.
The proposed QAct is consistent with all the existing techniques used in DINOv2, and hence can be easily incorporated into any ML framework.
- 5.
We also adapt QAct to design a classifier which returns better calibrated probabilities.

### Related Works on Domain Generalization (DG):

The problem of domain generalizations tries to answer the question – Can we use a classifier trained on one domain across several other related domains? The earliest known approach for this is *Transfer Learning* [, ], where a classifier from a single domain is applied to a different domain with/without fine-tuning. Several approaches have been proposed to achieve DG, such as extracting domain-invariant features over single/multiple source domains [, , , , ], Meta Learning [, ], Invariant Risk Minimization []. Self supervised learning is another proposed approach which tries to extract features on large scale datasets in an unsupervised manner, the most recent among them being DINOv2 []. Very large foundation models, such as GPT-4V, are also known to perform better with respect to distribution shifts []. Nevertheless, to the best of our knowledge, none of these models incorporates context distributions for classification.

## 2 Quantile Activation

### Rethinking Outputs from a Neuron:

To recall – if ${\mathbf{x}}$ denotes the input, a typical neuron does the following – (i) Applies a linear transformation with parameters $w,b$, giving $w^{t}{\mathbf{x}}+b$ as the output, and (ii) applies a rectifier $g$, returning $g(w^{t}{\mathbf{x}}+b)$. Typically, $g$ is taken to be the ReLU activation - $g_{relu}(x)=\max(0,x)$. Intuitively, we expect that each neuron captures an “abstract” feature, usually not understood by a human observer.

An alternate way to model a neuron is to consider it as predicting a latent variable ${\mathbf{y}}$, where ${\mathbf{y}}=1$ if the feature is present and ${\mathbf{y}}=0$ if the feature is absent. Mathematically, we have the following model:

$$ ${\mathbf{z}}=w^{t}{\mathbf{x}}+b+\epsilon\quad and\quad{\mathbf{y}}=I[{\mathbf {z}}\geq 0]$ (1) $$

This is very similar to the standard latent variable model for logistic regression, with the main exception being, the *outputs ${\mathbf{y}}$ are not known* for each neuron beforehand. If ${\mathbf{y}}$ is known, it is rather easy to obtain the probabilities – $P({\mathbf{z}}\geq 0)$. Can we still predict the probabilities, even when ${\mathbf{y}}$ itself is a latent variable?

The authors in [] propose the following algorithm to estimate the probabilities:

- 1.
Let $\{{\bm{x}}_{i}\}$ denote the set of input samples from the input distribution ${\mathbf{x}}$ and $\{z_{i}\}$ denote their corresponding latent outputs, which would be from the distribution ${\mathbf{z}}$
- 2.
Assign ${\mathbf{y}}=1$ whenever ${\mathbf{z}}>(1-\tau)^{th}$ quantile of ${\mathbf{z}}$, and 0 0 otherwise. For a specific sample, we have $y_{i}=1$ if $z_{i}>(1-\tau)^{th}$ quantile of $\{z_{i}\}$
- 3.
Fit the model $Q(x,\tau;\theta)$ to the dataset $\{(({\bm{x}}_{i},\tau),y_{i})\}$, and estimate the probability as,
$P(y_{i}=1)=\int_{\tau=0}^{1}I[Q(x,\tau;\theta)\geq 0.5]d\tau$
(2)

### The key idea:

Observe that in step 2., the labelling is done without resorting to actual ground-truth labels. This allows us to obtain the probabilities on the fly for any set of parameters, only by considering the quantiles of ${\mathbf{z}}$.

### Defining the Quantile Activation QAct

Let ${\mathbf{z}}$ denote the pre-activation of the neuron, and let $\{{\bm{z}}_{i}\}$ denote the samples from this distribution. Let $F_{{\mathbf{z}}}$ denote the cumulative distribution function (CDF), and let $f_{{\mathbf{z}}}$ denote the density of the distribution. Accordingly, we have that $F_{{\mathbf{z}}}^{-1}(\tau)$ denotes the $\tau^{th}$ quantile of ${\mathbf{z}}$. Using step (2) of the algorithm above, we define,

$$ $\textsc{QAct}{}({\bm{z}})=\int_{\tau=0}^{1}I[{\bm{z}}>F_{{\mathbf{z}}}^{-1}(1- \tau)]d\tau\overunderset{\text{Substitute}}{\tau\rightarrow(1-\tau)}{=}\int_{ \tau=0}^{1}I[{\bm{z}}>F_{{\mathbf{z}}}^{-1}(\tau)]d\tau$ (3) $$

### Computing the gradient of QAct :

However, to use QAct in a neural network, we need to compute the gradient which is required for back-propagation. Let $\tau_{{\bm{z}}}$ denote the quantile at which $F_{{\mathbf{z}}}^{-1}(\tau_{{\bm{z}}})={\bm{z}}$. Then we have that $\textsc{QAct}{}({\bm{z}})=\tau_{z}$ since $F_{{\mathbf{z}}}^{-1}(\tau)$ is an increasing function. So, we have that $\textsc{QAct}{}(F_{{\mathbf{z}}}^{-1}(\tau))=\tau$. In other words, we have that $\textsc{QAct}{}({\bm{z}})$ is $F_{{\mathbf{z}}}({\bm{z}})$, which is nothing but the CDF of ${\mathbf{z}}$. Hence, we have,

$$ $\frac{\partial\textsc{QAct}{}({\bm{z}})}{\partial z}=f_{{\mathbf{z}}}({\bm{z}})$ (4) $$

where $f_{{\mathbf{z}}}({\bm{z}})$ denotes the density of the distribution.

Figure: Algorithm 1 Forward Propagation for a single neuron

Figure: Algorithm 2 Backward Propagation for a single neuron

Figure: (a)
Refer to caption: extracted/5602265/img/article/Fig9/fig9a.png

### Grounding the Neurons:

With the above formulation, observe that since QAct is identical to CDF, it follows that, $\textsc{QAct}{}({\mathbf{z}})$ is always a uniform distribution between 0 0 and $1$, irrespective of the distribution ${\mathbf{z}}$. When training numerous neurons in a layer, this could cause all the neurons to learn the same behaviour. Specifically, if, half the time, a particular abstract feature is more prevalent than others, QAct (as presented above) would not be able to learn this feature. To correct this, we *enforce that positive values and negative values have equal weight*. Given the input distribution ${\mathbf{z}}$, We perform the following transformation before applying QAct. Let

$$ ${\mathbf{z}}^{+}=\begin{cases}{\mathbf{z}}&\text{ if }{\mathbf{z}}\geq 0\\ 0&\text{ otherwise}\end{cases}\hskip 80.00012pt{\mathbf{z}}^{-}=\begin{cases}{ \mathbf{z}}&\text{ if }{\mathbf{z}}<0\\ 0&\text{ otherwise}\end{cases}$ (7) $$

denote the truncated distributions. Then,

$$ ${\mathbf{z}}^{\ddagger}=\begin{cases}{\mathbf{z}}^{+}&\text{ with probability }0.5\\ {\mathbf{z}}^{-}&\text{ with probability }0.5\\ \end{cases}$ (8) $$

From definition of ${\mathbf{z}}^{\ddagger}$, we get that the median of ${\mathbf{z}}^{\ddagger}$ is 0 0. This grounds the input distribution to have the same positive and negative weight.

### Dealing with corner cases:

It is possible that during training, some neurons either only get positive values or only get negative values. However, for smooth outputs, one should still only give the weight of $0.5$ for positive values. To handle this, we include two values $c$ (large positive) and $-c$ (large negative) for each neuron. Since, the quantiles are conventionally computed using linear interpolation, this allows the outputs to vary smoothly. We take $c=100$ in this article.

### Estimating the Density for Back-Propagation:

Note that the gradient for the back propagation is given by the density of ${\mathbf{z}}^{\ddagger}$ (weighted distribution). We use the *Kernel Density Estimation* (KDE), to estimate the density. We, (i) First sample $N$ points with weights $w_{+},w_{-}$, and (ii) then estimate the density at all the input points $[{\bm{z}}_{i}]$. This is point-wise multiplied with the backward gradient to get the gradient for the input. In this article we use $N=1000$, which we observe gets reasonable estimates.

### Computational Complexity:

Computational Complexity (for a single neuron) is majorly decided by 2 functions – (i) Computing the quantiles has the complexity for a vector $[{\bm{z}}_{i}]$ of size $n$ can be performed in ${\mathcal{O}}(n\log(n))$. Since this is log-linear in $n$, it does not increase the complexity drastically compared to other operations in a deep neural network. (ii) Computational complexity of the KDE estimates is ${\mathcal{O}}(Sn_{\tau})$ where $S$ is the size of sample (weighted sample from $[{\bm{z}}_{i}]$) and $n_{\tau}$ is the number of quantiles, giving a total of ${\mathcal{O}}(n+Sn_{\tau})$. In practice, we consider $S=1000$ and $n_{\tau}=100$ which works well, and hence does not increase with the batch size. This too scales linearly with batch size $n$, and hence does not drastically increase the complexity.

### Remark:

Algorithms 1, and 2 provide the pseudocode for the quantile activation. For stable training, in practice, we prepend and append the quantile activation with BatchNorm layers.

### Why QAct is robust to distortions?

To understand the idea behind quantile activation, consider a simple toy example in figure 2. For ease of visualization, assume that the input features (blue) are in 2 dimensions, and also assume that the line of the linear projection is given by the red line in figure 2(a). Now, assume that the blue input features are rotated, leading to a different distribution (indicated here by orange). Since activations are essentially (unnormalized) signed distances from the line, we plot the histograms corresponding to the two distributions in figure 2(b). As expected, these distributions are different. However, after performing the quantile activation in equation 3, we have that both are uniform distribution. This is illustrated in figures 2(c) and 2(d). This behaviour has a normalizing effect across different distributions, and hence has better distribution generalization than other activations.

## 3 Training with QAct

In the previous section, we described the procedure to adapt a single neuron to its context distribution. In this section we discuss how this extends to the Dense/Convolution layers, the loss functions to train the network and the inference aspect.

### Extending to standard layers:

The extension of equation 3 to dense outputs is straightforward. A typical output of the dense layer would be of the shape $(B,N_{c})$ - $B$ denotes the batch size, $N_{c}$ denotes the width of the network. The principle is - *The context distribution of a neuron is all the values which are obtained using the same parameters*. In this case, each of the values across the ‘$B$’ dimension are considered to be samples from the context distribution.

For a convolution layer, the typical outputs are of the form - $(B,N_{c},H,W)$ - $B$ denotes the size of the batch, $N_{c}$ denotes the number of channels, $H,W$ denotes the sizes of the images. In this case we should consider all values across the 1st,3rd and 4th dimension to be from the context distribution, since all these values are obtained using the same parameters. So, the number of samples would be $B\times H\times W$.

### Loss Functions:

One can use any differentiable loss function to train with quantile activation. We specifically experiment with the standard Cross-Entropy Loss, Triplet Loss, and the recently proposed Watershed Loss [] (see section 4). However, if one requires that the boundaries between classes adapt to the distribution, then learning similarities instead of boundaries can be beneficial. Both Triplet Loss and Watershed Loss fall into this category. We see that learning similarities does have slight benefits when considering the embedding quality.

### Inference with QAct :

As stated before, we want to assign a label for classification based on the context of the sample. There exist two approaches for this – (1) One way is to keep track of the quantiles and the estimated densities for all neurons and use it for inference. This allows inference for a single sample in the traditional sense. However, this also implies that one would not be able to assign classes based on the context at evaluation. (2) Another way is to make sure that, even for inference on a single sample, we include several samples from the context distribution, but only use the output for a specific sample. This allows one to assign classes based on the context. In this article, we follow the latter approach.

### Quantile Classifier:

Observe that the proposed QAct (without normalization) returns the values in $[0,1]$ which can be interpreted as probabilities. Hence, one can also use this for the classification layer. Nonetheless, two changes are required – (i) Traditional softmax used in conjunction with negative-log-likelihoood loss already considers “relative” activations of the classification in normalization. However, QAct does not. Hence, one should use Binary-Cross-Entropy loss with QAct, which amounts to one-vs-rest classification. (ii) Also, unlike a neuron in the middle layers, the bias of the neuron in the classification layer depends on the class imbalance. For instance, with $10$ classes, one would have only $1/10$ of the samples labelled $1$ and $9/10$ of the samples labelled 0 0. To address this, we require that the median of the outputs be at $0.9$, and hence weight the positive class with $0.9$ and the negative class with $0.1$ respectively. In this article, whenever QAct is used, we use this approach for inference.

We observe that (figures 13 and 14) using quantile classifier on the learned features in general improves the consistency of the calibration error and also leads to the reducing the calibration error. In this article, for all networks trained with quantile activation, we use quantile classifier to compute the accracies/calibration errors.

## 4 Evaluation

To summarize, we make the following changes to the existing classification pipeline – (i) Replace the usual ReLU activation with QAct and (ii) Use triplet or watershed loss instead of standard cross entropy loss. We expect this framework to learn context dependent features, and hence be robust to distortions. (iii) Also, use quantile classifier to train the classifier on the embedding for better calibrated probabilities.

### Evaluation Protocol:

To evaluate our approach, we consider the two datasets developed for this purpose – CIFAR10C, CIFAR100C, TinyImagenetC [], MNISTC[]. These datasets have a set of 15 distortions at 5 severity levels. To ensure diversity we evaluate our method on 4 architectures – (overparametrized) LeNet, ResNet18[] (11M parameters), VGG[](15M parameters) and DenseNet [](1M parameters). The code to reproduce the results can be found at [https://anonymous.4open.science/r/QuantAct-2B41](https://anonymous.4open.science/r/QuantAct-2B41).

### Baselines for Comparison:

To our knowledge, there exists no other framework which proposed classification based on context distribution. So, for comparison, we consider standard ReLU activation [], pReLU [], and SELU [] for all the architectures stated above. Also, we compare our results with DINOv2 (small) [] (22M parameters) which is current state-of-the-art for domain generalization. Note that for DINOv2, architecture and datasets used for training are substantially different (and substantially larger) from what we consider in this article. Nevertheless, we include the results for understanding where our proposed approach lies on the spectrum. We consider the small version of DINOv2 to match the number of parameters with the compared models.

### Metrics:

We consider four metrics – Accuracy (ACC), calibration error (ECE) [] (both marginal and Top-Label) and mean average precision at $K$ (MAP@K) to evaluate the embedding. For the case of ReLU/pReLU/SELU activation with Cross-Entropy, we use the logistic regression trained on the train set embeddings, and for QAct we use the calibrated linear classifier, as proposed above. We do not perform any additional calibration and use the probabilities. We discuss a selected set of results in the main article. Please see appendix C for more comprehensive results.

Calibration error measures the reliability of predicted probabilities. In simple words, if one predicts 100 samples with (say) probability 0.7, then we expect 70 of the samples to belong to class $1$ and the rest to class 0 0. This is measured using either the marginal or top-label calibration error. We refer the reader to [] for details, which also provides an implementation to estimate the calibration error.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig2/Accuracy.png

### Remark:

For all the baselines we use the standard Cross-Entropy loss for training. For inference on corrupted datasets, we retrain the last layer with logistic regression on the train embedding and evaluate it on test/corrupted embedding. For QAct, we as a convention use watershed loss unless otherwise stated, for training. For inference, we train the Quantile Classifier on the train embedding and evaluate it on test/corrupted embedding.

### The proposed QAct approach is robust to distortions:

In fig. 3 we compare the proposed QAct approach with predominant existing pipeline – ReLU+Cross-Entropy and DINOv2(small) on CIFAR10C. In figure 3(a) we see that as the severity of the distortion increases, the accuracy of ReLU and DINOv2 drops significantly. On the other hand, while at small distortions the results are comparable, as severity increases QAct performs substantially better than conventional approaches. At severity $5$, QAct outperforms DINOv2. On the other hand, we observe that in figure 3(b), the calibration error stays consistent across distortions.

Figure: (a)
Refer to caption: extracted/5602265/img/article/Fig3/MAP.png

### How much does QAct depend on the loss function?

Figure 4(a) compares the watershed classifier with other popular losses – Triplet and Cross-Entropy. We see that all the loss functions perform comparably when used in conjunction with QAct. We observe that watershed has a slight improvement when considering MAP and hence, we consider that as the default setting. However, we point out that QAct is compatible with several loss functions as well.

### QAct vs ReLU/pReLU/SELU activations:

To verify that most existing activations do not share the robustness property of QAct, we compare QAct with other activations in figures 4(b) and 4(c). We observe that QAct is greatly more robust with respect to distortions in both accuracy and calibration error than other activation functions.

Figure: (a)
Refer to caption: extracted/5602265/img/article/Fig8/Accuracy.png

### Results on Larger Datasets:

To verify that our observations hold for larger datasets, we use CIFAR100C/TinyImagenetC to compare the proposed QAct+watershed with existing approaches. We observe on figure 5 that QAct performs comparably well as DINOv2, although DINOv2(s) has 22M parameters and is trained on significantly larger datasets. Moreover, we also observe that QAct has approximately constant calibration error across distortions, as opposed to a significantly increasing calibration error for ReLU or DINOv2.

## 5 Conclusion And Future Work

To summarize, traditional classification systems do not consider the “context distributions” when assigning labels. In this article, we propose a framework to achieve this by – (i) Making the activation adaptive by using quantiles and (ii) Learning a kernel instead of the boundary for the last layer. We show that our method is more robust to distortions by considering MNISTC, CIFAR10C, CIFAR100C, TinyImagenetC datasets across varying architectures.

The scope of this article is to provide a proof of concept and a framework for performing inference in a context-dependent manner. We outline several potential directions for future research:

- I.
The key idea in our proposed approach is that the quantiles capture the distribution of each neuron from the batch of samples, providing outputs accordingly. This poses a challenge for inference, and we have discussed two potential solutions: (i) remember the quantiles and density estimates for single sample evaluation, or (ii) ensure that a batch of samples from the same distribution is processed together. We adopt the latter method in this article. An alternative approach would be to *learn the distribution of each neuron* using auxiliary loss functions, adjusting these distributions to fit the domain at test time. This gives us more control over the network at test time compared to current workflows.
- II.
Since the aim of the article was to establish a proof-of-concept, we did not focus on scaling, and use only a single GPU for all the experiments. To extend it to multi-GPU training, one needs to synchronize the quantiles across GPU, in a similar manner as that for Batch-Normalization. We expect this to improve the statistics, and to allow considerably larger batches of training.
- III.
On the theoretical side, there is an interesting analogy between our quantile activation and how a biological neuron behaves. It is known that when the inputs to a biological neuron change, the neuron adapts to these changes []. Quantile activation does something very similar, which leads to an open question – can we establish a formal link between the adaptability of a biological neuron and the accuracy of classification systems?
- IV.
Another theoretical direction to explore involves considering distributions not just at the neuron level, but at the layer level, introducing a high-dimensional aspect to the problem. The main challenge here is defining and utilizing *high dimensional quantiles*, which remains an open question [].

### Broad Impact:

In this article, we propose an approach to maintain calibration and generalization across small distortions. While, we do not foresee any direct societal consequences of our work, we expect the potential future consequences of the technique to reduce the bias in the following ways – (i) Since we do not assume normal distribution, our approach is likely to handle long tails better than existing methods. This would help in reducing the dataset bias where marginal groups are less represented. (ii) Note that the output of each QAct layer is a uniform distribution. This can allow us to understand the working of each layer in isolation and possibly reduce the black-box nature of the current classification systems. (iii) Moreover, by directly modifying the context distribution of each neuron, one can easily make the networks more reliable without resorting to expensive re-training the entire network.

## Appendix A Experiment details for figure 1

We consider the features obtained from ResNet18 with both QAct and RelU activations for the datasets of CIFAR10C with gaussian_noise at all the severity levels. Hence, we have 6 datasets in total. To use TSNE for visualization, we consider $1000$ samples from each dataset and obtain the combined TSNE visualizations. Each figure shows a scatter plot of the 2d visualization for the corresponding dataset.

## Appendix B Compute Resources and Other Experimental Details

All experiments were performed on a single NVidia GPU with 32GB memory with Intel Xeon CPU (10 cores). For training, we perform an 80:20 split of the train dataset with seed $42$ for reproducibility. All networks are initialized using default pytorch initialization technique.

We use Adam optimizer with initial learning rate $1e-3$. We use ReduceLRonPlateau learning rate scheduler with parameters – factor=0.1, patience=50, cooldown=10, threshold=0.01, threshold_mode=abs, min_lr=1e-6. We monitor the validation accuracy for learning rate scheduling. We also use early_stopping when the validation accuracy does not increase by $0.001$.

## Appendix C Extended Results Section

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig2/Accuracy.png

### Comparing QAct + watershed and ReLU+Cross-Entropy:

Figure 6 shows the corresponding results. The first experiment compares QAct + watershed with ReLU + Cross-Entropy on two standard networks – ResNet18 and DenseNet. With respect to accuracy, we observe that while at severity 0 0, ReLU + Cross-Entropy slightly outperforms QAct + watershed, as severity increases QAct + watershed is far more stable. We even outperform DinoV2(small) (22M parameters) at severity $5$. Moreover, with respect to calibration error, we see a consistent trend across distortions. As [] argues, this helps in building more robust systems compared to one where calibration error increases across distortions.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig3/Accuracy.png

### Does loss function make a lot of difference?

Figure 7 compares three different loss functions Watershed, Triplet and Cross-Entropy when used in conjunction with QAct. We observe similar trends across all loss functions. However, Watershed performs better with respect to Mean Average Precision (MAP) and hence we use this as a default strategy.

Why Mean-Average-Precision? – We argue that the key indicator of distortion invariance should be the quality of embedding. While, accuracy (as measured by a linear classifier) is a good metric, a better one would be to measure the Mean-Average-Precision. With respect to calibration error, due to the scale on the Y-axis, the figures suggest reducing calibration error. However, the standard deviations overlap, and hence, these are assumed to be constant across distortions.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig4/Accuracy.png

### How well does watershed perform when used with ReLU activation?

Figure 8 shows the corresponding results. We observe that both the watershed loss and cross-entropy have large overlaps in the standard deviations at all severity levels. So, this shows that, when used in conjunction with ReLU watershed and cross-entropy loss are very similar. But in conjunction with QAct, we see that watershed has a slightly higher Mean-Average-Precision.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig5/Accuracy.png

### What if we consider an easy classification task?

In figure 9, we perform the comparison of QAct+Watershed and ReLU and cross-entropy on MNISTC dataset. Across different architectures, we observe a lot less variation (standard deviation) of QAct+Watershed compared to RelU and cross-entropy. This again suggests robustness against distortions of QAct+Watershed.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig6/Accuracy.png

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig7/Accuracy.png

### Comparing with other popular activations:

Figures 10 and 11 shows the comparison of QAct with ReLU, pReLU and SeLU. We observe the same trend across ReLU, pReLU and SeLU, while QAct is far more stable across distortions.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig8/Accuracy.png

### Results on CIFAR100/TinyImagenetC:

Figure 12 compares QAct+Watershed and ReLU+Cross-Entropy on CIFAR100C dataset. We also include the results of QAct+Cross-Entropy vs. ReLU+Cross-Entropy on TinyImagenetC. The results are consistent with what we observe on CIFAR10C, and hence, draw the same conclusions as before.

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig11/Accuracy.png

Figure: (a) Accuracy
Refer to caption: extracted/5602265/img/article/Fig12/Accuracy.png

### Effect of Quantile Classifier:

Figures 13 and 14 shows the effect of quantile classifier on standard ResNet10/DinoV2 outputs with CIFAR10C/CIFAR100C datasets. While the accuracy values are almost equivalent, we observe a “flatter” trend of the calibration errors, sometimes reducing the error as in the case of CIFAR100C.

## Appendix D Watershed Loss

The authors in [] proposed a novel classifier – *watershed classifier*, which works by learning similarities instead of the boundaries. Below we give the brief idea of the loss function, and refer the reader to the original paper for further details.

- 1.
Let $({\bm{x}}_{i},y_{i})$ denote the samples in each batch, and let $f_{\theta}$ denote the embedding network. $f_{\theta}({\bm{x}}_{i})$ denotes the corresponding embedding.
- 2.
Starting from randomly selected seeds in the batch, propagate the labels to all the samples. Let $\hat{y_{i}}$ denote the estimated samples. For each $f_{\theta}({\bm{x}}_{i})$ and for each label $l$, obtain the nearest neighbour in the samples in the set,
${\mathcal{S}}_{l}=\{f_{\theta}({\bm{x}}_{i})\mid\hat{y_{i}}=y_{i}=l\}$
(9)
that is, all the samples of class $l$ labelled correctly. Denote this nearest neighbour using $f_{\theta}({\bm{x}}_{i,l,1nn})$.
- 3.
Then the loss is given by,
$\text{Watershed Loss}=\sum_{i=1}^{n_{\text{samples}}}\sum_{l=1}^{L}I[y_{i}=l]
\log\left(\frac{\exp\left(-\|f_{\theta}({\bm{x}}_{i})-f_{\theta}({\bm{x}}_{i,l
,1nn})\|\right)}{\sum_{j=1}^{L}\exp\left(-\|f_{\theta}({\bm{x}}_{i})-f_{\theta
}({\bm{x}}_{i,j,1nn})\|\right)}\right)$
(10)

### Why Watershed Loss?:

Observe that the loss in equation 10 implicitly learns representations consistent with the RBF kernel, which is known to be translation invariant. Minimizing this loss function, hence, will learn translation invariant kernels. This is important for obtaining networks robust to distortions.

If one uses (say) Cross Entropy loss, then the features learned would be such that the classes are linearly separable. Contrast this with watershed, which instead learns a similarity between two points in a translation invariant manner.

### Remark:

Observe that the watershed loss is very similar to metric learning losses. The authors in [] claim that this offers better generalization, and show that this is consistent with 1NN classifier. Moreover, they show that this classifier (without considering $f_{\theta}$) has a VC dimension which is equal to the number of classes. While metric learning losses are similar, there is no such guarantee with respect to classification. This motivated our choice of using watershed loss over other metric learning losses.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the weakness.
soundness: 2
strengths: 1. This paper is well-written and easy to follow.
2. This paper considers an interesting scenario: existing ML models cannot correctly classify when several classes have the same probability.
summary: This paper proposes a novel activation function called Quantile Activation (QAct) to enhance model robustness against context distribution shift. Specifically, it adapts the outputs from each neuron to its context distribution by considering the relative quantile of the sample in its context distribution. Experiments demonstrate the superiority of QAct compared with other activation functions.
weaknesses: 1. The motivation and the experiments do not align. In Figure 1(a) and Figure 3(a), the authors emphasize the case where one class is the rotation version of the other. However, the datasets in the experiments are not the rotation version of each class, but the compression version[1]. 
2. In your experiments, the datasets (i.e., CIFAR-10C, CIFAR-100C, TinyImageNet-C, and MNISTC) are also considered by some papers [2, 3] as the covariate shift. Does your proposed method perform well on datasets with other distribution shifts?
3. Could you please re-explain how to obtain the weighted quantiles $\{q_i\}$? Do you augment $z_i$ by grounding the neuron and then re-wright them to obtain $q_i$, or grounding the neuron is the weighting process?

[1] Xie R, Wei H, Feng L, et al. On the importance of feature separability in predicting out-of-distribution error[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2] Viviers C, Valiuddin A, Caetano F, et al. Can Your Generative Model Detect Out-of-Distribution Covariate Shift?[J]. arXiv preprint arXiv:2409.03043, 2024.

[3] Mallinar N, Zane A, Frei S, et al. Minimum-Norm Interpolation Under Covariate Shift[J]. arXiv preprint arXiv:2404.00522, 2024.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: -Line 149 : « This work aims to address the failure mode described earlier. To the best of our knowledge, no existing literature specifically addresses this issue » But isn’t Challa et al. 2023 and the other related works addressing this issue?

-Do you think the approach could be used in the context of meta-learning? And if so, how would the quantile activation be used to gather information about a new context?

-Would the quantile approach be robust enough to handle unbalanced tasks?
soundness: 4
strengths: -Simplicity of the approach
-Thorough analysis (ex. computational complexity) and explanations of the approach
-Many explanations are provided as to why the idea makes sense in practice; think it is generally overlooked in machine learning, yet is key to grasping the inner workings of the approach.
-I don’t have much to say in the following sections: the work is straightforward, clearly explained, and well-supported by empirical evidence.
summary: The authors propose a approach for dealing with a specific failure mode in machine learning : when the same features are equally likely to belonging to various classes. This occurs especially when context or distribution shifts occur. In order to do so, the authors propose using Quantile activation function, which act as « normalizers » to generalize through distribution shifts.
weaknesses: 1. There might be a problem in the possibility to generalize the approach to more complex neural networks architecture. For instance, how would quantile activation be used in architecture involving transformers?

2. This isn’t quite a weakness in itself, but the idea is quite simple (almost naive), such that I’m surprised that this idea hasn’t been explored yet.

3. I find the toy examples quite interesting to understand the logic behind the approach, yet they describe quite unique situations that are unlikely to be met in practice. I think the toy examples would be more convincing if it described situations that are likely to be seen in practice.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 1
questions: See Weaknesses.
soundness: 1
strengths: - Reveal a failure mode in machine learning

- Propose a novel activation fucntion of QACT to deal with classifcation with distribution shift

- The proposed QACT has a good performance on CIFAR-10/100-C and TinyImageNet-C
summary: This manuscirpt reveals a failure mode cased by distribution shift in machine learning. Then the authors correspondingly propose a novel quantile activation fucntion of QACT to replace ReLU for addressing the failure model. The proposed QACT has a good performance on image datasets with distortions including CIFAR-10/100-C and TinyImageNet-C.
weaknesses: This paper does not have a good organization which makes the motivation unclear. The rationale of proposed quantile activation is also not well presented. In addition, only the small image datasets with distortion are not enough to demonstrate the effectiveness of QACT. Please see the details as below:

- **Q1:** This manuscript starts by showing a failure example in binary classification. For me, negative examples are distribution-shifted version of positive examples, thereby making this classification difficult. A formal definition should be provided to show this "failure" or "unlearnability". 

- **Q2:** What causes the "failure mode" and why quantile activation function can address this? Only some intuitive explanations are provided.  It will be much better to provide rigorous empirical or theoretically analysis w.r.t. these two questions.

- **Q3:** According to the manuscript, the computational complexity of QACT is $\mathcal{O} (n\log(n) +Sn_\eta)$, which means the method is hardly scalable to large datasets with significant sample size $S$.

- **Q4:** The experiments are not convincing for me: (1) few and narrow datasets: only the distorted image datasets are used. The proposed QACT should be at least verified on the original version of CIFAR-10/100 and TinyImageNet. Moreover, some other datasets with distirbution shift like ImageNet-Sketch [1] or ImageNet-A [2] can be used to make the results more convincing. (2) unsatisfied performance: even on the basic CIFAR-10/100-C, the proposed QACT does not show the superority when the severity is small (Figure 4(a) and 6(a)).




[1] Wang, Haohan, et al. "Learning robust global representations by penalizing local predictive power." Advances in Neural Information Processing Systems 32 (2019).

[2] Hendrycks, Dan, et al. "Natural adversarial examples." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: * In Algorithm 1, do the indices in $z_i$, $q_i$, and $\tau_i$ refer to the same index? It seems that $z_i$ and $(q_i, \tau_i)$ should have different indices, if I understand correctly.

* In the experiments, why is the watershed loss used only with QACT models? Is there a reason why the watershed loss might not be suitable for non-QACT models? If there is no such reason, non-QACT models trained with the watershed loss should be included as baselines.
soundness: 2
strengths: * Computing activation values based on the relative relationships between pre-activation values within a minibatch is interesting. This approach is somewhat similar to batch normalization (batchnorm) but does not include the trained coefficient and bias terms that batchnorm employs. I believe this is a key feature of the proposed activation function, enabling it to adapt to distribution shifts during inference.
summary: This paper proposes a novel activation function, QACT, which computes its output based on the cumulative distribution function (CDF) of the pre-activation values within a minibatch. In the backward computation of QACT, the probability density function (PDF) of the CDF is estimated using kernel density estimation because the gradient of the CDF is given by the PDF.
The proposed activation function is evaluated using existing datasets with controlled levels of corruption, such as CIFAR-10C, demonstrating that DNN models with QACT achieve better classification accuracy and lower calibration error under strong distortion.
weaknesses: * Terms such as "context," "context distribution," and "failure mode" are used in the manuscript, but their definitions are unclear. Specifically, I could not understand what "context" refers to in the toy example presented in Section 1. The authors should provide clear definitions for these terms, especially for "context," which is a key term in the paper.

* The concept of "context" in the paper appears to be similar to "distribution shift" as described in machine learning literature. Does it relate to distribution shift? If "context" in this paper is related to distribution shift, it should be compared with additional baseline methods developed for distribution shift, in addition to Dino-v2.

Meta-review: The submission proposed a new activation function to address a special-case distribution shift problem. The motivation and contributions are not solid enough to make it accepted. Please try to make your writing more scientific in the next version, for example, to formally define every key concept besides giving intuition about the concept.

Decision: Reject