Online Feature Updates Improve Online (Generalized) Label Shift Adaptation

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 09:33:36

Abstract: This paper addresses the prevalent issue of label shift in an online setting with missing labels, where data distributions change over time and obtaining timely labels is challenging. While existing methods primarily focus on adjusting or updating the final layer of a pre-trained classifier, we explore the untapped potential of enhancing feature representations using unlabeled data at test-time. Our novel method, Online Label Shift adaptation with Online Feature Updates (OLS-OFU), leverages self-supervised learning to refine the feature extraction process, thereby improving the prediction model. By carefully designing the algorithm, theoretically OLS-OFU maintains the similar online regret convergence to the results in the literature while taking the improved features into account. Empirically, it achieves substantial improvements over existing methods, which is as significant as the gains existing methods have over the baseline (i.e., without distribution shift adaptations).

Main paper: 
 introduction: 
 The effectiveness of most supervised learning models relies on a key assumption that the train data and test data share the same distribution. However, this assumption rarely holds in real-world scenarios, giving rise to distribution shift. Previous research has primarily focused on understanding distribution shifts in offline or batch settings, where a single shift occurs between the train and test distributions. In contrast, real-world applications often involve test data arriving in an online fashion, and the distribution shift can continuously evolve over time. Additionally, there is another challenging issue of missing and delayed feedback labels, stemming from the online setup, where gathering labels for the streaming data in a timely manner becomes a challenging task.To tackle the distribution shift problem, prior work makes further assumptions on the nature of the shift, such as label shift or covariate shift. In this paper, we focus on the common (generalized) label shift problem in an online fashion with missing labels. Specifically, the learner is given a fixed set of labeled training data D 0 ∼ P train in advance and trains a model f 0 . At test-time, only a small batch of unlabelled test data S t ∼ P test t arrives in an online fashion (t = 1, 2, • • • ). For online label shift, we assume the label distribution P test t (y) may change over time t while the conditional distribution stays the same, i.e. P test t (x|y) = P train (x|y). MRI image classifiers for concussion detection can be challenging due to label shift caused by seasonal changes in the image distribution. A classifier trained during skiing season may perform poorly when tested afterward, as the image distribution changes continuously between skiing and non-skiing seasons. The generalized label shift relaxes this assumption by assuming there exists a transformation h of the covariate, such that the conditional distribution P test t (h(x)|y) = P train (h(x)|y) stays the same. Reiterating our example, a classifier on MRI images can be used in different clinics during training and testing, where the MRI machines might have different versions at both hardware and software levels. Then, the images may have some variations such as brightness, resolution, etc. However, a feature extractor h exists, capable of mapping these variations to the same point in the latent space. The goal of the learner is to adapt to the (generalized) label shift within the non-stationary environment, continually adjusting the model's predictions in real time.Existing online label shift adaptation algorithms (OLS) primarily adopt one of two strategies: either directly reweighting of the pretrained classifier f 0 , or re-training only the final linear layer of f 0typically keeping the feature extractor frozen. Recent work has demonstrated that feature extractors can still be improved, even during test-time and in the absence of labels. We hypothesize that a f 0Labeled train data 8i, {(x i 0 , y i 0 ) ⇠ P trainUnlabeled test samples D 0 = {(x 1 0 , y 1 0 ), similar effect can be leveraged in the (generalized) label shift setting and propose to improve the feature representation during testing. In online label shift, updating the feature extractor offers two possible advantages. First, it enables the utilization of additional unlabeled samples to enhance the sample efficiency of the feature extractor. Second, it allows the adaptation of the feature extractor to label shift. Latter is important because the optimal feature extractor is not necessarily independent of the label distribution. In particular, in generalized label shift, the feature transformation h is typically unknown, and additional unlabeled test samples can ease the learning of h.Building upon this insight, this paper introduces the Online Label Shift adaptation with Online Feature Updates (OLS-OFU) framework, aimed at enhancing feature representation learning in the context of online label shift adaptation. Specifically, each instantiation of OLS-OFU incorporates a self-supervised learning method associated with a loss function denoted as l ssl and an existing online label shift adaptation (OLS) algorithm. In each time step, it executes a modified version of OLS, updates the feature extractor through self-supervised learning, and subsequently refines the last linear prediction layer. OLS-OFU is easy to implement and can be seamlessly integrated with any off-the-shelf OLS algorithms. Theoretically, we demonstrate that OLS-OFU reduces the loss of the overall algorithm by leveraging self-supervised learning techniques to enhance the feature extractor, thereby improving predictions for test samples at each time step t. Empirical evaluations on CIFAR-10 and CIFAR-10C datasets, considering both online label shift and online generalized label shift scenarios, affirm the effectiveness of OLS-OFU. It outperforms its counterpart OLS method across various distribution shift settings and consistently across all OLS methods, underscoring its robustness. Notably, the advantages of OLS-OFU become particularly pronounced in cases of domain shifts that happened in the generalized label shift, as demonstrated in the ablation study. 
problem setting & related work: 
 We start with some basic notations. Let ∆ K-1 be the probability simplex. Let f : X → ∆ K-1 denote a classifier. Given an input x from domain X , f (x) outputs a probabilistic prediction on K classes. For example, f can be the output from the softmax operation after any neural network. If we reweight a model f by a vector p ∈ R K , we refer to this model as g(•; f, p) with g denotes the method of reweighting. For any two vectors p and q, p/q denotes the element-wise division.Online distribution shift adaptation. The effectiveness of any machine learning model f relies on a common assumption that the train data D 0 and test data D test are sampled from the same distribution, i.e., P train = P test . However, this assumption is often violated in practice, which leads to distribution shift. This can be caused by various factors, such as data collection bias and changes in the data generation process. Moreover, once a well-trained model f 0 is deployed in the real world, it enters the test stage, which can be a sequence of time periods or time steps. The test distribution at time step t, P test t , from which test data x t is sampled, may vary over time. One example is that an MRI image classifier might be trained on MRI images during skiing season (which may have a high frequency of head concussions) but tested afterward (where the frequency of concussion is lower). The test stage can be several months long until the next classifier is trained. The MRI image distribution during the test stage may change continuously between non-skiing and skiing seasons.As the test-time distribution changes over time, the challenge lies in how to adjust the model continuously from f t-1 to f t in an online fashion to adapt to the current distribution P test t . We call this problem online distribution shift adaptation and illustrate it in Figure 1 . Given a total T steps in the online test stage, we define the average loss for any online algorithm A through the loss of the sequence of models f t , t ∈ [T ] that are produced from A, i.e.,EQUATIONwhere ℓ(f ; P) = E (x,y)∼P ℓ sup (f (x), y) and ℓ sup is the loss function, for example, 0-1 loss or cross-entropy loss for classification tasks.If we have knowledge of P test t or have enough samples from P test t , then the problem reduces to offline distribution shift for each single time step, or if we have a few labeled samples, the problem can be treated as classical online learning, and we can run stochastic gradient descent at the end of every time step t to meet certain theoretical guarantees. However, it is more realistic but technically highly non-trivial when we only have very few unlabeled samples, or even a single sample x t . An effective algorithm under this scenario has to utilize information from all historical data (both train data D 0 , validation data D ′ 0 and test data up to time t), as well as previously deployed models f 0,1,••• ,t-1 . In this paper, we consider this challenging case where at each time step t, only a small batch of unlabeled samplesS t = {x 1 t , • • • , x B t } is received.We formalize the algorithm A as:EQUATIONIn contrast to the classical online learning setup, this scenario presents a significant challenge as classical online learning literature usually relies on having access to either full or partial knowledge of the loss at each time step, i.e., ℓ sup (f (x t ), y t ). In this setting, however, only a batch of unlabeled samples is provided at each time step. This lack of both label information and loss values significantly amplifies the difficulty of accurately estimating the true loss in Equation 1.Related work have studied the online distribution shift adaptation under various assumptions. 2020), assumes that there exists an unknown transformation h of the feature x, such that the conditional distribution P(h(x)|y) remains invariant between the training and test time. Discussed in Section 5, our methodology applies beyond label shift scenarios, extending to distribution shift more broadly.Online label shift adaptation. In online label shift, the conditional distribution P test t (x|y) is invariant and equivalent to P train (x|y) for all t ∈ [T ], while the marginal distribution of the label P test t (y) changes over time. This assumption is most typical when the label y is the causal variable and and the feature x is the observation (Schölkopf et al., 2012) and the example of concussion detection from MRI images lies in the case. Most previous methods tackle this problem by a non-trivial reduction to the classical online learning problem. Given this, most of the online label shift algorithms (Wu et al., 2021; Bai et al., 2022; Baby et al., 2023) study the theoretical guarantee of the algorithm via the convergence of the regret function, either static regret or dynamic regret. In prior studies, the hypothesis class F of the prediction function f is typically chosen in one of two ways. The first approach involves defining F as a family of post-hoc reweightings of f 0 , with the parameter space comprising potential reweight vectors. Notable examples within this category include ROGD (Wu et al., 2021) , FTH (Wu et al., 2021) , and FLHFTL (Baby et al., 2023) . The second approach defines F as a family of functions that share the same parameters in f 0 except the last linear layer, for example, UOGD (Bai et al., 2022) and ATLAS (Bai et al., 2022) .Online generalized label shift adaptation. In the context of MRI image classification, where head MRI images serve as the feature x, MRI machines in different clinics may use different versions of software or hardware. The resulting images may exhibit differences in terms of brightness, contrast, resolution, and other characteristics. In such scenarios, the conditional probability distribution P(x|y) is no longer invariant. However, when a feature extractor h is robust enough, it can map the images into a feature space where the images from different machines have the same distributions P(h(x)|y) in this transformed feature space. The concept of generalized label shift, as introduced in Tachet des Combes et al. (2020) , formalizes this situation by postulating the existence of an unknown function h, such that the conditional probability distribution P(h(x)|y) remains invariant. The primary challenge in this context is to find the underlying transformation h. Building upon this, online generalized label shift assumes that, for every time step t, the test distribution P test t exhibits a form of generalized label shift from the training distribution P train , and this shift is governed by the same underlying transformation h.Self-supervised learning. Inspired by the body of work in semi-supervised learning (Grandvalet & Bengio, 2004b; Lee et al., 2013; Laine & Aila, 2016; Gidaris et al., 2018; Miyato et al., 2018) and unsupervised representation learning (Chen et al., 2020a; He et al., 2019; Grill et al., 2020; He et al., 2022) , self-supervised learning (SSL) techniques emerge as promising tools for enhancing feature extraction from unlabeled data, e.g. for image classification. Likewise, when dealing with online label shift adaptation, it is crucial to leverage the unlabeled test samples S 1 ∪ • • • ∪ S t-1 obtained from previous time steps. Ideally, these unlabeled samples could improve the feature representation learning and ultimately lead to better prediction for the test samples S t at a time step t.qt+1 := ALG(s 1 , • • • , s t ) 3. Let f ′ t+1 be a reweighting version of f ′′ t by the weight qt+1[k] q0[k] : k = 1, • • • K ▷ In the original FLHFTL, it is f0 rather than f ′′ t .Output at time t: f ′ t+1 . end for• FLHFTL and FTH subsequently employ a vector qt3 to reweight the initial pretrained model f 0 . Now, instead of reweighting the original pretrained model f 0 , the algorithm reweights the improved model f ′′ t , utilizing its enhanced predictive performance. • ROGD, UOGD, and ATLAS initially update the model through an unbiased gradient estimator within a hypothesis space that either includes a linear model after a fixed feature extractor in f 0 or a reweighted version of f 0 . Now, we have the flexibility to shift to a comparable hypothesis space, replacing f 0 with the improved model f ′′ t , and continue applying updates using the same unbiased gradient estimator.Algorithm 2 provides a formal illustration of the revision specific to FLHFTL. Further details regarding the revisions for ROGD, FTH, UOGD, and ATLAS can be found in Appendix B. We use f ′ t+1 to denote the model after running the revised OLS.(2) Updating the Feature Extractor. We now introduce how to utilize an SSL loss ℓ ssl to update the feature extractor for any incoming unlabeled test data batch S t at timestep t. Specifically, let θ feat t denote the parameters of the feature extractor in f ′ t+1 . The update of θ feat t+1 at time t is given by:θ feat t+1 := θ feat t -η • ∇ θ feat ℓ ssl (S t ; f ′ t+1 ).We replace the feature extractor in f ′ t+1 by θ feat t+1 .(3) Re-training Last Linear Layer. Given the updated feature extractor θ feat t+1 , it is necessary to re-train the last linear layer θ linear t+1 to adapt to the new feature extractor. We start the re-raining from random initialization, while keeping the feature extractor frozen. The train objective of θ linear t+1 is to minimize the average cross-entropy loss under train data D 0 . We denote the model with the frozen feature extractor θ feat t+1 as f (•|θ feat t+1 , θ linear ). The objective for re-training last linear layer can be written as follows:θ linear t+1 := arg min θ linear (x,y)∈D0 ℓ ce f (x|θ feat t+1 , θ linear ), y .We calibrate the model f (•|θ feat t+1 , θ linear t+1 ) by temperature calibration (Guo et al., 2017) using the validation set D ′ 0 and denote the model after calibration as f ′′ t+1 . This re-training step is needed to ensure the model is a calibrated model of estimating P train (y|x) for any given input x.In the end, we are going to define f t+1 for the next time step. If the parameter space of the original OLS is a reweighting version of the prediction model (ROGD, FTH, FLHFTL), suppose the reweighting vector in f ′ t+1 is p t+1 and we define f t+1 := g(•; f ′′ t+1 , p t+1 ); else (UOGD, ATLAS), we define f t+1 := f ′ t+1 .t+1 = 1 t t τ =1 s τ 3. Let f t+1 be a reweighting version of f ′′ t by the weight pt+1[k] q0[k] : k = 1, • • • K ▷ In the original FTL, it is f0 rather than f ′′ t .Output at time t: f ′ t+1 . end for Algorithm 5 Revised UOGD for online feature updates (UOGD-R). See the original version in Equation 9in Bai et al. (2022) .Require: The learning rate η.for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, and intermediate model {f ′′ 1 , • • • , f ′′ t } from step 3 in Algorithm 1, the validation set D ′ 0 , the train label marginal q 0 := P train (y). 1. Compute the unbiased estimator for label marginal distribution:s t = 1 |St| xt∈St C -1 f ′′ t ,D ′ 0 f ′′ t (x t ).▷ In the original UOGD, it is f0 rather than f ′′ t .2. Grab the weight w t from the last linear layer of f t .Ĝt (w) ∈ R K such that E[ Ĝt (w)|S 1:t-1 ] = E ((x,y)∼P test t ) ∇ w ℓ ce f (x|θ feat f ′′ t , w), y .Hence we proceed to verify these two facts in our setup. Fact 1 is true because the cross-entropy loss is convex in any subset of the simplex and the last linear layer weights only defines an affine transformation which preserves convexity.For fact 2, note that the f ′′ t only uses the data until round t -1. So by the same arguments in Bai et al. (2022) , using the BBSE estimator defined from the classifier f ′′ t , the unbiased estimate of risk gradient can be defined. Let w t be the weight of the last layer maintained by UOGD at round t. Let u 1:T be any sequence in W. Consequently we have for any round,EQUATIONRest of the proof is identical to Bai et al. (2022) . = {η i = O σ √ KT • 2 i-1 |i ∈ [N ]} (N = 1 + ⌈ 1 2 log 2 (1 + 2T )⌉), then under the same assumptions as Lemma 1 in Bai et al. (2022) , UOGD-OFU satisfies thatEQUATION) where V T,ℓ := T t=2 sup w∈W |ℓ(f (•; θ feat f ′′ t , w); P test t ) -ℓ(f (•; θ feat f ′′ t-1 , w); P test t-1 )|, σ denotes the min- imum singular value of the invertible confusion matrices {C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′0 } and K is the number of classes and the expectation is taken with respect to randomness in the revealed co-variates.The proof is similar to that of Theorem 5 and hence omitted. 
method: 
 We introduce a novel online label shift adaptation algorithm that uses self-supervised learning (SSL) to improve representation learning. This approach is general and can be used with any existing OLS algorithm. We show theoretically and empirically the performance improvement in OLS algorithms and their effectiveness for the more challenging online generalized label shift adaptation problem.In this section, we discuss how to utilize various SSL techniques to improve the feature representation learning of existing online label shift algorithms. To illustrate the concept, we narrow our focus to three particular SSL techniques for classification tasks: rotation degree prediction (Gidaris et al., 2018; Sun et al., 2020) , entropy minimization (Grandvalet & Bengio, 2004b; Wang et al., 2020) and MoCo (He et al., 2019; Chen et al., 2020b; 2021) . It is important to note that this concept extends beyond these three SSL techniques, and the incorporation of more advanced SSL techniques has the potential to further elevate the performance. Specifically, rotation degree prediction involves initially rotating a given image by a specific degree from the set {0, 90, 180, 270} and the classifier is required to determine the degree by which the image has been rotated. Entropy minimization utilizes a minimum entropy regularizer, with the motivation that unlabeled examples are mostly beneficial when classes have a small overlap. MoCo is a more advanced representation learning technique, using a query and momentum encoder to learn representations from unlabeled data by maximizing the similarity of positive pairs and minimizing the similarity of negative pairs. The details of these SSL are introduced in Appendix D.8. Now we formally introduce Online Label Shift adaptation with Online Feature Updates (OLS-OFU; Algorithm 1), which requires a self-supervised learning loss ℓ ssl and an online label shift adaptation algorithm (OLS) that either reweights the offline pretrained model f 0 or updates the last linear layer1 . In the train stage, we train f 0 by minimizing the supervised and self-supervised loss together defined on train data. In the test stage, OLS-OFU comprises three steps at each time step t: (1) running the refined version of OLS, which we refer to as OLS-R, (2) updating the feature extractor, (3) re-training the last linear layer. The details of these three steps are elaborated below.(1) Running the Revised OLS. First, we review common OLS methods (FLHFTL, FTH, ROGD, UOGD, and ATLAS) and identify two specific points in the algorithm where we can employ the updated prediction model f ′′ t (with the improved feature extractor 2 ) to supplant the pretrained model f 0 , hence enhancing the original OLS algorithm. Denote C f,D ∈ [0, 1] K×K the confusion matrix evaluated on dataset D for the model f with C f,D [i, j] = P (x,y)∼D arg max(f (x)) = j|y = i . At the outset, all OLS methods rely on an unbiased estimator s t of the label distribution q t with q t [y] = P test t (y), where st := 1 |St| xt∈St C -1 f0,D0 f 0 (x t ). This is the first point that we can replace f 0 with the improved prediction model f ′′ t to enhance the estimation of label marginal distribution. For the second point, Algorithm 1 Online label shift adaptation with online feature updates (OLS-OFU).Require: An online label shift adaptation algorithm OLS (∈{ROGD, FTH, UOGD, ATLAS, FLHFTL}), a self-supervised learning loss ℓ ssl . A pretrained model f 0 and initializef 1 = f 0 . for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, train set D 0 and validation set D ′ 0 . 1.Run the revised version of OLS, that is, OLS-R, and get f ′ t+1 . (Algorithm 2 is the revised FLHFTL; See Appendix B for revisions of ROGD, FTH, UOGD, ATLAS) 2. Update the feature extractor θ feat t in f ′ t+1 byθ feat t+1 := θ feat t -η • ∇ θ feat ℓ ssl (S t ; f ′ t+1 ).Replace the feature extractor θ feat t in f ′ t+1 by θ feat t+1 . 3. Within the feature extractor θ feat t+1 , re-train the last linear layer from random initialization by minimizing the average loss among train data D 0 :θ linear t+1 := arg min θ linear (x,y)∈D0 ℓ ce f (x|θ feat t+1 , θ linear ), y .Calibrate the model f (•|θ feat t+1 , θ linear t+1 ) by temperature calibration using the validation set D ′ 0 and denote the model after calibration as f ′′ t+1 .If the parameter space of OLS is a reweighting version of the prediction model (ROGD, FTH, FLHFTL), suppose the reweighting vector in f ′ t+1 is p t+1 and we define f t+1 := g(•; f ′′ t+1 , p t+1 ); else (UOGD, ATLAS), we define f t+1 := f ′ t+1 . end for Algorithm 2 Revised FLHFTL for online feature updates (FLHFTL-R); See the original version in Algorithm 2 in Baby et al. (2023) .Require: Online regression oracle ALG.for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, and intermediate model {f ′′ 1 , • • • , f ′′ t } from step 3 in Algorithm 1, the validation set D ′ 0 , the train label marginal q 0 := P train (y). 1. Compute the unbiased estimator for label marginal distribution:s t = 1 |St| xt∈St C -1 f ′′ t ,D ′ 0 f ′′ t (x t )▷ In the original FLHFTL, it is f0 rather than f ′′ t .The original Online Label Shift (OLS) methods exhibit theoretical guarantees in terms of regret convergence for online label shift setting, where P test t (x|y) = P train (x|y) is invariant. With the incorporation of the additional online feature update step, OLS-OFU demonstrates analogous theoretical results. By comparing the theoretical results between OLS and OLS-OFU, we can gain insights into potential enhancements from OLS to OLS-OFU. Due to limited space, we provide the theoretical results pertaining to FLHFTL-OFU here, and present the theoretical results for ROGD-OFU, FTH-OFU, UOGD-OFU, and ATLAS-OFU in Appendix C.Theorem 1 [Regret convergence for FLHFTL-OFU] Suppose we choose the OLS subroutine in Algorithm 2 to be FLH-FTL from Baby et al. (2023) . Let f flhftl-ofu t be the output at time step t -1 from Algorithm 1, that is g(•; f ′′ t , qt /q 0 ). Let σ be the smallest among the the minimum singular values of invertible confusion matrices{C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′ 0 }.Then under Assumptions 1 and 2 in Baby et al. (2023) , FLHFTL-OFU has the guarantee for online label shift below:E 1 T T t=1 ℓ(f flhftl-ofu t ; P test t ) - 1 T T t=1 ℓ(g(•; f ′′ t , q t /q 0 ); P test t ) ≤ O K 1/6 V 1/3 T σ 2/3 T 1/3 + K σ √ T ,(3) where V T := T t=1 ∥q tq t-1 ∥ 1 , K is the number of classes, and the expectation is taken w.r.t. randomness in the revealed co-variates. This result is attained without prior knowledge of V T .To ease comparison, we state the theorem for the original OLS algorithm FLHFTL.Theorem 2 [Regret convergence for FLHFTL (Baby et al., 2023) ] Under Assumptions 1 and 2 in Baby et al. (2023) , FLHFTL has the guarantee for online label shift below:EQUATIONwhere V T := T t=1 ∥q tq t-1 ∥ 1 , σ denotes the minimum singular value of invertible confusion matrices C f0,D ′ 0 , K is the number of classes, and the expectation is taken with respect to randomness in the revealed co-variates. Further, this result is attained without prior knowledge of V T .Recall that the objective function for the online label shift problem is defined as the average loss in Equation 1. Both theorems establish the convergence of this average loss. In the event that f ′′ t (t ∈ [T ]) from the online feature updates yield improvements:E 1 T T t=1 ℓ(g(•; f ′′ t , q t /q 0 ); P test t ) < 1 T T t=1ℓ(g(•; f 0 , q t /q 0 );EQUATIONthen it guarantees that the loss of FLHFTL-OFU will converge to a smaller value, resulting in enhanced performance compared to FLHFTL. We substantiate this improvement through empirical evaluation in Section 4. For other OLS algorithms such as ROGD-OFU, FTH-OFU, UOGD-OFU, and ATLAS-OFU, a similar analysis can be derived and we present them in Appendix C.When we have knowledge of the feature map h under generalized label shift, the problem simplifies to the classical online label shift scenario. However, the more challenging situation arises when h remains unknown. Due to the violation of the label shift assumption, standard OLS algorithms can perform arbitrarily bad. Fortunately, existing research in test-time training (TTT) (Sun et al., 2020; Wang et al., 2020; Liu et al., 2021; Niu et al., 2022) demonstrates that feature updates driven by SSL align the source and target domains in feature space. When the source and target domains achieve perfect alignment, such feature extractor effectively serves as the feature map h as assumed in generalized label shift. Therefore, the sequence of feature extractors in f 1 , • • • , f T generated by Algorithm 1 progressively approximates the underlying h. This suggests that, compared to the original OLS, OLS with online feature updates (Algorithm 1) experiences a milder violation of the label shift assumption within the feature space and is actually expected to have better performance in the setting of online generalized label shift. Indeed, as demonstrated in Section 4, we can observe significant improvements in OLS-OFU compared to standard OLS for online generalized label shift.t+1 := Proj ∆ K-1 p t -η • J p (p t ) ⊤ s t , where J p,f ′′ t (p t ) = ∂ ∂p (1 -diag(C f ′′ t ,D0,p ))| p=pt , and let f t+1 be a reweighting version of f ′′ t by the weight pt+1[k] q0[k] : k = 1, • • • K ▷ In the original ROGD, it is f0 rather than f ′′ t .Output at time t: f ′ t+1 . end for Algorithm 4 Revised FTH for online feature updates (FTH-R). See the original version in Equation 9in Wu et al. (2021) .for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, and intermediate model {f ′′ 1 , • • • , f ′′ t } from step 3 in Algorithm 1, the validation set D ′ 0 , the train label marginal q 0 := P train (y). 1. Compute the unbiased estimator for label marginal distribution:s t = 1 |St| xt∈St C -1 ft,D ′ 0 f ′′ t (x t ).▷ In the original FTL, it is f0 rather than f ′′ t .t+1 := w t -η • ∂ ∂w J w (w t ) ⊤ s t , where J w (w t ) = ∂ ∂w ( R1 t (w), • • • , RK t (w))| w=wt , Rk t (w) = 1 |D k 0 | (x,y)∈D k 0 ℓ ce (f (x|θ feat t , θ linear = w), y), D k 0 denotes the set of data with label k in D 0 . ▷ In the original UOGD, it is θ feat 0 rather than θ feat t . 4. Let f t+1 be f (•|θ feat t , w t+1 ).Output at time t: f ′ t+1 . end for Algorithm 6 Revised ATLAS for online feature updates (ATLAS-R). See the original version in Equation 9 in Bai et al. (2022) .Require: The learning rate pool H with size N; Meta learning rate ε; ∀i ∈[N ], p 1,i = 1/N and w 1,i = θ linear f0 . for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, and intermediate model {f ′′ 1 , • • • , f ′′ t } from step 3 in Algorithm 1, the validation set D ′ 0 , the train label marginal q 0 := P train (y). 1. Compute the unbiased estimator for label marginal distribution:s t = 1 |St| xt∈St C -1 ft,D ′ 0 f ′′ t (x t ).▷ In the original ATLAS, it is f0 rather than f ′′ t .fori ∈ [N ] do 2. Update w t+1,i := w t,i -η i • ∂ ∂w J w (w t,i ) ⊤ s t , where J w (w t,i ) = ∂ ∂w ( R1 t (w), • • • , RK t (w))| w=wt,i , Rk t (w) = 1 |D k 0 | (x,y)∈D k 0 ℓ ce (f (x|θ feat t , w), y), D k 0 denotes the set of data with label k in D 0 . ▷ In the original ATLAS, it is θ feat 0 rather than θ feat t .end for 3. Update weight p t+1 according to p pt,i ∝ exp(-εt-1 τ =1 Rτ (w τ,i )) 3. Compute w t+1 = N i=1 p t+1,i w t+1,i . Let f t+1 be f (•|θ feat t , w t+1 ). Output at time t: f ′ t+1 . end for 
experiment: 
 In this section, we empirically evaluate how OLS-OFU improves the original OLS methods on both online label shift and online generalized label shift4 . The experiment is performed on CIFAR-10 ( Krizhevsky et al., 2009) and CIFAR-10C (Hendrycks & Dietterich, 2019) , where we vary the shift processes and SSL techniques to evaluate the efficacy of the method.Dataset and shift process set-up. For online label shift, we evaluate the efficacy of our algorithm on CIFAR-10, which has 10 categories of images. We split the original train set of CIFAR-10 into the offline train and validation sets, which have 40,000 and 10,000 images respectively. At the online test stage, the unlabeled batches are sampled from the test set of CIFAR-10. For online generalized label shift, the offline train and validation sets are the same CIFAR-10 images, but the test unlabeled batches are sampled from CIFAR-10C. CIFAR-10C is the benchmark that has the same objects in CIFAR-10 but with multiple types of corruption. We experiment with three types of corruptions (i.e., domain shifts): Gaussian noise, Fog and Pixelate. Besides CIFAR-10 and CIFAR-10C, we also experiment with three additional datasets for the setting of online label shift and present the results in Appendix D.2 due to the space limits. See more details of dataset set-up in Appendix D.1.We follow Bai et al. (2022) and Baby et al. (2023) to simulate the online label distribution shift in two online shift patterns: Sinusoidal shift and Bernoulli shift. Given two label distribution vectors q and q ′ , the label marginal distributions at time t is q t := α t q + (1α t )q ′ . In Sinusoidal shift, α t = sin iπ L (periodic length L = √ T , i = t mod L) while in Bernoulli shift, α t is a random bit (either 0 or 1), where the bit switches α t = α t -1 if the coinflip probability exceeds p = 1 √ T . The q and q ′ are 1 K (1, • • • , 1) and (1, 0, • • • , 0) in the experiment. To sample the batch test data at time t, we first sample a batch of labels (not revealed to the learner) according to q t . Then given each label we can sample an image from the test set, and collect this batch of images without labels as S t . We experiment with T = 1000 and batch size B = 10 at each time step, following Baby et al. (2023) .Evaluation metric. We report the average error 1 ) for the evaluation efficiency. This approximation is valid for large T due to its exponential concentration rate by the Azuma-Hoeffding inequality.T B T t=1 xt∈St 1 (f t (x t ) ̸ = y t ),Online algorithms set-up. We perform an extensive evaluation of 6 OLS algorithms in the literature: FTFWH, FTH, ROGD, UOGD, ATLAS, and FLH-FTL. FTFWH is an empirical OLS proposed in Wu et al. (2021) . We further report the performance of OLS-OFU (Algorithm 1) on top of each OLS. OLS-OFU is implemented with 3 common SSL methods: rotation degree prediction, entropy minimization, and MoCo. Additionally, we report two baseline scores. The first, denoted as Base, uses the fixed pretrained model f 0 to predict the labels at all test time steps. The second is online feature updates (OFU) only, where at time step t we only update the features (Step 2 in Algorithm 1) without utilizing OLS algorithms.How does OLS-OFU perform compared with original OLS methods? Figure 2 (a) compares the performance of various OLS-OFU algorithms with their respective OLS counterparts, under classical online label shift. The experiment is performed on CIFAR10 with Sinusoidal shift. Appendix D.3 shows similar conclusions for the Bernoulli shift. First of all, both OLS and OLS-OFU show significantly better performance than Base and OFU. This demonstrates the inherent advantages of OLS methods in effectively addressing the online label shift problem. Additionally, it's worth noting that OLS-OFU, when integrated with three distinct SSL methods, consistently outperforms OLS across all six OLS methods. This underscores the importance of improving feature representation learning within label shift adaptation in OLS-OFU. Furthermore, we report the performance of OLS-OFU in the context of online generalized label shift. In this scenario, the test images exhibit three types of domain shifts in CIFAR-10C -Gaussian noise, Fog, and Pixelation-with mild severity. We present the result under Sinusoidal shift with the SSL method as rotation degree prediction. Results for other SSL methods and online shift patterns show a similar pattern in Appendix D.4. As shown in Figure 2 (b), OFU has huge improvements from Base, showing the necessity of updating feature extractors to facilitate the learning of h. Moreover, some OLS methods perform worse than OFU. This is expected as the underlying assumption of label shift no longer holds in this generalized label shift setting. Lastly, it is worth highlighting that OLS-OFU consistently outperforms the original OLS methods by a larger gap than the one occurring at the classic label shift setting. We hypothesize that the SSL methods within OLS-OFU aid the feature extractor in learning the unknown mapping h within the generalized label shift assumption. Intuitively, the label shift assumption in the feature space of θ feat t is violated lighter than the one in the feature space of θ feat 0 . Given the relatively mild violation, the OLS module within OLS-OFU continues to help adapt the online label shift, explaining OLS-OFU's advantage over OFU. How does OLS-OFU perform under high severity of domain shift in the setting of generalized online label shift? In Figure 3 , it is clear that as the domain shift severity increases, OLS-OFU significantly enlarges the gap compared with OLS. However, it might be worth pointing out that neither OLS nor OLS-OFU are better than OFU. This is because the label shift assumption is violated so severely -even OFU cannot reduce the violation under an acceptable level, especially when the OLS module exists in OLS-OFU, as the adaptation to distribution shift is far off. For results involving higher levels of severity, different corruption types, and SSL techniques, please refer to Appendix D.5. 5hold empirically? In Section 3.2, we argued that when the inequality in Equation 5holds, the loss of FLHFTL-OFU exhibits a tighter upper bound compared to FL-HFTL. Figure 4 presents the RHS (corresponds to OLS) and LHS (corresponds to OLS-OFU with SSL loss as rotation degree prediction) of Equation 5. We perform the study over cross eight different settings, varying types of domain shift and online shift pattern, which empirically validates that OLS-OFU yields improvements on the baseline of the regret as shown in Equation 5. Appendix D.6 validates this inequality for other SSL techniques. Does the order of prediction and update matter? In the default online distribution shift framework (Figure 1 ), model updates occur after making predictions for samples at timestep t. This raises the question of whether the model should be updated before making predictions. We conducted empirical evaluations for both the "predict first" and "update first" approaches and found no compelling evidence to favor one over the other (additional results in Appendix D.7). However, it's noteworthy that within the "predict first" framework, OLS and OLS-OFU benefit from robust theoretical guarantees, hence we recommend the "predict first" approach in practical applications. 
conclusion: 
 We focus on online (generalized) label shift adaptation and introduce a novel algorithm Online Label Shift adaptation with Online Feature Updates (OLS-OFU), which harnesses the power of self-supervised learning to enhance feature representations dynamically during testing, leading to improved predictive models and better test time performance. Our theoretical analyses show that OLS-OFU successfully achieves the improved theoretical guarantees under online label shift. We also validate the performance of OLS-OFU on both online label shift and generalized label shift scenarios, demonstrating OLS-OFU's superiority over prior online label shift algorithms. This underscores its efficacy and robustness, particularly when confronted with domain shifts.Discussion and future work. In this paper, we have shown how online feature updates help with online (generalized) label shift adaptation. One promising direction is to extend this idea to online covariate shift -the algorithm in Zhang et al. (2023) freezes the feature extractor and only updates the linear layer. Another possible extension is to consider a more realistic domain shift within the generalized label shift setting -domain shift types may vary over time or they can be more challenging, such as shifting from cartoon images to realistic images. More advanced SSL techniques in the deep learning literature might be needed to handle these more intricate domain shifts. 
 
 Appendix: Offline distribution shift and domain shift. Offline label shift and covariate shift have been studied for many years. Some early work (Saerens et al., 2002; Lin et al., 2002) assumes the knowledge of how the distribution is shifted. Later work (Shimodaira, 2000; Zadrozny, 2004; Huang et al., 2006; Gretton et al., 2009; Lipton et al., 2018; Alexandari et al., 2020; Azizzadenesheli et al., 2019; Garg et al., 2020) relaxes this assumption and estimates this knowledge from unlabeled test data.Online distribution shift with provable guarantees. There has been several work modeling online distribution shift as the classic online learning problem (Wu et al., 2021; Bai et al., 2022; Baby et al., 2023; Zhang et al., 2023) , which leverage the classical online learning algorithms (Shalev-Shwartz, 2012; Besbes et al., 2015; Baby & Wang, 2022) to bound the static or dynamic regret. However, none of them updates the feature extractor in a deep learning model but only the last linear layer or the post-hoc linear reweighting vectors. Our proposed method OLS-OFU utilizes the deep learning SSL to improve the feature extractor, which brings better performance.Domain shift adaptation within online streaming data. When we consider the most authentic online learning setup where the learner only receives the unlabeled samples, the most representative idea is test-time training (Sun et al., 2020; Wang et al., 2020; Liu et al., 2021; Niu et al., 2022) , which utilizes a (deep learning) self-supervised loss to online update the model. However, it focuses on how to adapt to a fixed domain shifted distribution from online streaming data and is not designed for how to adapt to continuous distribution changes during the test stage, while our algorithm concentrates the later problem. Besides test-time training, Hoffman et al. (2014) and Mullapudi et al. (2019) study the online domain shift for specific visual applications.Algorithm 3 Revised ROGD for online feature updates ROGD-R. See the original version in Equation 7 and Equation 8in Wu et al. (2021) .Require: Learning rate η.for t = 1, • • • , T do Input at time t: Samples S 1 ∪ • • • ∪ S t , models {f 1 , • • • , f t }, and intermediate model {f ′′ 1 , • • • , f ′′ t } from step 3 in Algorithm 1, the validation set D ′ 0 , the training label marginal q 0 := P train (y). 1. Compute the unbiased estimator for label marginal distribution:s t = 1 |St| xt∈St C -1 f ′′ t ,D ′ 0 f ′′ t (x t ).▷ In the original ROGD, it is f0 rather than f ′′ t .2. Grab the weight p t from f t .The revised algorithms to be used in the main algorithm OLS-OFU (Algorithm 1) are FTH-R (Algorithm 4), UOGD-R (Algorithm 5), ROGD-R (Algorithm 3), ATLAS-R (Algorithm 6).In this section, we present the theoretical results of FLHFTL-OFU, ROGD-OFU, FTH-OFU, UOGD-OFU, ATLAS-OFU and their proofs. The proofs are mostly the same as the proofs for the original algorithms with small adjustments. As our results are not straight corollaries for the original theorems, we write the full proofs here for the completeness.Before proving Theorem 1 (in Section 3.2 ) we recall the assumption from Baby et al. (2023) for convenience. We refer the reader to Baby et al. (2023) for justifications and further details of the assumptions.Assumption 1 Assume access to the true label marginals q 0 ∈ ∆ K of the offline train data and the true confusion matrix C ∈ R K×K . Further the minimum singular value σ min (C) = Ω(1) is bounded away from zero.Assumption 2 (Lipschitzness of loss functions) Let D be a compact and convex domain. Let r t be any probabilistic classifier. Assume thatL t (p) := E [ℓ(g(•; r t , p/q 0 )|x t ] is G Lipschitz with p ∈ D ⊆ ∆ K , i.e, L t (p 1 ) -L t (p 2 ) ≤ G∥p 1 -p 2 ∥ 2 for any p 1 , p 2 ∈ D.The constant G need not be known ahead of time.Theorem 1 [Regret convergence for FLHFTL-OFU] Suppose we choose the OLS subroutine in Algorithm 2 to be FLH-FTL from Baby et al. (2023) . Let f flhftl-ofu t be the output at time step t -1 from Algorithm 1, that is g(•; f ′′ t , qt /q 0 ). Let σ be the smallest among the the minimum singular values of invertible confusion matrices{C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′ 0 }.Then under Assumptions 1 and 2 in Baby et al. (2023) , FLHFTL-OFU has the guarantee for online label shift below:EQUATIONwhere V T := T t=1 ∥q tq t-1 ∥ 1 , K is the number of classes, and the expectation is taken w.r.t. randomness in the revealed co-variates. This result is attained without prior knowledge of V T .The algorithm in Baby et al. (2023) requires that the estimate s t in Line 1 of Algorithm 2 is unbiased estimate of the label marginal q t . Since f ′′ t in Algorithm 2 is independent of the sample S t , and since we are working under the standard label shift assumption, due to Lipton et al. (2018) we have that Lipton et al. (2018) , the reciprocal of standard deviation of this estimate is bounded below by minimum of the singular values of confusion matricesC -1 f ′′ t ,D ′ 0 • 1 |St| xt∈St f ′′ t (x t ) forms an unbiased estimate of E x∼P test t [f ′′ t (x)]. Further, from{C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′ 0 }.Let qt be the estimate of the label marginal maintained by FLHFTL. By Lipschitzness, we have thatEQUATIONwhere the last line is via Assumption 2. Rest of the proof is identical to that of Baby et al. (2023) . We reproduce it below for completeness.EQUATIONwhere the second line is due to Jensen's inequality, third line by Cauchy-Schwartz and last line by Proposition 16 in Baby et al. (2023) . This finishes the proof.Proof of Theorem 2: The proof is similar to the arguments in the proof of Theorem 1. The only point of deviation is that we choose r t = f 0 instead of f ′′ t in Assumption 2. The rest of the arguments follow via Lipschitzness.We state the assumptions first for the later theorems. These assumptions are similar to Assumption 1-3 in Wu et al. (2021) . EQUATIONEQUATIONProof: For any fixed p, ℓ(f ) = (p t -p) • J p,f ′′ t (p t ) ⊤ E St [s t |S 1 , • • • , S t-1 ] = E St [(p t -p) • J p,f ′′ t (p t ) ⊤ s t |S 1 , • • • , S t-1 ],where the last inequality holds by the fact that(p t -p)•J p,f ′′ t (p t ) ⊤ is independent of {S 1 , • • • , S t-1 }. To bound (p t -p) • J p,f ′′ t (p t ) ⊤ s t , ∥p t+1 -p∥ 2 2 = ∥Prof ∆ K-1 (p t -η • J p,f ′′ t (p t ) ⊤ s t ) -p∥ 2 2 ≤ ∥p t -η • J p,f ′′ t (p t ) ⊤ s t -p∥ 2 2 = ∥p t -p∥ 2 2 + η 2 ∥J p,f ′′ t (p t ) ⊤ s t ∥ 2 2 -2η(p t -p) • (J p,f ′′ t (p t ) ⊤ s t ). This implies (p t -p) • (J p,f ′′ t (p t ) ⊤ s t ) ≤ 1 2η (∥p t -p∥ 2 2 -∥p t+1 -p∥ 2 2 ) + η 2 ∥J p,f ′′ t (p t ) ⊤ s t ∥ 2 2 Thus E S1,••• ,S T 1 T T t=1 ℓ(f rogd-ofu t ; P test t ) - 1 T T t=1 ℓ(g(•; f ′′ t , p/q 0 ); P test t ) ≤ E S1,••• ,S T 1 T T t=1 1 2η (∥p t -p∥ 2 2 -∥p t+1 -p∥ 2 2 ) + η 2 ∥J p,f ′′ t (p t ) ⊤ s t ∥ 2 2 ≤ 1 2ηT ∥p 1 -p∥ 2 2 + η 2T T t=1 E S1,••• ,St [∥J p,f ′′ t (p t ) ⊤ s t ∥ 2 2 ] ≤ 1 ηT + ηL 2 2 = 2 T L.This bound holds for any p. Thus,E S1,••• ,S T 1 T T t=1 ℓ(f rogd-ofu t ; P test t ) -min p∈∆ K-1 E S1,••• ,S T 1 T T t=1 ℓ(g(•; f ′′ t , p/q 0 ); P test t ) ≤ 2 T L.We begin with two assumptions. {C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′0 }, under Assumption 6 and 7 with δ = 0, FTH-OFU satisfies the guarantee that with probability at least 1 -2KT -7 over samplesEQUATIONwhere K is the number of classes.Proof: Denote q t := (P test t (y = k) : k ∈ [K]). By the Hoeffding and union bound, we haveP ∀t ≤ T, ∥p t+1 - 1 t t τ =1 q τ ∥ ≤ √ Kε t ≥ 1 - T t=1 2M exp -2ε 2 t t/σ 2 .This implies that with probability at least 1 - q τ /q 0 );P test t ) - T t=1 ℓ(g(•; f ′′ t , p/q 0 ); P test t ) + L √ M • T t=1 ε t ≤ T t=1 ℓ(g(•; f ′′ t , 1 t -1 t-1 τ =1 q τ /q 0 ); P test t ) - T t=1 ℓ(g(•; f ′′ t , 1 t t τ =1q τ /q 0 );P test t ) + L √ M • T t=1 ε t ≤ T t=1 L 1 t -1 t-1 τ =1 q τ - 1 t t τ =1 q τ + L √ M • T t=1 ε t ≤ T t=1 L t 1 t -1 t-1 τ =1 q τ -q t + L √ M • T t=1 ε t ≤ T t=1 2L t + L √ M • T t=1 ε t .If we take ε t = 2σ ln T T , the above is equivalent to: with probability at least 1 -2KT Bai et al. (2022) and let step size be η, then under the same assumptions as Lemma 1 in Bai et al. (2022) , UOGD-OFU satisfies thatEQUATION) where V T,ℓ := T t=2 sup w∈W |ℓ(f (•; θ feat f ′′ t , w); P test t ) -ℓ(f (•; θ feat f ′′ t-1 , w); P test t-1 )|, σ denotes the min- imum singular value of the invertible confusion matrices {C f ′′ 1 ,D ′ 0 , • • • C f ′′ T ,D ′0 } and K is the number of classes and the expectation is taken with respect to randomness in the revealed co-variates.Proof Sketch: Recall that ℓ(f (•; θ feat f ′′ t , w); P test t ) := E (x,y)∼P test t ℓ ce f (x|θ feat f ′′ t , w), y .This guarantee follows from the arguments in Bai et al. (2022) from two basic facts below:1. The risk ℓ(f (•; θ feat f ′′ t , w); P test t) is convex in w over a convex and compact domain W.Severity of CIFAR-10C in the experiment. For each type of corruption in CIFAR-10C, we select a mild level and a high level of severity in the experiment section. Here we introduce the exact parameters of mild and high levels of severity for those corruptions. For Gaussian Noise, the severity levels for [mild, high] are [0.03, 0.07]. For Fog, the severity levels for [mild, high] are [(0.75,2.5), (1.5,1.75) ]. For Pixelate, the severity levels for [mild, high] are [0.75, 0.65] .Details of additional datasets. In addition to CIFAR-10 and CIFAR-10C, we evaluate on three more datasets: STL10 (Coates et al., 2011) , CINIC (Darlow et al., 2018) , and EuroSAT (Helber et al., 2019) . Similar to CIFAR-10, we split the original train sets of these datasets into the train set and the validation set by the ratio 4 : 1 and use the original test sets for sampling test images in the online test stage.Figure 5 shows the results for three additional datasets: STL10, CINIC, and EuroSAT, also with OLS-OFU (rotation degree prediction) under sinusoidal shift setting. In-line with CIFAR-10, we observe that OLS and OLS-OFU can perform better than Base and OFU, and OLS-OFU can outperform OLS. We find this pattern to be more consistent for EuroSAT and STL10 than CINIC.Figure 6 shows the results on CIFAR-10 for Bernoulli shift cross three SSL methods in OLS-OFU. Similar to Figure 2 (a), OLS-OFU within all three SSL methods outperforms all baseline methods. time step can be the negative examples. However, the batch size for a time step is small, e.g. 10 in our experiment. MoCo updates with such a small batch won't work. Thus, we experiment with MoCo by applying a batch accumulation strategy, which we introduce next.Prior test-time training (rotation degree prediction, entropy minimization) methods operate with the OFU framework easily, with feature extractor updates taking place in each time step. Rotation degree prediction originates as a self-supervised training method (Gidaris et al., 2018) , thus naturally we evaluate more recent self-supervised training methods, namely MoCo (He et al., 2019; Chen et al., 2020b; 2021) . No prior work shows how to use MoCo (or self-supervised learning in general) in a test-time training setting. Given that self-supervised training is sensitive to batch size, the intuition is that a larger batch size (much larger than the number of online samples available per time step) is required to perform a valid gradient update for a MoCo checkpoint. As such, we evaluated a batch accumulation strategy OLS-OFU (BA=τ ), where we continue evaluating online samples per time step, but only perform the online feature update at every τ steps (having accumulated the online samples throughout the τ steps in one batch). In particular, we perform feature extractor update every τ = 50 steps (for 1000 steps, feature update occurs 20 times, online samples evaluation occurs 1000 times), evaluating with 10 online samples per time step, using a smaller learning rate (0.0005) but test-time train with 10 epochs. Notice that τ = 1 is the default setting in Algorithm 1. OLS-OFU with MoCo presented in Figure 2 , Figure 6 and Figure 7 is equivalent to OLS-OFU (BA=50).To show the necessity of large τ , we evaluate both OLS-OFU (BA=1) and OLS-OFU (BA=50) on CIFAR-10 (Figure 11 ) and CIFAR-10C (Figure 12 ). Firstly, we can observe that OLS-OFU (BA=1) is even worse than OLS. We hypothesize this is because small batch size of MoCo will hurt the performance and larger batch size in MoCo is necessary. Hence, we increase τ from 1 to 50 and then we can observe the significant improvement from OLS-OFU (BA=1) to OLS-OFU (BA=50). Now, OLS-OFU (BA=50) can outperform OLS.We would like to compare the improvements in the literature for the OLS problem and the improvements from the SOTA in the literature to our method OLS-OFU. We can calculate the difference ∆ error between the errors of method A and B to show the improvement A → B. To measure the improvements in the literature for the OLS problem, we measure ∆ error between FTH (proposed in the first OLS paper) and FLHFTL (SOTA in the literature). We also calculate ∆ error between FLHFTL and FLHFTL-OFU.Figure 13 shows the ∆ error for FTH→FLHFTL and FLHFTL→FLHFTL-OFU across all data settings (dataset and online shift pattern) when the SSL in OF U is set as rotation degree prediction. The improvement of FLHFTL→FLHFTL-OFU is as significant as the improvement of FTH→FLHFTL. Moreover, the improvement of FLHFTL→FLHFTL-OFU is more consistent. This consistency demonstrates the potential that OLS-OFU can have improvement for any future OLS algorithm.Pseudo labelling (Lee, 2013) , a common self-training technique, generates pseudo labels for unlabelled data and uses them to update the model. Though we are not able to use ground-truth labels to compute feature extractor updates, we can use the model at time t to make predictions with respect to the online samples at time t, and train on the inputs with their assigned (pseudo) labels. An issue that arises in self-training is confirmation bias, where the model repeatedly overfits to incorrect pseudo-labels. As such, different methods can be used to select which samples will be pseudolabelled and used in updating the model, e.g. using data augmentation (Arazo et al., 2020) , usingAs pointed out in Section2, this is relatively general as most previous online label shift algorithms belong to these two categories.2 We detail how to obtain this enhanced model f ′′ t using SSL techniques in the following paragraphs.The reweighting factor qt is a function of unbiased estimators s1, • • • , st in FLHFTL and FTH.Code released at https://anonymous.4open.science/r/online_label_shift_with_online_feature_updates-3A1F/. 


Decision: This paper studies the problem of online generalized label shift adaptation, proposing a feature learning stage integrated with existing methods for online label shift. The authors theoretically demonstrate the soundness of their methods and conduct a variety of experiments to show their practical effectiveness.

The paper initially received mixed evaluations, so I invited another expert reviewer. After checking all the reviews and also paper myself, I find the paper's contributions interesting and valuable to the community. Specifically, integrating the feature learning (i.e., the representation ability of deep neural networks) into online label shift adaptation is highly useful for practical applications. Therefore, I recommend the acceptance. However, several reviewers have raised constructive feedback, including suggestions to enhance the experiments to make their claims more convincing. Moreover, it is also important to discuss the storage complexity in the initial stage. Other interesting issues include how to manage new classes during the online adaptation phase, etc. The authors are encouraged to take those comments into account when revising the paper. Some issues may be acknowledged as limitations or noted for future work. It is common in scientific research to focus on certain points and leave some limitations in future studies, but it is important to discuss and acknowledge them clearly.

Reviews: review: paper_summary: paper_summary: This paper addresses the problem of online (generalized) label shift, where label information is unavailable during testing, and the label distributions change over time. The main contribution of this work is the proposal of a unified framework that integrates feature learning into the online learning process, enabling the method to leverage the strengths of deep models. Theoretical analysis and experiments demonstrate the effectiveness of the proposed approach
strength_weakness: strength_weakness: - One of my main concerns is the theoretical analysis of the feature learning component. While Theorem 1 is commendable for demonstrating that the proposed method is comparable to the best model adapted from $f_t^{''}$, it remains unclear how effective the feature extractor obtained via self-supervised learning is (from a theoretical view). The self-supervised learning technique appears to be used as a black box, with no theoretical insight provided into its performance.

- Regarding storage costs: It seems that the proposed method requires storing the training set $D_0$. Such a requirement is somewhat unfavorable in practice, as the training set $D_0$ typically consists of a large volume of data.

- Concerning Principle 1: I understand that Principle 1 is essential for achieving theoretical guarantees in the online label shift problem. However, I am uncertain whether such a requirement is necessary in practice. Should the update procedure depend on the amount of data available at each round? If we have a reasonable amount of data at each time, wouldn't it be more reasonable to update the feature extractor before the online learning process to gather more information? Conversely, if the data at each iteration is limited, I'm unsure if the difference between update procedures is significant. It would be beneficial if the authors could provide a more detailed discussion on this matter.

I am happy to discuss these concerns with the authors and update my score if the questions are adequately addressed.

===post-rebuttal===

I appreciate the authors' efforts in addressing my questions and am satisfied with their feedback. I encourage the authors to incorporate the discussion from the rebuttal period into the main paper. I have raised my score to 7.
questions: questions: Could you provide more theoretical insight on the feature extractor? For instance, is there any guidance on selecting the SSL method or determining the step size for performing the gradient step?

- Is there any method to reduce the requirement of storing the train set $D_0$?

- Could you provide more justification for Principle 1? (Please refer to the three points of concern outlined above for more details.)
limitations: limitations: I did not identify the negative societal impact of this work.
paper_summary: paper_summary: This paper introduces a novel method for addressing label shifts in an online setting, where data distributions change over time, and obtaining timely labels is challenging. Unlike traditional approaches, this paper explores enhancing feature representations using unlabeled data during test time. The proposed method, Online Label Shift adaptation with Online Feature Updates (OLS-OFU), leverages self-supervised learning to refine the feature extraction process, thereby improving the prediction model. This approach is designed to maintain similar online regret convergence to existing results while incorporating improved features. Empirical evaluations show that OLS-OFU achieves substantial improvements over current methods, demonstrating that integrating online feature updates is as effective as the fundamental online label shift methods themselves. The results are consistent across various datasets and scenarios, highlighting the robustness and generality of OLS-OFU. The paper suggests that this method could be extended to more complex scenarios, such as online covariate shift and varying domain shifts over time.
strength_weakness: strength_weakness: - From a critical perspective, the online generalized label shift problem presented in the paper could be viewed as a sub-field of concept drift, where the relationship between x and y changes over time. Therefore, a detailed explanation of the differences between the proposed online generalized label shift scenario and the traditional concept drift is required. Additionally, it would be beneficial to conduct experimental comparisons to determine whether methods developed for concept drift could be effective in the experimental settings used in this study.
questions: questions: - What is the difference between concept drift and online generalized label shift situations?
limitations: limitations: n/a
paper_summary: paper_summary: This paper addresses the online label shift (OLS) adaptation problem, which involves continually adapting a pretrained offline model to test data with various and evolving label distribution shifts. The proposed method integrates existing self-supervised learning (SSL) techniques into current OLS methods, based on three proposed method-combination principles. Experimental results demonstrate that incorporating SSL methods leads to performance improvements.
strength_weakness: strength_weakness: 1. **Limited Novelty and Significance**: The proposed method appears to be a straightforward combination of existing approaches, specifically SSL and OLS methods. The three proposed principles for combining these approaches are relatively trivial:

    - Since SSL affects extracted features, OLS is performed first.

    - As SSL changes the feature extractor, the classifier is retrained.

    - Given that SSL requires additional training resources, it is applied after accumulating enough data.

    These principles seem too basic to be considered significant technical contributions.

2. **Problems in Paper Structure**: The paper has serious structural issues.
    The proposed method relies heavily on existing OLS methods, yet there is a lack of brief or detailed introduction of related OLS methods. Only three sentences describe OLS and SSL methods.

    Conversely, 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.

    As a result, the experimental section is compressed into two pages, limiting the space to present results adequately. The conclusion is also excessively brief, reduced to just one sentence.

3. **Unreasonable Problem Setting**: The addressed problem seems unreasonable for several reasons.

    - **Data Privacy Concern**: Adapting the model at each time step requires storing all historical training data, testing data, and models, raising significant data privacy concerns. For instance, in the MRI example mentioned in the paper, it is questionable whether you are allowed to carry MRI data from 999 clinics just to adapt the model to the 1000th clinic.

    - **Adequate Offline Training Data**: The offline training data appears overly sufficient, making the problem less challenging. In the CIFAR-10 experiment, all training data are used for offline training, with the only challenge being the variation in label distribution in the test data. Without any adaptation, the model performs well on the testing data.

4. **Low Baseline Performance**: The reported baseline result without any adaptation appears too low. A ResNet-18 model on CIFAR-10 typically achieves around 93\% accuracy on test data without any special data augmentation (see https://github.com/kuangliu/pytorch-cifar), which is significantly higher than the reported 84\% and most adaptation results.

5. **Minor Typos**:
    - Line 83: How do you reweight a model f? Do you mean reweight the model output?
    - Line 179: OGD or ROGD?
    - Lines 232-233: "Either...or" should be used correctly.
    - Line 249: Should it be "is" or "as"?
questions: questions: See Weaknesses.
limitations: limitations: Already addressed by Authors.
paper_summary: paper_summary: This paper addresses the problem of online label shift and proposes a novel algorithm that exploits feature representation learning to enhance performance. Inspired by test-time training literature, the proposed method uses self-supervised learning to refine the feature extraction process. The algorithm comes with strong theoretical guarantees, and experiments demonstrate its effectiveness.
strength_weakness: strength_weakness: 1. The proposed method requires the storage of previous historical data for self-supervised learning (batch accumulation in the paper). This may not be feasible in certain problems with privacy concerns.

2. The experiments are primarily conducted on simulated benchmark datasets, such as CIFAR-10 and CINIC, rather than real-world applications or datasets. This limits the understanding of the method's performance in practical settings.

3. It is recommended to include some recent work on online label shift in the paper, such as [1], which addresses the appearance of new class data in the scenario of online label shift.


[1] 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.
questions: questions: See weakness above.

In addtion, for the storage of historical data, is it feasible to apply data augmentation techniques to the limited number of data per time stamp and use the augmented data for self-supervised learning? How does the number of historical data stored affect the proposed algorithm?
limitations: limitations: the authors adequately addressed the limitations
paper_summary: paper_summary: The authors of this paper focus on the task of label distribution shift in the online setting without true labels. Following the importance of improving feature extractors even during test-time found in the current literature, they propose to update the feature extractor online with unlabeled instances during testing for online label shift and then refine the last classification layer with labeled instances. The proposed method is simple and intuitive, and extensive experiments over several label shift cases and benchmark datasets have shown its effectiveness.
strength_weakness: strength_weakness: 1.	The motivation of this paper is not clear. The authors are just motivated by “the potential for improving feature extractors”, and hypothesize “that a similar effect can be harnessed in the context of (generalized) label shift”. This may be too intuitive. Why is improving feature extractors helpful for only label shift? What can improve feature extractors gain? Are there some theoretical or empirical results to prove that? The theory of regret convergence may be not enough.
2.	The novelty of the proposed method may be limited. Improving feature extractors has been proved its importance during testing time for distribution shifts in the current literature, such as [1]. What is the difference between the proposed method and them? What are your main contributions in this paper?
3.	What are the limitations and broader impacts?

[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.
questions: questions: 1.	What is the difference between (online) label distribution shift and domain adaption? Can the proposed method adapt to the domain adaption task?
2.	What will happen if there are some instances related to new unseen categories during testing? And how to handle this case? Can improving the feature extractor be helpful?
limitations: limitations: See the above “*Weaknesses*” and “*Questions*” parts.

Comments: title: Rebuttal by Authors
Rebuttal by Authors
Rebuttal by Authors
comment: 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!
**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

|                               | FTFWH  | FTH    | ROGD   | ATLAS  | UOGD   | FLHFTL |
|-------------------------------|--------|--------|--------|--------|--------|--------|
| OLS-OFU($\tau=1$)             | 11.3%  | 11.2%  | 13.9%  | 11.6%  | 11.4%  | 11.2%  |
| OLS-OFU-difforder($\tau=1$)   | 12.33% | 12.12% | 14.35% | 12.10% | 11.91% | 12.08% |
| OLS-OFU($\tau=100$)           | 7.33%  | 8%     | 10.8%  | 10.1%  | 8.35%  | 7.45%  |
| OLS-OFU-difforder($\tau=100$) | 7.37%  | 8.05%  | 10.82% | 10.11% | 8.36%  | 7.48%  |

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.

|                               | FTFWH | FTH   | ROGD   | ATLAS  | UOGD  | FLHFTL |
|-------------------------------|-------|-------|--------|--------|-------|--------|
| OLS-OFU($\tau=100$)           | 7.23% | 7.91% | 10.81% | 10.12% | 8.33% | 7.53%  |
| OLS-OFU-difforder($\tau=100$) | 8.02%	 | 8.63% | 11.19% | 10.55% | 8.80% | 8.29%  |

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

| % stored training data | FTH-OFU | ROGD-OFU | FLHFTL-OFU | FTFWH-OFU |
|------------------------|---------|----------|------------|-----------|
| 100% (original)        | 8%      | 10.8%    | 7.45%      | 7.33%     |
| 80%                    | 8.18%   | 10.93%   | 7.62%      | 7.48%     |
| 60%                    | 8.68%   | 11.84%   | 8.04%      | 7.92%	     |
| 40%                    | 9.49%   | 12.50%   | 8.91%      | 8.40%     |
| 20%                    | 9.54%   | 12.63%   | 9.51%      | 8.91%     |
| 10%                    | 9.67%   | 12.82%   | 10.11%     | 9.86%     |
| 5%                     | 9.81%   | 12.94%   | 10.34%     | 10.03%    |
| 0%                     | 10.24%  | 13.50%   | 10.43%     | 10.41%    |
| OLS only               | 12.04%  | 13.65%   | 12.02%     | 11.9%     |

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).
Thank you for your reply. After reading other comments and rebuttals, I would like to keep my score.
Thanks for you response. Despite the feedback from the authors, most of the problems in this paper remain there.

**Limited novelty and significance.** The proposed method is a simple combination of existing OLS and SSL methods. The main contribution comes from the proposed three combination principles. While we agree that other options are not optimal, these principles are too trivial and lack of significance.

**Paper structure.** 
- Counting all the sentences the author mentioned, there are totally 5 sentences introducing the related works about OLS methods. 
- This paper has nothing to do with online distribution shift, including the detailed description about it instead of OLS methods is unreasonable and misleading, especially given that the proposed method heavily depends on previous OLS methods. We would not mind including online distribution shift if existing OLS methods have been thorougly described.
- One paragraph for the analysis of experimental results from line 317 to 330 are far away from enough. One sentence describing three figures are far away from enough.
- One sentence for conclusion is far away from enough.

**Data privacy concern.** The data privacy concern is raised not only from the batch accumulation, but also the requirement of the offline training data, all the training data from previous time step (<t). Here comes the same question, are you allowed to carry MRI data from the first 999 clinics just to adapt the model to the 1000th clinic?

**Too much offline training data.** 
- Both training and validation sets count as data used for training. Thus, all the training data from CIFAR-10 has been used for offline model training, leaving the addressed OLS setting unchallenging at all. 

**Low baseline performance**
- The reported low baseline accuracy has not been explained. From the included link (https://github.com/kuangliu/pytorch-cifar), ResNet-18 on CIFAR-10 achieves 93\% accuracy. From the original ResNet paper, ResNet-18 on CIFAR-10 achieves an accuracy of 91.25\%. CIFAR-10 dataset has its own fixed testing set. Without any adaptation, the offline model outperforms almost all the reported results. This raises the question, is OLS adaptation necessary? Why wouldn't we train the offline model adequately in advance?

Based on the forementioned reasons, I would recommend a clear Reject for this submission.
Thanks for your responses. I have read them and would like to keep my score.

Here are some clarifications for my score: I still have doubts regarding the experimental validity of concept drift. So, I leaned toward the accept side but stayed on the borderline.
Thanks for your responses. I have read them, and would like to keep my score.