Online Feature Updates Improve Online (Generalized) Label Shift Adaptation

We thank the reviewer for providing insightful comments! Here are our responses.

Motivation and its justification. We would like to answer your questions in this point separately.

• “Why is improving feature extractors helpful for only label shift?” In fact, updating feature extractor can help with other types of shift too. As discussed at line 359-361, the existing algorithm to solve online covariate shift, which has strong theoretical guarantee, also doesn’t take feature extractor update into account. We believe by adopting online feature update, the performance of the algorithm can get certain improvement as well, and our work shows this possibility as a first work to introduce feature extractor update into those theoretical algorithms.
• “What can prove feature extractors gain? Are there some theoretical or empirical results to prove that?” Equation (7) describes the gain from the feature extractor improvement rather than the label shift adaptation, because it is “almost” the best loss at time $t$ given the updated feature extractor $f_t’’$ or the old feature extractor $f_0$ together with ** the knowledge of label distribution $q_t$**. On the one hand, the theoretical guarantee can be improved from the feature extractor when Equation (7) holds; on the other hand, in the experiment, we empirically validate the holdness of equation 7. We believe these two theoretical and empirical evidence explain how the feature extractor updates improve our target problem.

We will write those arguments more explicitly in our revision!

Relation to the related work [1]. Thanks for connecting our work to [1]. First, as cited in section 2, we introduce it as a self-supervised feature update algorithm, which can serve for our motivations discussed at line 135-141. Moreover, we discussed in Section 3.4 on how it is particularly powerful to address online generalized label shift. We would like to clarify this point here:

According to the definition of generalized label shift, it can be reduced to label shift once the learner knows the underlying feature mapping $h$, and online feature update actually helps learn the underlying feature mapping $h$ that makes $P(h(x_t)|y_t)$ invariant. In the literature of test-time training [1] and more [2, 3], they found the self-supervised learning of feature updates can explicitly align the feature space of new distribution to the feature space of training distribution; they even only trained with the feature extractor without retrain the last linear layer, this strongly supports the straightforward feature space alignment between original distribution and shifted distribution.

We would like to add this explanation in Section 3.4, which we believe can better explain the role of [1] when tackling the online generalized label shift.

Our main contribution. Our main contribution is as a first work to bring self-supervised learning into the online label shift problem, which was mainly studied from the theoretical aspects. Our algorithm provides a way to leverage self-supervised learning, which largely improves the performance as validated through the experiments and still keeps the similar theoretical guarantees.

Limitations and broader impacts. The goal of this work is to advance the robustness of models for real-world deployment. Our work contributes to the mitigation of different adverse effects of online label shift, such as out-of-date or miscalibrated models (e.g. healthcare or finance settings, autonomous systems such as self-driving). Adaptation to changing shifts has positive ethical implications, such as in fairness (e.g. improving the models with updated, fair data over time) or privacy (e.g. unlearning data owned by specific groups).

Difference between label distribution shift and domain adaptation. Domain adaptation is a broader concept, including the case where the support of x can be totally different. Label distribution shift has an explicit assumption where $P(x|y)$ doesn’t change. Our algorithm and previous label shift adaptation algorithms are studied given this explicit assumption and the theoretical results are based on this assumption. While those algorithms work well for label shift, there is no guarantee that they can work well without the label shift assumption.

How to handle unseen categories. This is actually an interesting question! As pointed out by Reviewer aV82, [4] focuses on this special setting by unsupervised estimating the portion of unseen data. We will add the discussion of this setting and the related work [4] into the related work section.

[1] Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229–9248. PMLR, 2020.

[2] Wang, Dequan, et al. "Tent: Fully test-time adaptation by entropy minimization." arXiv preprint arXiv:2006.10726 (2020).

[3] Liu, Yuejiang, et al. "Ttt++: When does self-supervised test-time training fail or thrive?." Advances in Neural Information Processing Systems 34 (2021): 21808-21820.

[4] Qian Y Y, Bai Y, Zhang Z Y, et al. Handling New Class in Online Label Shift[C]//2023 IEEE International Conference on Data Mining (ICDM). IEEE, 2023: 1283-1288.

We thank the reviewer for providing insightful comments! Here are our responses.

The storage of previous historical data. This is actually an insightful point! We acknowledge this can bring additional privacy concerns, depending on the practical scenarios, and we would like to discuss this in the revision. At the same time, as shown in Table 1, we can find that even without batch accumulation (batch size $\tau$=1), OLS-OFU still outperforms OLS. Parameter $\tau$ can be treated as the performance and privacy trade-off.

Dataset set-up. Our selection of datasets and the way to simulate the shift pattern mainly follow the previous online label shift literature [1, 2] and additionally, we experiment with additional domain adaptation dataset CIFAR-10C. Although we agree experimenting with real-world shifts can be meaningful, we believe our current experiments are sufficient enough to prove the improvement from the literature.

Related work. Thank you for the pointer! This work particularly solves the unseen categories in the online label shift setting and is very relevant. It tackles the new unseen class by unsupervised estimating the portion of unseen data. We will add this literature in our related work section.

[1] Baby, Dheeraj, et al. "Online label shift: Optimal dynamic regret meets practical algorithms." Advances in Neural Information Processing Systems 36 (2024).

[2] Bai, Yong, et al. "Adapting to online label shift with provable guarantees." Advances in Neural Information Processing Systems 35 (2022): 29960-29974.

We would like to thank the reviewer for providing some useful comments about data privacy concerns and some typos. However, we respectfully disagree with the criticism on novelty, paper structure and the experiment set-up (data split and baseline performance). We reply to the weaknesses and questions point by point as below.

Novelty of three principles. While we agree that it turns out the steps are simple, the choice of them rather than other options is out of careful considerations. We list the possible other options and re-iterate why our final design stands out from other options.

1. Update the feature extractor first or run original OLS first. Through our analysis, we found only running original OLS first can still guarantee the theoretical results, while the other option cannot.
2. Re-training the linear layer under training distribution or other distribution. A more natural option is actually to retrain the linear layer with the most recent estimated label distribution. However, by checking with the original OLS methods, they require the model to be in the training distribution, or at least a distribution having all classes, but test distribution is not guaranteed to have this property since it can be just one class.
3. Batch accumulation or not. In fact, as shown in Table 1, even with batch size $\tau=1$ OLS-OFU has consistent improvement over OLS, and the necessity of batch accumulation can be missed. Through the experiment, as shown in Table 1, the benefit of batch accumulation brings both the performance gain and more time efficiency. Moreover, through theoretical and empirical analysis, we show that our algorithm provides a way to leverage self-supervised learning, which largely improves the performance as validated through the experiments and still keeps the similar theoretical guarantees.

Paper structure.

“Only three sentences describe OLS and SSL methods.”

We actually spent more energy to describe OLS and SSL methods. Given that our focus is on how our algorithm makes the bridge between the two, we included the sufficient details in the main paper and enumerated more details in the appendix.

• For OLS, we kindly refer the reviewer to these positions:
  1. At line 115 - 125 in Section 2, original OLS methods do not update feature extractor and this motivates the methodology of our paper.
  2. Line 175-181 describes the details of the condition where the theoretical guarantees would hold for original OLS methods, and this motivated our first design of our algorithm.
  3. Line 187-191 describes the choice of hypothesis space of original OLS methods, and the underlying reason for this choice in the literature motivates our second design of our algorithm.
  4. Moreover, for completeness, we included detailed algorithms of all 5 OLS methods as Algorithm 2-5 in the appendix.
• For SSL, the details of SSL are not our focus, as long as it has the form of $\ell_{\rm ssl}(S)$ for any batch of data $S$. Our algorithm OLS-OFU should work for general SSL methods, which is one of the information we would like to convey from our experiments. Therefore, we experiment with different SSLs as introduced at line 304-308 and their numerous details in the Appendix E.1.

“The problem setting of online label shift, which can be summarized in one sentence as "the test data have a different label distribution from training data", is explained with excessive and redundant content, including irrelevant details about online distribution shift.”

We kindly disagree with the description of online distribution shift and online label shift is redundant. Rather than just describe the mathematical problem, other components are necessary in introducing a problem too to make the paper have a broader audience, which include the motivating example (line 85-96), the description of objective function in the online setting (Line 97-103), the learning setting of unsupervised adaptation and limited unlabeled batch (Line 104-110), the assumption of label shift (Line 110-113), the summarization of existing online label shift methods and their limitations (Line 114-128) and introducing online generalized label shift adaptation as a first work (line 128 until line 129.)

“The experimental section is compressed into two pages, limiting the space to present results adequately”

We believe our experiment results are strong enough to demonstrate the effectiveness of our algorithms, which can be reflected by other reviews. We disagree that a two-pages experiment is a limitation, while we would appreciate any concrete advice to improve our result presentation.

Data privacy concern. This is actually an insightful point. We acknowledge this can bring additional privacy concerns, depending on the practical scenarios, and we would like to discuss this in the revision. At the same time, as shown in Table 1, we can find that even without batch accumulation (batch size $\tau$=1), OLS-OFU still outperforms OLS. Parameter $\tau$ can be treated as the performance and privacy trade-off.

Data split and baseline performance. There seems to be some misunderstanding about our data split and baselines, where we basically followed the literature. For the data split, as described at Line 293, our model was trained by 80% training set offline, with a 20% split as a validation set. We followed the code released from [1] for training the $f_0$ and our performance of “base” matches their numbers in Table 1, which are around 16% error for the CIFAR10.

Typos. Thanks for catching them! We will fix the typos. The definition of reweight is to reweight the soft-max probability by element-wise multiplying another vector.

[1] Baby, Dheeraj, et al. "Online label shift: Optimal dynamic regret meets practical algorithms." Advances in Neural Information Processing Systems 36 (2024).

We thank the reviewer for providing insightful comments! Here are our responses.

Discussion with concept drift. Yes, you are right that the (online) generalized label shift problem is a sub-field of concept drift and has its particular assumption. According to the definition of generalized label shift, it can be reduced to label shift once the learner knows the underlying feature mapping $h$. We would like to explain how each component of our method can solve each part specifically according to the literature and hence our method .

1. Online feature update learns the underlying feature mapping $h$ that makes $P(h(x_t)|y_t)$ invariant. In the literature of test-time training [1, 2, 3], they found the self-supervised learning of feature updates can explicitly align the feature space of new distribution to the feature space of training distribution; they even only trained with the feature extractor without retrain the last linear layer, this strongly supports the straightforward feature space alignment between original distribution and shifted distribution.
2. Built upon the updated feature extractor from 1, the problem can be reduced to online label shifts and we adopt the particular online label shift method, which follows the well-studied offline label shift method. Thank you for raising this point about generalized label shifts. We will add these explanations to Section 3.4 and we believe this can help improve the understanding of how our method can tackle the online generalized label shift.

[1] Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229–9248. PMLR, 2020.

[2] Wang, Dequan, et al. "Tent: Fully test-time adaptation by entropy minimization." arXiv preprint arXiv:2006.10726 (2020).

[3] Liu, Yuejiang, et al. "Ttt++: When does self-supervised test-time training fail or thrive?." Advances in Neural Information Processing Systems 34 (2021): 21808-21820.

We would like to thank the reviewer for these meaningful questions! Here are our responses.

Theoretical insights of SSL and the choice of gradient step size. We have to admit that the theoretical study for feature learning is generally hard; Instead, the analysis for the SSL in the literature is through exhaustive empirical validation. One hypothesis of the SSL methods in our paper has been validated in the literature: they can generally improve the feature representation rather than only work for specific tasks / distributions. Equivalently, for most data distributions, the accuracy of the classifier based on the “improved” feature representations could be higher. Moreover, this hypothesis supports our Equation 7, which shows the loss comparison between the classifiers where given the test distributions the classifiers are learned with the updated feature extractor or not; we also include the empirical validation for Equation 7 in the experiment section. As analyzed in Section 3.3, Equation 7 is a sufficient condition for FLHFTL-OFU having a better loss upper bound than FLHFTL. These arguments show how the hypothesis of the SSL methods, that has been validated in the literature, can indicate our OLS-OFU to be a better algorithm than OLS in the specific task of online label shift.

As for the choice gradient step size, we would have two empirical suggestions for practice. The first is to check the proper batch size for the SSL offline, for example large batch is necessary for the contrastive learning such as MoCo. The second is that as suggested in our experiment, $\tau=100$ is generally good for all three SSL methods, different datasets and different types of shifts. We recommend $\tau=100$ as a good starting point for choosing this parameter.

We will explicitly add the discussion of the SSL hypothesis and the gradient step size in our revision for more intuitions behind our algorithm.

The requirement to the training set. Thank the reviewer for raising this insightful question! We would like to first clarify that the step of retraining the last linear layer is the only place that needs training data in our algorithm. Moreover, as discussed in Principle 2 and line 248-252, retraining the last linear layer is designed only for three previous OLS methods (ROGD, FTH, FLHFTL); our algorithm for two other OLS methods (UOGD and ATLAS) in the literature are independent of this step. Therefore, including feature extractor updates by our algorithm for UOGD and ATLAS actually doesn’t require the training data.

As for our algorithm for ROGD, FTH, or FLHFTL, we further study how the amount of training data stored for the online test adaptation influences the effectiveness of our OLS-OFU. We evaluate OLS-OFU with $0\%-100\%$ stored training data; $0\%$ means that we still update the feature extractor but reuse the pretrained linear classifier. The results are reported in the following table.

We can observe that with less stored training data for retraining the last linear layer, the error of OLS-OFU would increase gradually. However, an important finding is that even with 0% stored training data, which means that we keep reusing the pretrained linear classifier together with updated feature extractor, the error of OLS-OFU is still lower than the OLS without feature extractor updates. This can actually explained by the original test-time training papers [1,2], where they only update the feature extractor without refining the last linear layer and the only feature updates still brings substantial benefit.

From the results, we can conclude that even if we remove the requirement of storing training data, our algorithm OLS-OFU can still outperform the OFU; more stored training data can further boost the performance. We will add this results in our revision!

[1] Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229–9248. PMLR, 2020.

[2] Wang, Dequan, et al. "Tent: Fully test-time adaptation by entropy minimization." arXiv preprint arXiv:2006.10726 (2020).

The influence of principle 1 in practice. We agree that empirical evidence can be important as well to illustrate the necessity for Principle 1. Therefore, we further make this ablation study to justify it: we compare OLS-OFU with its other variant named OLS-OFU-difforder where we update SSL first and run OLS later (which violates Principle 1). We compare these two algorithms across all previous 6 OLS methods and two choices of batch accumulation $\tau=1$ and $\tau=100$. The dataset is CIFAR-10, the SSL is rotation degree prediction and the shift pattern is sinusoidal shift; we observed similar performance for other settings of dataset, SSL and shift patterns and will include the full results in the reivision. The results are reported in the following table

We can observe that when $\tau=1$, the difference between OLS-OFU and OLS-OFU-difforder can be $0.9\%$ for OLS=FTL for example, which cannot be neglected. This means that Principle 1 indeed is important in practice when the batch is small.

As for the $\tau=100$, at first it seems there is no difference between OLS-OFU and OLS-OFU-difforder. This is actually because now we only update the feature extractor every $\tau=100$ time steps; this means that Principle 1 introduces no difference between OLS-OFU and OLS-OFU-difforder for the remaining 99 steps every 100 steps. We take a closer look at the average error of OLS-OFU and OLS-OFU-difforder for only the steps where we update the feature extractor and Principle 1 has been applied. Here are the numbers.

We can observe that OLS-OFU where we apply Principle 1 has non-neglecte improvements, which conclude that Principle 1 is important in practice even when the batch is large. To understand this, it is because when the feature extractor has dependency for the to-be-adapted test data, the later estimation for the distribution of these test data in the OLS method can be more inaccurate and hence this can hurt the performance.

Overall, Principle 1 is important not only just for the theoretical results, but also matters in practice. We will add this further analysis to our experiment section!

We hope we were successful in addressing your concerns. Please let us know if you have any additional concerns. We look forward to hearing from you!