Self-Boost via Optimal Retraining: An Analysis via Approximate Message Passing

Thanks for your review and feedback! You have raised good points and we respond to your concerns and questions below.

Extension of theory to multi-class case. This can be done by following similar ideas in the paper for the binary case, but will be heavier in notation and will involve more tedious analysis. Here we outline how this extension can be made for the GMM case in Section 3. We first describe the setting. We define a label matrix $Y \in \mathbb{R}^{n \times k}$, where $n$ is the number of samples and $k$ is the number of classes, using one-hot encoding, where column $i$ has $1$ for the samples belonging to class $i$ and $0$ otherwise. Next, we define $M \in \mathbb{R}^{k \times d}$ where row $i$ is the mean feature vector of class $i$. Equation (1) then transforms to $X = YM+Z$, where $Z$ is random noise; so this is a low rank matrix ($YM$ is of rank $k$) plus noise. We can then generalize the AMP iterations for this low-rank model. The forms of eqs. (4) and (5) in this case will be the same with $g_t$ redefined appropriately for the matrix case (here it will be applied row-wise instead of coordinate-wise). Also the coefficient $C_t$ in line 126 will now be a $k \times k$ matrix, defined similarly using the Jacobian of $g_t$ instead of derivative. Likewise, the state evolution recursion (7) will be of the same form, with the modification that $m_t, \sigma_t, \bar{m_t}, \bar{\sigma}_t$ are now vectors of length $k$. The insights from the analysis will be similar but the analysis will be on multi-dimensional objects. A similar extension can be made for GLMs in Section 4, where $\beta$ should now be defined as a matrix of size $k \times d$. The extension here for the AMP iterates (16-17) and the state evolution (21) are similar to the GMM case described above.

Extension of theory to non-linear models. We note that the AMP techniques can be used to analyze the statistical behavior of more general M-estimators (see for e.g., reference [14] in the paper). The corresponding update rule will involve applying the aggregator function (aggregating current predictions and the given noisy labels) as well as a gradient descent type of update (similar to eq. (4)). The exact details of these updates and the analysis will be quite tedious and are left for future work.

Extension in practice to different types of label noise (beyond uniform) and general ML models. Although our theoretical results focused on linear models, the algorithm is Section 5 proposed for use in practice, i.e. BayesMix RT, is based on insights that carry over to general ML models. Specifically, eq. (27) is obtained by fitting a bimodal GMM on the distributions of the unnormalized logits (following the derivation in eq. (30) for theory). Following the same approach (and even in multi-class setting), we can rewrite the Bayesian posterior of the class memberships, after observing the current model's unnormalized logits with the prior based on the label noise model (in our paper, we have considered the uniform prior). We can then derive the corresponding aggregating function $g$ similar to eq. (27) in the paper. In short, BayesMix RT can be used in deep learning models and even with non-uniform label noise priors (e.g., when the flipping probability from class i to class j is some $p_{i, j}$). As an example, in the binary case, consider the non-uniform noise model: $\mathbb{P}(\hat{y} = -1| y = +1) = p$ and $\mathbb{P}(\hat{y} = +1| y = -1) = q$ with $p \neq q$. Then, eq. (27) changes to: $g(z, \hat{y}) = \frac{2}{1+ \Big(\frac{q(1-q)}{p(1-p)}\Big) \Big(\frac{p q}{(1-p)(1-q)}\Big)^{\hat{y}/2} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1.$ We show a result under the above non-uniform noise model with $p=0.45$ and $q=0.2$ for the MedMNIST Pneumonia dataset in the same setting as Section 5 using the formula above for BayesMix RT; the table below lists the average test accuracies $\pm$ standard deviation (over 3 runs) in this case. As we see below, BayesMix RT performs the best.

For more general types of label noise (e.g., sample-dependent noise) or in scenarios where the noise model is not known, we propose the following simpler aggregation function (instead of eq. (27)): $g(z, \hat{y}) = \frac{2}{1+ \gamma^{\hat{y}} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1,$ for some constant $\gamma \in (0,1)$ which can be tuned. We call this simpler version BayesMix-Simple RT (everything else is the same as BayesMix RT in the paper); in our implementation, we set $\gamma = 0.7$ throughout without any tuning. Also note that this can be used with general ML models.

(1) To show the efficacy of BayesMix-Simple RT, we first consider an adversarial noise model. Specifically, an adversary gets to determine the "important" samples and flip their labels. The adversary determines the important samples by training the same model which will be trained by us with clean labels and orders the samples in the decreasing order of the absolute value of the unnormalized logits. Note that a higher logit absolute value implies that the model is very confident on the sample, so flipping the label of such a sample would be more detrimental to the model's training as it would need to adjust the decision boundary by a large amount for this incorrectly labeled sample. The adversary flips the labels of the top $\alpha$ fraction of the samples with largest absolute values of the unnormalized logits. Under this adversarial noise model (we don't assume knowledge of this process in training to be clear) with $\alpha=0.25$ for MedMNIST Pneumonia in the same setting as Section 5, we show the efficacy of BayesMix-Simple RT below compared to the baselines (again, the table below lists the average test accuracies $\pm$ standard deviation over 3 runs).

While we did not have time to find and work with naturally noisy datasets, we believe the above result under adversarial noise clearly demonstrates the strength of our method. But in the future, we will test our algorithm on naturally noisy datasets as you suggested.

(2) We now show the efficacy of BayesMix-Simple RT even for deep network training. Recall that our earlier experiments are for linear probing with a pre-trained ResNet-50. Specifically, here we consider training an entire pre-trained ResNet-50. Because of the heavy computational cost of training an entire deep network, we only do 3 iterations of retraining here. We show results for two cases of Food-101:

i. Pho vs. Ramen (just like in the paper) corrupted by uniform noise with $p=0.4$

ii. Spaghetti Bolognese vs. Spaghetti Carbonara (not in the paper) corrupted by non-uniform noise described above with $p=0.45$ and $q=0.2$

Note that in the above two tables (for entire ResNet-50 training), BayesMix-Simple RT leads to significant gains over the initial model, while Full RT and Consensus-based RT are actually detrimental.

So in summary, our proposed method is effective with different label noise models as well as for general ML models.

While we did not have enough time to concretely work on the multi-class case, we believe our current contributions for binary classification are worthy and fundamental enough in their own right and pave the way for the multi-class case in the future.

We agree our paper is pretty dense and math-heavy. We will provide some intuition about the theoretical results in the next version. We encourage you to take a look at our discussion for "Interpretation of the Aggregator Function" in our response to Reviewer NMag (sorry we couldn't include it here due to lack of space).

Thank you for pointing out the SELC paper; it is indeed relevant and we will cite it in the next version. However, please note that SELC is an empirical paper and it does not theoretically prove that their proposed method is optimal in any sense. On the other hand, our paper is primarily theoretical and our main goal is to derive the Bayes' optimal way to combine the given labels and the model's predictions for linear models.

We will add all this discussion in the next version of our paper. We hope to have resolved your concerns and questions, but we're happy to discuss further if needed. If you're satisfied with our responses, we sincerely hope you will raise your score!

Thanks for your review and feedback! You have raised great points and we address your questions below.

Extension in practice to different types of label noise (beyond uniform) and general ML models (Questions 1 and 4 & Weaknesses 1, 2 and 4). Although our theoretical results focused on linear models, the algorithm is Sec. 5 proposed for use in practice, i.e. BayesMix RT, is based on insights that carry over to general ML models. Specifically, eq. (27) is obtained by fitting a bimodal GMM on the distributions of the unnormalized logits (following the derivation in eq. (30) for theory). Following the same approach, we can rewrite the Bayesian posterior of the class memberships, after observing the current model's unnormalized logits with the prior based on the label noise model (in our paper, we have considered the uniform prior). We can then derive the corresponding aggregating function $g$ similar to eq. (27) in the paper. In short, BayesMix RT can be used in deep learning models and even with non-uniform label noise priors (e.g., when the flipping probability from class i to j is some $p_{i, j}$). As an example, in the binary case, consider the non-uniform noise model: $\mathbb{P}(\hat{y} = -1| y = +1) = p$ and $\mathbb{P}(\hat{y} = +1| y = -1) = q$ with $p \neq q$. Then, eq. (27) changes to: $g(z, \hat{y}) = \frac{2}{1+ \Big(\frac{q(1-q)}{p(1-p)}\Big) \Big(\frac{p q}{(1-p)(1-q)}\Big)^{\hat{y}/2} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1.$ We show a result under the above non-uniform noise model with $p=0.45$ and $q=0.2$ for the MedMNIST Pneumonia dataset in the same setting as Sec. 5 using the formula above for BayesMix RT; the table below lists the average test accuracies $\pm$ std. (over 3 runs) in this case. Note BayesMix RT performs the best.

For more general types of label noise (e.g., sample-dependent noise) or in scenarios where the noise model is not known, we propose the following simpler aggregation function (instead of (27)): $g(z, \hat{y}) = \frac{2}{1+\gamma^{\hat{y}} \exp\Big(\frac{(z-\mu_{+})^2}{2\sigma_{+}^2} - \frac{(z-\mu_{-})^2}{2\sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1,$ for some constant $\gamma \in (0,1)$ which can be tuned. We call this simpler version BayesMix-Simple RT (everything else is the same as BayesMix RT in the paper); in our implementation, we set $\gamma=0.7$ throughout w/o any tuning. Also note that this can be used with general ML models.

(1) To show the efficacy of BayesMix-Simple RT, we first consider an adversarial noise model. Specifically, an adversary gets to determine the "important" samples and flip their labels. The adversary determines important samples by training the same model which will be trained by us with clean labels and orders the samples in the decreasing order of the absolute value of the unnormalized logits. Note that a higher logit absolute value implies that the model is very confident on the sample, so flipping the label of such a sample would be more detrimental to the model's training as it would need to adjust the decision boundary by a large amount for this incorrectly labeled sample. The adversary flips the labels of the top $\alpha$ fraction of the samples with largest absolute values of the unnormalized logits. Under this adversarial noise model (we don't assume knowledge of this process in training to be clear) with $\alpha=0.25$ for MedMNIST Pneumonia in the same setting as Sec. 5, we show the efficacy of BayesMix-Simple RT below compared to the baselines (table below lists the average test accuracies $\pm$ std. over 3 runs).

(2) We now show the efficacy of BayesMix-Simple RT even for deep network training. Recall that our earlier experiments are for linear probing with a pre-trained ResNet-50. Specifically, here we consider training an entire pre-trained ResNet-50. Because of the heavy computational cost of training an entire ResNet-50, we only do 3 iterations of retraining here. We show results for two cases of Food-101:

i. Pho vs. Ramen corrupted by uniform noise with $p=0.4$

ii. Spaghetti Bolognese vs. Spaghetti Carbonara corrupted by non-uniform noise described above with $p=0.45$ and $q=0.2$

Note that in the above two tables (for entire ResNet-50 training), BayesMix-Simple RT leads to significant gains over the initial model, while Full RT and Consensus-based RT are actually detrimental.

So in summary, our proposed method is effective with different label noise models as well as for general ML models.

Extension of theory to non-linear models (Question 4 & Weakness 1). We note that the AMP techniques can be used to analyze the statistical behavior of more general M-estimators (see for e.g., reference [14] in the paper). The corresponding update rule will involve applying the aggregator function (aggregating current predictions and the given noisy labels) as well as a gradient descent type of update (similar to eq. (4)). The exact details of these updates and the analysis will be quite tedious and are left for future work.

Extension of theory to multi-class case (Weakness 1). This can be done by following similar ideas in the paper for the binary case, but will involve more tedious analysis. Here we outline how this extension can be made for the GMM case in Sec. 3. We first describe the setting. We define a matrix $Y \in \mathbb{R}^{n \times k}$, where $n$ is the # of samples and $k$ is the # of classes, using one-hot encoding, where column $i$ has $1$ for the samples belonging to class $i$ and $0$ otherwise. Next, we define $M \in \mathbb{R}^{k \times d}$ where row $i$ is the mean feature vector of class $i$. Eq. (1) then transforms to $X = YM+Z$, where $Z$ is random noise; so this is a low rank matrix ($YM$ is of rank $k$) plus noise. We can then generalize the AMP iterations for this low-rank model. The forms of eqs. (4) & (5) in this case will be the same with $g_t$ redefined appropriately for the matrix case (here it will be applied row-wise instead of coordinate-wise). Also the coefficient $C_t$ in line 126 will now be a $k \times k$ matrix, defined similarly using the Jacobian of $g_t$ instead of derivative. Likewise, the state evolution recursion (7) will be of the same form, with the modification that $m_t, \sigma_t, \bar{m_t}, \bar{\sigma}_t$ are now vectors of length $k$. The insights from the analysis will be similar but the analysis will be on multi-dimensional objects. A similar extension can be made for GLMs in Sec. 4, where $\beta$ should now be defined as a matrix of size $k \times d$. The extension here for the AMP iterates (16-17) and the state evolution (21) are similar to the GMM case described above.

Interpretation of the Aggregator Function (Question 2). Consider the aggregator in eq. (15) for the GMM case (similar insights hold for GLM and BayesMix RT). Recall that $g(y,\hat{y})$ is the soft label obtained based on the soft prediction $y$ and the given noisy label $\hat{y} \in$ {$+1,-1$}. It is easy to see that $g(y,\hat{y})\in[-1,1]$, it is monotonic increasing in $y$ and $\hat{y}$ (since $p<1/2$). Also it is increasing in $\pi_+/\pi_-$, accounting for class probabilities (so if $\pi_+>\pi_-$, our aggregated labels are biased towards the positive class). Also note the denominator can be rewritten as $1+\exp\Big(-\log(\frac{1-p}{p})\hat{y} - 2y\Big) \frac{\pi_-}{\pi_+}$; so the term inside the exponential is a negative linear combination of $\hat{y}$ and $y$. Now if the label flipping probability $p$ is small, we put a larger weight on the given labels $\hat{y}$ relative to the predictions $y$, but as $p$ gets larger we reduce the weight on $\hat{y}$, thereby relying more on the predictions $y$. These behavior trends of the aggregator are consistent with what we expect intuitively.

Scalability and Runtime Considerations (Question 3). We focus on computational overhead of BayesMix(-Simple) RT as that's used in practice. The major extra cost compared to Full RT & Consensus RT is in fitting the bimodal GMM to obtain the aggregation function $g$. But this is negligible when compared to the training (backward & forward propagation) cost of every epoch, especially for large models. So our method doesn't introduce too much extra overhead. Note that even the AMP update rules for theory in Sec. 3 & 4 don’t introduce much extra overhead; the cost is dominated by the 2 matrix-vector multiplications involving $X$ and $X^\top$.

Theory for BayesMix RT is left for future work.

We hope to have resolved your concerns and questions, but we're happy to discuss further if needed. If you're satisfied with our responses, we hope you will raise your score!

Thanks for engaging in discussion! In our response “Extension in practice to different types of noise and general ML models”, we discussed how BayesMix RT and its underlying intuition (namely fitting a gaussian mixture model to soft predictions at each step and using Bayes' optimal estimates) can be used for general ML models. To that end, we also presented empirical results for deep learning in our rebuttal above (see the last two tables for Pho vs. Ramen and Spaghetti Bolognese vs. Spaghetti Carbonara). Here we focus our response on AMP techniques, its applicability to nonlinear models and the challenges.

AMP techniques can be used to analyze the statistical behavior of the so-called “general first order methods (GFOM)”. Consider estimation problems where the data is given as a random matrix $X$ and the goal is to estimate an unknown model parameter. These algorithms keep state sequences of vectors (or in more general form matrices) $u^1,\dotsc, u^t$ and $v^1,\dotsc, v^t$ which are updated by two types of operations: row-wise application of a (possibly non-linear) function, or multiplication by $X$ or $X^\top$. This cover a broad range of important practical problems, for example classical first order optimization (for convex and non-convex objectives), such as gradient descent and its accelerated versions where the loss is a function of $X\theta$ (note that the gradient updates here become $X^\top \nabla L(X\theta, y)$ and follows the above GFOM form). Other examples include sparse PCA, low rank matrix estimation, phase retrieval, nonlinear power methods, finding planted hidden clique, among others (see references below). Each of these problems is an important application which has attracted an immense interest and significant research effort. So in short AMP provides a very powerful class of iterative algorithms, for non-linear problems, and comes with an asymptotic "precise" statical characterization. That said, it cannot be applied to general deep NNet architecture, for which the updates do not follow the form of general first order methods, discussed above.

To outline the challenges for using these techniques (in its general form) for our problem, we should note that in AMP algorithms (in supervised learning), the given labels are fixed across iterations and can be used in the nonlinear function updates. However, here we iteratively use the current model to make predictions and then aggregate them with the initially given noisy labels to retrain the model. This complication requires new derivations, which we did for two data models: Gaussian Mixture model and generalized linear model, with linear classifier. Given the success of AMP in other nonlinear setting we expect the derivations to be extendable to other nonlinear classifiers, but a rigorous study is out of the scope of the current work.

We hope that we have addressed your concerns. Please let us know if you have any other question!

[1] The estimation error of general first order methods, M Celentano, A Montanari, Y Wu, Conference on Learning Theory, 2020

[2] Information-theoretically optimal sparse PCA, Y Deshpande, A Montanari, IEEE International Symposium on Information Theory, 2014

[3] Finding Hidden Cliques of Size $\sqrt{n/e}$ in Nearly Linear Time, Y Deshpande, A Montanari, Foundations of Computational Mathematics 15 (4), 1069-1128

[4] Statistically optimal first order algorithms: a proof via orthogonalization, A Montanari, Y Wu, Information and Inference: A Journal of the IMA, 2024

[5] A unifying tutorial on 342 approximate message passing, O. Y. Feng, R. Venkataramanan, C. Rush, R. J. Samworth, Foundations and Trends in Machine Learning, 15(4):335–536, 343 2022.

[6] Optimization-Based AMP for Phase Retrieval: The Impact of Initialization and $\ell_2$ Regularization, J Ma, J Xu, A Maleki IEEE Transactions on Information Theory, 2019

Thanks for your detailed review and positive assessment of our work! You have raised great points and we address your questions and concerns below.

Extension of theory to multi-class case. This can be done by following similar ideas in the paper for the binary case, but will involve more tedious analysis. Here we outline how this extension can be made for the GMM case in Section 3. We first describe the setting. We define a label matrix $Y \in \mathbb{R}^{n \times k}$, where $n$ is the # of samples and $k$ is the # of classes, using one-hot encoding, where column $i$ has $1$ for the samples belonging to class $i$ and $0$ otherwise. Next, we define $M \in \mathbb{R}^{k \times d}$ where row $i$ is the mean feature vector of class $i$. Eq. (1) then transforms to $X = YM+Z$, where $Z$ is random noise; so this is a low rank matrix ($YM$ is of rank $k$) plus noise. We can then generalize the AMP iterations for this low-rank model. The forms of eqs. (4) and (5) in this case will be the same with $g_t$ redefined appropriately for the matrix case (here it will be applied row-wise instead of coordinate-wise). Also the coefficient $C_t$ in line 126 will now be a $k \times k$ matrix, defined similarly using the Jacobian of $g_t$ instead of derivative. Likewise, the state evolution recursion (7) will be of the same form, with the modification that $m_t, \sigma_t, \bar{m_t}, \bar{\sigma}_t$ are now vectors of length $k$. The insights from the analysis will be similar but the analysis will be on multi-dimensional objects. A similar extension can be made for GLMs in Section 4, where $\beta$ should now be defined as a matrix of size $k \times d$. The extension here for the AMP iterates (16-17) and the state evolution (21) are similar to the GMM case described above.

Dealing with unknown label flip probability $p$. Great question! We propose the following way to estimate $p$. We focus on the samples for which the model is the most confident after training with the given noisy labels. Specifically, we pick the set of samples $S_\text{conf}$ with the highest $\rho$ (this should be small) fraction of absolute value of unnormalized logits; a higher value indicates higher confidence. The philosophy of self-training based ideas is that the model is probably correct on the samples for which it is very confident; so we treat the model's predicted labels for the samples in $S_\text{conf}$ as the ground truth labels and calculate the % of samples in $S_\text{conf}$ for which the predicted label matches the given label. This gives us an estimate for $p$. We tested this approach for MedMNIST Pneumonia dataset in the same setting as Sec. 5; here are the results averaged over 3 runs.

So this gives us a good estimate. We also saw that using the estimated $p$'s above basically didn't change the results at all. Noting this and also taking into account that the estimated $p$ may be off by a lot sometimes, we tested the following simpler aggregator function (as a substitute for eq. (27)) independent of $p$: $g(z, \hat{y}) = \frac{2}{1+ \gamma^{\hat{y}} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1,$ for some constant $\gamma \in (0,1)$ which can be tuned. Note that $\gamma$ essentially plays the role of $p/(1-p)$. We call this simpler version BayesMix-Simple RT (everything else is the same as BayesMix RT in the paper); in our implementation, we set $\gamma = 0.7$ throughout w/o any tuning. We compare test accuracies (mean $\pm$ std. averaged over 3 runs) for MedMNIST Pneumonia obtained after 10 iterations of Full RT, Consensus RT, BayesMix RT with the exact $p$ and BayesMix-Simple RT (with $\gamma = 0.7$).

Notice that BayesMix-Simple RT is competitive with BayesMix RT even though it does not use the value of $p$.

Regarding your question about fitting a GMM to the logits, note that in reality the distribution is unlikely to be a GMM exactly. Yet our method consistently leads to good performance (we show results with more noise models below).

Extension in practice to different types of label noise and general ML models. BayesMix RT is based on insights that carry over to general ML models. Specifically, eq. (27) is obtained by fitting a bimodal GMM on the distributions of the unnormalized logits (following the derivation in eq. (30) for theory). Following the same approach, we can rewrite the Bayesian posterior of the class memberships, after observing the current model's unnormalized logits with the prior based on the label noise model. We can then derive the corresponding aggregating function $g$ similar to eq. (27) in the paper. In short, BayesMix RT can be used in deep learning models and even with non-uniform label noise priors (e.g., when the flipping probability from class i to class j is some $p_{i, j}$). As an example, in the binary case, consider the non-uniform noise model: $\mathbb{P}(\hat{y} = -1| y = +1) = p$ and $\mathbb{P}(\hat{y} = +1| y = -1) = q$ with $p \neq q$. Then, eq. (27) changes to: $g(z, \hat{y}) = \frac{2}{1+ \Big(\frac{q(1-q)}{p(1-p)}\Big) \Big(\frac{p q}{(1-p)(1-q)}\Big)^{\hat{y}/2} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1.$ We show a result under the above non-uniform noise model with $p=0.45$ and $q=0.2$ for the MedMNIST Pneumonia dataset in the same setting as Sec. 5 using the formula above for BayesMix RT; the table below lists the average test accuracies $\pm$ std. (over 3 runs) in this case. As we see below, BayesMix RT performs the best.

For more general types of label noise or in scenarios where the noise model is not known, we propose to use the BayesMix-Simple RT algorithm described above as it is independent of any noise distribution parameter and is competitive with BayesMix RT in the uniform noise case. Also note that this can be used with general ML models.

(1) To show the efficacy of BayesMix-Simple RT, we first consider an adversarial noise model. Specifically, an adversary gets to determine the "important" samples and flip their labels. The adversary determines important samples by training the same model which will be trained by us with clean labels and orders the samples in the decreasing order of the absolute value of the unnormalized logits. Note that a higher logit absolute value implies that the model is very confident (and probably correct) on the sample, so flipping the label of such a sample would be more detrimental to the model's training as it would need to adjust the decision boundary by a large amount for this incorrectly labeled sample. The adversary flips the labels of the top $\alpha$ fraction of the samples with largest absolute values of the unnormalized logits. Under this adversarial noise model (we don't assume knowledge of this process in training to be clear) with $\alpha=0.25$ for MedMNIST Pneumonia in the same setting as Sec. 5, we show the efficacy of BayesMix-Simple RT below compared to the baselines (table below lists average test accuracies $\pm$ std. over 3 runs).

(2) We now show the efficacy of BayesMix-Simple RT for deep network training. Recall that our earlier experiments are for linear probing with a pre-trained ResNet-50. Specifically, here we consider training an entire pre-trained ResNet-50. Because of the heavy computational cost of training an entire ResNet-50, we only do 3 iterations of retraining here. We show results for 2 cases of Food-101:

i. Pho vs. Ramen corrupted by uniform noise with $p=0.4$

ii. Spaghetti Bolognese vs. Spaghetti Carbonara corrupted by non-uniform noise described above with $p=0.45$ and $q=0.2$

Note that in the above two tables (for entire ResNet-50 training), BayesMix-Simple RT leads to significant gains over the initial model, while Full RT and Consensus RT are actually detrimental; the latter also answers your question about whether retraining can hurt performance in some cases. It is possible that BayesMix-Simple RT hurts performance in some cases.

But in summary, our proposed method is effective with different label noise models as well as for general ML models.

We hope to have resolved your concerns and questions, but we're happy to discuss further if needed.

Thanks for your review and feedback! You have raised good points and we respond to your concerns and questions below.

Extension of theory to multi-class case. This can be done by following similar ideas in the paper for the binary case, but will be heavier in notation and will involve more tedious analysis. Here we outline how this extension can be made for the GMM case in Section 3. We first describe the setting. We define a label matrix $Y \in \mathbb{R}^{n \times k}$, where $n$ is the number of samples and $k$ is the number of classes, using one-hot encoding, where column $i$ has $1$ for the samples belonging to class $i$ and $0$ otherwise. Next, we define $M \in \mathbb{R}^{k \times d}$ where row $i$ is the mean feature vector of class $i$. Equation (1) then transforms to $X = YM+Z$, where $Z$ is random noise; so this is a low rank matrix ($YM$ is of rank $k$) plus noise. We can then generalize the AMP iterations for this low-rank model. The forms of eqs. (4) and (5) in this case will be the same with $g_t$ redefined appropriately for the matrix case (here it will be applied row-wise instead of coordinate-wise). Also the coefficient $C_t$ in line 126 will now be a $k \times k$ matrix, defined similarly using the Jacobian of $g_t$ instead of derivative. Likewise, the state evolution recursion (7) will be of the same form, with the modification that $m_t, \sigma_t, \bar{m_t}, \bar{\sigma}_t$ are now vectors of length $k$. The insights from the analysis will be similar but the analysis will be on multi-dimensional objects. A similar extension can be made for GLMs in Section 4, where $\beta$ should now be defined as a matrix of size $k \times d$. The extension here for the AMP iterates (16-17) and the state evolution (21) are similar to the GMM case described above.

Extension of theory to non-linear models. We note that the AMP techniques can be used to analyze the statistical behavior of more general M-estimators (see for e.g., reference [14, section 4.4] in the paper). The corresponding update rule will involve applying the aggregator function (aggregating current predictions and the given noisy labels) as well as a gradient descent type of update (similar to eq. (4)). The exact details of these updates and the analysis will be quite tedious and are left for future work.

Extension in practice to different types of label noise (beyond uniform) and deep learning models. Although our theoretical results focused on linear models, the algorithm is Section 5 proposed for use in practice, i.e. BayesMix RT, is based on insights that carry over to general ML models. Specifically, eq. (27) is obtained by fitting a GMM on the distribution of the unnormalized logits (following the derivation in eq. (30) for theory). Following the same approach (and even in multi-class setting), we can rewrite the Bayesian posterior of the class memberships, after observing the current model's unnormalized logits with the prior based on the label noise model (in our paper, we have considered the uniform prior). We can then derive the corresponding aggregating function $g$ similar to eq. (27) in the paper. In short, BayesMix RT can be used in deep learning models and even with non-uniform label noise priors (e.g., when the flipping probability from class i to class j is some $p_{i, j}$). As an example, in the binary case, consider the non-uniform noise model: $\mathbb{P}(\hat{y} = -1| y = +1) = p$ and $\mathbb{P}(\hat{y} = +1| y = -1) = q$ with $p \neq q$. Then, eq. (27) changes to: $g(z, \hat{y}) = \frac{2}{1+ \Big(\frac{q(1-q)}{p(1-p)}\Big) \Big(\frac{p q}{(1-p)(1-q)}\Big)^{\hat{y}/2} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1.$ We show an experimental result under the above non-uniform noise model with $p=0.45$ and $q=0.2$ for the MedMNIST Pneumonia dataset in the same setting as Section 5 using the formula above for BayesMix RT; the table below lists the average test accuracies $\pm$ standard deviation (over 3 runs) in this case. As we see below, BayesMix RT performs the best.

For more general types of label noise (e.g., sample-dependent noise) or in scenarios where the noise model is not known, we propose the following simpler aggregation function (instead of eq. (27)): $g(z, \hat{y}) = \frac{2}{1+ \gamma^{\hat{y}} \exp\Big(\frac{(z - \mu_{+})^2}{2 \sigma_{+}^2} - \frac{(z - \mu_{-})^2}{2 \sigma_{-}^2}\Big) \frac{\pi_-}{\pi_+}} - 1,$ for some constant $\gamma \in (0,1)$ which can be tuned. We call this simpler version BayesMix-Simple RT (everything else is the same as BayesMix RT in the paper); in our implementation, we set $\gamma = 0.7$ throughout without any tuning though. Please note that this can be used with general ML models.

(1) To show the efficacy of BayesMix-Simple RT, we first consider an adversarial noise model. Specifically, an adversary gets to determine the "important" samples and flip their labels. The adversary determines important samples by training the same model which will be trained by us with clean labels and orders the samples in the decreasing order of the absolute value of the unnormalized logits. Note that a higher logit absolute value implies that the model is very confident on the sample, so flipping the label of such a sample would be more detrimental to the model's training as it would need to adjust the decision boundary by a large amount for this incorrectly labeled sample. The adversary then flips the labels of the top $\alpha$ fraction of the samples with largest absolute values of the unnormalized logits. Under this adversarial noise model (we don't assume knowledge of this process in training to be clear) with $\alpha=0.25$ for MedMNIST Pneumonia in the same setting as Section 5, we show the efficacy of BayesMix-Simple RT below compared to the baselines (again, the table below lists the average test accuracies $\pm$ standard deviation over 3 runs).

(2) We now show the efficacy of BayesMix-Simple RT even for deep network training; recall that our earlier experiments are for linear probing with a pre-trained ResNet-50. Specifically, here we consider training an entire pre-trained ResNet-50; this is a typical deep learning setting. Because of the heavy computational cost of training an entire deep network, we only do 3 iterations of retraining here. We show results for two cases of Food-101:

i. Pho vs. Ramen corrupted by uniform noise with $p=0.4$

ii. Spaghetti Bolognese vs. Spaghetti Carbonara corrupted by non-uniform noise described above with $p=0.45$ and $q=0.2$

Note that in the above two tables (for full ResNet-50 training), BayesMix-Simple RT leads to significant gains over the initial model, while Full RT and Consensus-based RT are actually detrimental.

So in summary, our proposed method is effective with different label noise models as well as for general ML models.

Regarding your question about practical computational overhead of our methods BayesMix RT and BayesMix-Simple RT, note that the major extra cost compared to Full RT and Consensus-based RT is in fitting the bimodal GMM to obtain the aggregation function $g$. But this is negligible when compared to the training (backward and forward propagation) cost of every epoch, especially for large models. So our proposed methods do not introduce too much extra overhead. It is worth mentioning that for larger models, having multiple rounds of retraining will be expensive. Also, just to clarify the AMP updates in eqs. (4-5) and (16-17) are for theory only. While the insights obtained from the AMP analysis are used to propose BayesMix/BayesMix-Simple RT, they do not use the specific AMP updates. Having said that, even if we consider the AMP updates in eqs. (4-5) for GMM (Sec 3), our algorithm has the same effective per-iteration complexity as Full RT and Consensus-based RT of $O(nd)$ due to the two matrix-vector multiplications (those involving $X$ and $X^\top$). The only difference is in applying the optimal aggregator $g_t$ and also computing the correction (Onsager) term $C_t$; both of these are $O(n)$ which is negligible compared to $O(nd)$ for large $d$. A similar argument holds for the GLM case (Sec 4).

We will add all this discussion in the next version of our paper. We hope to have resolved your concerns and questions, but we're happy to discuss further if needed. If you're satisfied with our responses, we sincerely hope you will raise your score!