One-step Label Shift Adaptation via Robust Weight Estimation

1 Introduction

The success of traditional machine learning methodologies is generally contingent on the closed environment hypothesis, which presumes that the training and testing data are independently and identically distributed [Mohri, 2018; Bengio et al., 2021]. However, real-world learning tasks frequently occur in open environments [Zhou, 2022], where data distributions may undergo temporal variations, potentially resulting in significant performance degradation of closed-environment systems [Sugiyama and Kawanabe, 2012; Huang and Ren, 2024]. Label shift [Fan et al., 2023; Li et al., 2024; Fan et al., 2024a], which represents a typical scenario in open environments, assumes that the label distributions in the source and target domains are distinct (Ps(Y ) ̸= Pt(Y )), whereas the conditional distributions given the labels remain consistent (Ps(X|Y ) = Pt(X|Y )). To illustrate, in the realm of COVID-19 diagnosis, incidence rates fluctuate across diverse regions, yet the symptomatic manifestations of pneumonia remain consistent. Another pertinent instance concerns bird identification, as exemplified in Figure 1. The migratory patterns of birds can induce seasonal fluctuations in their distribution within the same geographical area. Notably, despite these seasonal shifts, their morphological characteristics remain invariant. Given the aforementioned examples, label shift demonstrates considerable potential for application, thereby emerging as a pivotal area of research in recent years [Wei et al., 2024; Luo and Ren, 2024].

To mitigate the consequences of label shift and attain precise target output results, existing methodologies typically employ a two-step strategy to address the label shift issue [Zhao et al., 2021; Tian et al., 2023]. Initially, they calculate importance weight, which is subsequently followed by the adjustment of posterior probabilities for target samples. In the first step, certain research endeavors utilize a pre-trained source classifier to approximate the confusion matrix and class probabilities, and estimate the importance weight based on distribution transformation theory [Lipton et al., 2018; Azizzadenesheli et al., 2019]. Meanwhile, other studies derive the importance weight by minimizing Kullback-Leibler (KL) divergence between the weighted source distribution P w s (X) and the target distribution Pt(X) [Alexandari et al., 2020; Sipka et al., 2022]. Additionally, some methods aim to enhance the accuracy of weight estimation by incorporating the prior parameters of the target distribution [Sulc and Matas, 2019; Ye et al., 2024b]. In the second step, existing approaches generally involve either direct or indirect adjustment of the posterior probabilities for target samples. Specifically, indirect adjustment methods leverage the equivalence between weighted source and target loss expectations to minimize the importanceweighted empirical risk, thereby obtaining the final target classifier and correcting the target outputs [Garg et al., 2020; Fan et al., 2024a]. On the other hand, direct adjustment methods combine the importance weight with the outputs generated by the source classifier, ensuring unbiased target outputs directly under the framework of distribution transformation theory [Wen et al., 2024].

Although the aforementioned methods have demonstrated outstanding performance within the label shift framework, there are at least two critical challenges that require careful consideration. Firstly, when tackling the label shift problem, precise importance weight is instrumental in deriving accurate target outputs, and vice versa. However, prevalent label shift methods typically follow a sequential process of weight estimation followed by output adjustment, overlooking the intricate interplay between these two components. Secondly, numerous traditional methodologies are heavily contingent upon the efficacy of the pre-trained source classifier; in instances where this classifier exhibits suboptimal performance, the resultant bias in weight estimation can be substantial, ultimately culminating in a deterioration of the final target outputs.

In this paper, to effectively address the aforementioned challenges, we propose a novel bi-level approach denominated One-step Label Shift Adaptation (OLSA) via Robust Weight Estimation, as shown in Figure 2. Specifically, the upper-level task minimizes the unbiased target loss utilizing importance weight, whereas the lower-level objective estimates this importance weight based on target outputs. Therefore, weight estimation and classifier training can be jointly done through a bi-level optimization strategy. To enhance the robustness of our proposed model, we incorporate a bias correction term in the upper-level task. Under mild assumptions, theoretical analysis demonstrates that OLSA can attain the same expected risk as conventional label shift methods. Furthermore, we introduce a regularization term in the lower-level objective to mitigate the dependency of weight estimation on the trained classifier, thereby ensuring stable performance across diverse label shift scenarios. We derive a generalization bound for the final classifier. Finally, comprehensive experimental results are presented to validate the efficacy of OLSA. In summary, the contributions of our research are outlined as follows:

• Our research focuses on developing a method to address the label shift issue in a single step. This strategy avoids the necessity for intermediary steps, thereby enabling a more streamlined and efficient process for training the target classifier and estimating importance weight.
• To enhance the robustness of our proposed model, we incorporate a debiasing term into the upper-level task and devise a regularization term for the lower-level objective. Under mild assumptions, theoretical proof has been provided for the efficacy of our approach.
• We demonstrate the effectiveness of our approach across diverse datasets. The experimental results consistently indicate that our approach outperforms other comparative methods in the majority of cases, particularly in scenarios involving significant label shifts.

2 RelatedWork

Label shift is a prevalent scenario in open environments, sparking significant interest among researchers [Tasche, 2017; Bai et al., 2022]. This phenomenon occurs when there exists a disparity in label distributions between the source and target domains, while preserving the same conditional distribution [Wu et al., 2021; Fan et al., 2023]. Previous works, such as BBSE [Lipton et al., 2018] and RLLS [Azizzadenesheli et al., 2019], estimate importance weight using the confusion matrix and predicted target labels, followed by retraining a new target classifier within the Weighted Empirical Risk Minimization (WERM) framework. In addition, MLLS [Alexandari et al., 2020] and SCML [Sipka et al., 2022] estimate importance weight by minimizing the KL divergence between the weighted source and target distributions, and directly correct the target outputs through a posterior adjustment strategy. Building on these foundations, MAP [Sulc and Matas, 2019] and MAPLS [Ye et al., 2024b] introduce the Bayesian posterior of the target label distribution parameters given data. From a novel perspective of matching the source label distribution, CPM [Wen et al., 2024] maintains the same theoretical guarantees as traditional feature distribution matching frameworks, while significantly enhancing computational efficiency due to the direct matching of label variables. However, the aforementioned approaches employ a two-step strategy, which neglects the interaction between output adjustment and weight estimation.

One-step Domain Adaptation is designed to directly learn an effective target classifier under distribution shift. Vapnik's principle [Vapnik, 2013] emphasizes the avoidance of solving a general problem as an intermediary step when confronted with limited information, thereby highlighting the significance of one-step approach. DAOT [Peng et al., 2018] conceptualizes the feature distribution alignment process as a one-step transformation and implements it through a single-layer convolutional neural network. OCSA [Zhang et al., 2020] introduces a one-step strategy for accommodating covariate shifts, without the intermediate step of estimating the ratio between training and test input densities. Despite these advancements, the above methods have not taken into account how to perform one-step learning in the label shift scenario. ADM-OS [Fan et al., 2024b] emerges as a tailored one-step method for label shift, which jointly learns the predictive model and importance weight. However, empirical findings have indicated that ADM-OS may exhibit limited robustness under specific shifting conditions.

3 Problem Setting

3.1 Traditional WERMFramework

We employ stochastic variables X ∈ X , Y ∈ Y to model the features and labels respectively, where X = Rd and Y is a discrete set as {1, 2, ...,K}. d is the feature dimension, and K represents the number of categories. We have access to labeled source data (Xs, Ys) = {xi, yi}ni=1 and unlabeled target data Xt = {xj}mj=1 , which are independently and identically drawn from the source distribution Ps and target distribution Pt respectively. We define L(·, ·) : ∆K-1 × Y → R1 as the loss function and h ∈ H: X ↦→ ∆K-1 as the classifier, where ∆K-1 denotes the standard K-dimensional probability simplex. Under the label shift assumption, i.e., Ps(X|Y ) = Pt(X|Y ) and Ps(Y ) ̸= Pt(Y ), if we give a hypothesis space H and a loss function L, the goal of label shift setting is to find an optimal target classifier ht ∈ H which minimizes the following weighted loss [Lipton et al., 2018; Azizzadenesheli et al., 2019; Garg et al., 2020; Fan et al.,

where w∗ ∈ RK is the importance weight and w∗ (Y ) = Pt(Y )/Ps(Y ). In practical scenarios, the veritable weight w∗

is commonly undisclosed. If w∗ is approximated as ˆw, the target classifier can be derived by minimizing the following empirical loss:

3.2 The OLSAApproach

In this section, we aim to address the label shift issue using a one-step strategy. Inspired by the theory of hyper-parameter optimization [Liu et al., 2022; Liu et al., 2024b], we treat the importance weight ˆw as a hyper-parameter to be optimized, and then introduce a bi-level optimization technique for its solution. Specifically, the upper-level loss is employed for training the target classifier, while the lower-level loss is utilized for estimating importance weight. Now, let us introduce the specific forms of upper-level and lower-level losses respectively.

Upper-level Loss. Existing label shift methods for estimating importance weight require the availability of target outputs. Motivated by this observation, we attempt to investigate an innovative approach that incorporates unlabeled target data into training loss. The most prevalent approach is semi-supervised learning [Schmutz et al., 2023; Ye et al., 2024a], which leverages both labeled and unlabeled data simultaneously to train the classifier. However, due to the discrepancy between the source and target distributions, it is unreasonable to directly apply traditional methods. Therefore, we introduce the following weighted semi-supervised loss:

where Ht(·) := ∫Y Pt(Y |·)L(h(·), Y )dY represents a labelindependent component, while β ∈ (0, 1) is a balance parameter that governs the equilibrium between the labeled and unlabeled terms. However, the introduction of the second item naturally increases the risk of traditional label shift loss

ˆRT. Therefore, we attempt to leverage labeled data on labelindependent term to mitigate potential bias and design the following upper-level loss:

where Hs(·; ˆw) := ˆw(Y )∫Y Ps(Y |·)L(h(·), Y )dY and the source conditional distribution Ps(Y |·) can be approximated by the output of source classifier with the softmax layer. Fortunately, our approach achieves the same expected risk as traditional label shift methods, as demonstrated in Theorem 1.

Theorem 1. Assume the importance weight ˆw reaches its optimal value Pt(Y )/Ps(Y ), we have

The comprehensive proof of the aforementioned theorem is provided in supplementary materials. On this basis, given the assumption that the importance weight ˆw is known, we can directly optimize the loss function ˆRDeM(h; ˆw) in order to obtain the target classifier. Subsequently, the critical task is to determine an accurate approach for estimating the importance weight ˆw.

Lower-level Loss. In the first, according to the properties of distributions, we give a simple estimation strategy for the weights ˆw:

Since the labels of the source data are known, we can estimate the source label distribution Ps(Y ) through statistical class frequency, i.e.,

On this basis, the key issue becomes how to estimate the conditional distribution Pt(Y |X). A classic approach is to approximate the true conditional distribution by considering the output of softmax layer in a deep neural network, i.e., ht(x) ≈ Pt(Y |X). However, previous studies [Guo et al., 2017; Liu et al., 2024a] have demonstrated that conventional neural networks encounter challenges in accurately estimating class posterior probabilities, despite their proficiency as classifiers. Most label shift works use bias-corrected temperature scaling (BCTS) [Alexandari et al., 2020; Sipka et al., 2022] to calibrate outputs. However, as it is a post-processing technique, it is not applicable to our one-step approach. Therefore, for good classifier calibration, we introduce γ-loss, which is define as follows:

By minimizing the expectation of the above loss [Sypherd et al., 2022], we derive the optimal outputs as

The parameter γ has the ability to 'soften' the outputs. When γ = 1, the output is consistent with traditional softmax output, and as γ → 0, the output h(x) gravitates towards 1/K, indicating maximum uncertainty. By combining Eq. (6), we estimate the importance weight ˆw as

Thus, we can design the lower-level loss as follows:

However, the aforementioned estimation is heavily influenced by the classifier outputs h(x). If the calibration effect is poor, it can lead to significant weight bias. Thus, we try to add a regularization term to alleviate this situation.

Where λ controls how much the label distribution is skewed. When λ → ∞, w = 1 KPs(Y ) and Pt(Y ) = 1K . This regularization term ensures that the target distribution is balanced even when the estimation deviation is large, thus ensuring the stability of training. Through the transformation of Eq. (12), we get the final lower-level loss as

Total Loss. By combining the upper-level loss Eq. (4) and lower-level loss Eq. (13), we have the total loss as follows:

Then, we prove an upper bound for the generalization error of OLSA as follows.

Theorem 2. Give n samples drawn from the source distribution Ps, m samples drawn from the target distribution Pt and bound loss function L and H. Then, there exists two constants κ1, κ2 > 0, that depends on β, w∗ , L and H, the following generalization bound holds with probability at least 1 - δ,

RT(h;w∗ ) ≤ ˆRDeM(h; ˆw) +2Radn+m

where Radn+m is the Rademacher complexity.

The detailed proof can be found in supplementary materials.

3.3 Optimization

The aforementioned loss constitutes a bi-level problem, necessitating the utilization of an optimization framework grounded in implicit gradients (IG). This strategy enables precise modeling and optimization of the intricate interplay between the estimated weight ˆw and model parameters h. Specifically, the IG derived from Eq. (15) can be mathematically formulated as follows:

where ∇h and ∇ˆw represent the partial derivatives of the bivariate function. Traditional bi-level optimization approaches derive the IG formula through the meticulous application of implicit function theory. However, the computation of IG presents inherent difficulties, primarily arising from the complexities involved in matrix inversions, second-order partial derivatives, and constraints. In order to effectively compute IG, it is evident that the optimal solution ˆw∗ of the unconstrained lower-level loss function is formulated as follows:

Since we adopt the softmax layer in our model, the output h(x) satisfies that h(x) ⪰ 0 and ∑Ki=1 hi(x) = 1 for any x ∼ Pt(X). Thus we have

On this basis, we can derive the optimal solution for the constrained lower-level loss as ˆw∗ , and compute the IG (d ˆw∗T/dh) directly. In conclusion, the procedure for optimizing Eq. (14) is summarized in Algorithm 1.

Algorithm 1 Procedure of OLSA approach

Input: The labeled source samples {xi, yi}ni=1 , unlabeled target samples {xi}n+m i=n+1 , trade-off parameters λ, β and γ, upper-level learning rate υ and a pre-trained source classifier hs(x) ≈ Ps(Y |x).

Initialize: Initialize network parameters h0 = hs;

for Iteration j = 1, ..., T do

Lower-level: Update the weight ˆwj by optimizing Eq. (17);

Upper-level: Update the network parameters ˆhj via stochastic gradient descent calling Eq. (16); end for

Use the updated classifier ˆhT to make predictions {ˆ yi}n+m i=n+1 on the unlabeled target samples;

Output: The ultimate classifier ˆhT and the predicted labels {ˆ yi}n+m i=n+1 .

End procedure

4 Experiments

In this section, we commence by introducing the datasets, comparative methodologies, network architectures, and parameter settings. Building upon these foundational elements, we undertake a comprehensive evaluation of the performance and effectiveness of the proposed OLSA approach, focusing on two key aspects. For the initial component, we conduct a comparative analysis of OLSA with traditional label shift methods across diverse shift scenarios and evaluation metrics. In the second component, we offer detailed outcomes derived from in-depth examinations of OLSA across various dimensions, encompassing an analysis of the deviation in importance weight estimation, an exploration of various parameter configurations and a visualization of label distribution bias.

4.1 Configuration

Dataset. In our study, we assess the performance and efficacy of OLSA on the MNIST [LeCun et al., 1998], FasionMNIST [Xiao et al., 2017], CIFAR10 [Krizhevsky et al., 2009] and CIFAR100 [Krizhevsky et al., 2009] datasets, incorporating numerous artificial shifts. Specifically, we introduce two distinct shift categories in our experimental setup: (1) TweakOne shift, where the probability of a specific source class is altered to ρ (larger values of ρ result in more extreme label shift), while the probabilities of remaining classes maintain their original proportions. (2) Dirichlet shift, which generates a Dirichlet distribution utilizing the concentration parameter α (smaller values of α result in more extreme label shift), and aligns the source label distribution with this Dirichlet distribution. For the MNIST and Fasion-MNIST datasets, we allocate 2000 samples each for the training and validation sets, and 10,000 samples for the test set. Analogously, for the CIFAR10 dataset, we assign 4000 samples each for training and validation, and 10,000 samples for testing. For the CIFAR100 dataset, the distribution is 10,000 samples for training, 5000 for validation, and 20,000 for testing. Notably, to guarantee the robustness of training process, each class within the shift set is represented by a minimum of 40 samples.

Compared methods. In the main experiment subsection, we demonstrate the performance of OLSA by conducting a comparative analysis with the following traditional label shift techniques. (1) Wout-W illustrates the performance of the base classifier in the scenario where the estimated weight is excluded from consideration. (2) BBSE [Lipton et al., 2018] and RLLS [Azizzadenesheli et al., 2019] represent two label shift methods that are grounded in the utilization of hard confusion matrices, while BBSE-S and RLLS-S are based on the employment of soft confusion matrices. (3) MLLS [Alexandari et al., 2020], CML [Sipka et al., 2022], SCML [Sipka et al., 2022], MAPLS [Ye et al., 2024b] and CPMCN [Wen et al., 2024] are five advanced label shift methods that incorporate classifier calibration and obtain the target outputs directly. (4) ADM-OS is a one-step label shift method that alternately iterates between importance weight estimation and classifier training.

Network architectures and evaluation indicators. We can adopt any model as the base model for the source classifier, and the choice of the model has implications for the accuracy of subsequent importance weight estimation. In our primary experimental configuration, we employ a two-layer fully connected neural network for the MNIST and Fasion-MNIST datasets, while concurrently utilizing ResNet-18 [He et al., 2016] for both the CIFAR-10 and CIFAR-100 datasets. We sample the data 5 times using the specified shift parameter and calculate the average as final outputs. To assess the effectiveness of OLSA, we employed accuracy (Acc), F-score, and mean squared error (MSE) as evaluation metrics, where MSE is defined as follows:

Parameter settings. The parameters for compared methods are chosen based on the technologies delineated in their respective references. In our study, we report the results of different shift parameters α ∈ [0.1, 0.5, 1.0, 5.0] and ρ = [0.3, 0.5, 0.7, 0.9]. In addition, we fix the trade-off parameter β = 0.1 empirically, while the calibration parameter γ is selected from the discrete set [0.8, 0.9, 1, 2] and the regularization parameter λ is chosen form the discrete set [0, 0.1, 1, 10] through the validation set results.

4.2 Main Results

To demonstrate the efficacy of our OLSA approach, all methods run on framework with Python 3.7 and PyTorch based on the same pre-trained classifier. We present the Acc (shown in supplementary materials) and F-score (shown in Table 1) on both the Dirichlet and Tweak-One shift datasets, and have the following observations.

1. While the performance of different comparative methods varies across diverse datasets, our OLSA approach demonstrates a consistent improvement in the performance of existing label shift methods in most cases. For example, on the Dirichlet shift CIFAR10 dataset with α = 0.1, OLSA exhibits significant improvements, achieving an increase of nearly 6% in the F-score.
2. When compared to the performance outcomes observed under conditions of small label shifts, OLSA demonstrates a more significant enhancement in performance under scenarios involving large label shifts. For example, on the Tweak-One shift Fasion-MNIST dataset with small shift ρ = 0.3, OLSA exhibits a decrement of 1% in the F-score, while under large shift ρ = 0.9, achieves a notable increase of approximately 4% in the F-score.

These findings underscore the robustness and stability of OLSA in addressing large label shift scenarios.

1. It is noteworthy that, on the CIFAR100 dataset, OLSA attains a substantial improvement in the F-score, albeit with less satisfactory results in terms of Acc. This disparity arises from the significant weight estimation bias exhibited by the compared methods, which results in highly precise predictions for classes with abundant samples while overlooking the less represented classes.

4.3 Performance Analyses

MSE comparison

To demonstrate the efficacy of our shift strategy, we introduce MSE indicator, which represents the deviation between the estimated and true weights. The results are depicted in Fig. 3 (a-d), from which several key observations have been derived. (1) The performance of OLSA remains superior to compared methods in the majority of cases, especially in cases where the degree of shift is significant. This fully demonstrates the stability of our method, indicating its applicability to most shift scenarios. (2) We observe that accurate weights do not necessarily lead to a better classifier among the comparative methods. However, such a scenario rarely occurs in our approach. This indirectly demonstrates the effectiveness of OLSA, which simultaneously trains the classifier and estimates the importance weight.

Parameter sensitivity analysis

Here, we check the Acc of the trade-off parameters λ and γ on the Tweak-One shift MNIST and Dirichlet shift FasionMNIST datasets respectively. The results are visually presented in Fig. 3 (e-f). We find that for different datasets and shift scenarios, the optimal value of γ varies, which underscores the necessity of adjusting it, showcasing the effectiveness of the calibration loss. Furthermore, when the value of λ is large, the model's performance typically remains good, which illustrates the effectiveness of introducing the lowerlevel regularization term.

Visualization of label distribution bias

In order to demonstrate that OLSA approach can alleviate label shift, we present some visualization results in Figure 3 (g-h), which shows the label shift degree ˆ Pt(Y ) - Pt(Y ) of the Wout-W and OLSA methods. A smaller degree of label shift indicates a more accurate estimation of the target label distribution. As can be observed from figures, OLSA approach demonstrates robust performance, even under conditions of significant label shift, with a deviation from the true label distribution remaining within 0.05.

5 Conclusion

In this paper, we introduce an innovative one-step methodology to address the label shift problem, which is crucial but rarely studied. Utilizing bilevel optimization techniques, OLSA strengthens the interplay between weight estimation and classifier training. Dependent on the introduced regularization terms, OLSA approach opens up an interesting frontier for the robust one-step modeling of scenarios involving large label shifts. The solid theoretical analysis and enriched experimental analysis fully demonstrate the effectiveness of our OLSA approach. In future work, attempts can be made to enhance the robustness of weight estimation by refining the regularization term. This may be achieved by leveraging data structures, among other approaches.

Acknowledgments

This work was partially supported by the National Key Research and Development Program (No. 2022ZD0114803), the NSF for Distinguished Young Scholars under Grant No. 62425607 and the Key NSF of China under Grant No. 62136005. Chenping Hou and Tingjin Luo are the corresponding authors.