Correcting Covariate Shift with the Frank-Wolfe Algorithm

1 Introduction

Many machine learning algorithms assume that training and test data come from the same distribution, which is often violated in practical applications. Researchers have thus endeavoured to resolve the distribution shift problem under varied assumptions [Zadrozny, 2004; Bickel et al., 2007; Quionero-Candela et al., 2009]. In this work, we focus on the covariate shift scenario [Shimodaira, 2000] in which the marginal data distributions are different between training and test domains (ptrpxq � ptepxq) while their conditional distributions remain the same (ptrpy|xq � ptepy|xq). Correcting covariate shift has a wide range of applications, such as in natural language processing [Jiang and Zhai, 2007], off-policy reinforcement learning [Hachiya et al., 2009], computer vision [Yamada et al., 2012] and signal processing [Yamada et al., 2010].

A common approach to correcting covariate shift is importance reweighting: each individual training point is assigned a positive weight intended to diminish the discrepancy between training and test marginals by some criterion [Sugiyama et al., 2012]. If a reweighting function is modelled properly, covariate shift can be effectively corrected in learning. For example, for a given prediction function fpxq and a loss function lpfpxq, yq, reweighting training points with weight wpxq � ptepxq{ptrpxq can minimize the loss over the test distribution:

which suggests that one should seek the function f � that minimizes this expected loss. While these correction techniques remain effective for many real world applications, most of them do not exploit the structure of the solution and as a result, do not scale well in practice.

Recently, the Frank-Wolfe (FW) algorithm has begun to gain popularity in the machine learning community [Zhang et al., 2012; Bach, 2015]. It has been proven to be an efficient algorithm for many optimization problems, particularly when the solution has sparse structure [Jaggi, 2013], and has also been shown effective and efficient in many applications [Joulin et al., 2014; Salamatian et al., 2014]. However, whether Frank-Wolfe can be applied to the covariate shift problem has remained unexplored.

In this work, we show that herding [Welling, 2009; Chen et al., 2010] (as a Frank-Wolfe algorithm) can be applied to the covariate shift scenario, and point out a connection between kernel herding and Kernel Mean Matching (KMM) [Gretton et al., 2009]. Moreover, we exploit the structure of another commonly used covariate shift correction technique, Kullback-Leibler Importance Estimation Procedure (KLIEP) [Sugiyama et al., 2008], to speed up the algorithm using Frank-Wolfe. We also analyse the convergence rate and complexity of KMM(FW) and KLIEP(FW), and present an empirical study that demonstrates the efficiency of Frank-Wolfe in solving KMM and KLIEP over the traditional projected gradient approach.

2 RelatedWork

This work is closely related to kernel herding [Chen et al., 2010], which is a sampling technique for moment approximation. In the original setting, data points are iteratively generated to approximate the population mean (first moment), either from a population when the distribution p is known, or selected from existing dataset when p is unknown. This work is different in that we apply herding to the covariate shift scenario, where the goal is to approximate the test mean by sub-sampling existing training data.

Even though there are many ways to correct covariate shift, few methods scale well, which hampers their usage in practice. Tsuboi et al. [2009] modifies KLIEP with a different parametric form of the weight function. Although the computation time of the new model is independent of the number of test points, the parametric form is less interpretable than the original KLIEP. Here we apply Frank-Wolfe to the original KLIEP, where the weight function is parametrized by a mixture of Gaussians. Another attempt has been made to solve KMM and KLIEP via online learning [Agarwal et al., 2011]. However, the convergence rate of that approach, Op1{?tq, is slower than our proposed method, Op1{tq, and assumes the availability of explicit feature representation of φpxq, contrary to the idea of Hilbert space estimation.

There are also several other techniques that correct covariate shift. For instance, RuLSIF [Yamada et al., 2011] minimizes the α-relative Pearson divergence using least squares, while RCSA [Wen et al., 2014] corrects covariate shift in an adversarial setting. We do not extensively investigate the possibility of applying Frank-Wolfe to these techniques, since Frank-Wolfe is efficient primarily when a convex problem has sparse solution structure and the corresponding linearised problem can be solve easily.

3 A Brief Review of Herding

Kernel herding [Chen et al., 2010] is a greedy algorithm that sequentially generates instances, φpxtq in the tth iteration, to minimize a squared error E 2 T :

where φp�q is a feature map to a reproducing kernel Hilbert space (RKHS) H and µp � Ex�prφpxqs is the population mean. Algorithmically, it generates instances according to:

for some properly chosen u0, where X is the sampling space. In the original kernel herding algorithm, the expectation is taken with respect to the training population, i.e., p � ptr.

Bach et al. [2012] showed that herding is equivalent to a Frank-Wolfe algorithm (also known as conditional gradient descent) for the objective

where M is the marginal polytope (the convex hull of all φpxq for x P X). Specifically, it employs the following updates

where ρt P r0, 1s is step size. They pointed out that herding corresponds to using a step size ρt � 1{pt � 1q, while other choices (e.g., via line search) are also valid. With ρt � 1{pt� 1q, st can be seen as φpxtq while p µt � µ�ut{t as in (1).

4 Frank-Wolfe for Covariate Shift

4.1 Kernel Herding and KMM

In this work, we apply kernel herding to the covariate shift scenario, where the training and test marginal distributions are different (ptrpxq � ptepxq) while their conditional distributions are the same (ptrpy|xq � ptepy|xq). We use herding to generate instances from the training pool in order to approximate the test mean, i.e., where X is the training set but p � pte is the test marginal distribution. Intuitively, herding actively selects representative training points to approximate the test population mean. With proper substitution of empirical estimates, our objective is

where x M is marginal polytope of training dataset and m is the number of (unlabelled) test points. According to (2), in each update step, a representative training point φpxtq will be chosen from x M and p µt will be moved toward that direction. This objective is very similar to that of Kernel Mean Matching (KMM) [Gretton et al., 2009]: 1

where n training points and m test points are given. Problem (3) (a linearly constrained quadratic program) is generally solved by standard solver with gradient descent and feasibility projection. However, according to the interpretation in (2), herding for KMM is projection-free: with a proper step size ρt P r0, 1s, the estimated mean p µt is a convex combination of previous chosen training points, and thus the constraint W in (3) is automatically satisfied.

4.2 Frank-Wolfe for KLIEP

Once herding is viewed as a Frank-Wolfe algorithm, it can be applied to other covariate shift correction procedures. For instance, the Kullback-Leibler Importance Estimation Procedure (KLIEP) [Sugiyama et al., 2008], a popular choice for covariate shift correction, can be efficiently solved by FrankWolfe. The objective of KLIEP is to minimize the KullbackLeibler divergence from ptepxq to p ptepxq, where p ptepxq is attained by reweighting the training marginal:

The reweighting function is further parametrized as a mixture of Gaussians:

1 Here the predefined bound on the weight (wi / B) and the derivation tolerance (| �ni�1 wi �1| / ϵ) are ignored for the sake of clarity and simplicity. However, they can be easily incorporated into the herding scheme.

,

where αl are the mixing coefficients, cl are predefined centres (usually test points) and σ is a parameter selected by some criterion. After some derivation and substitution of empirical marginals, the final objective becomes

This is a convex optimization problem. To apply FrankWolfe, we first derive the gradient of the objective

Considering the constraint, one can observe that the possible range for αl is r0, n{ �ni�1 ϕlpxiqs. For the tth iteration, we choose lt such that

then update the current αt by

where ρt is the step size and el � r0, � � � , 0, 1, 0, � � � , 0sT (all zeros except for a one in the lth entry). Compared to KMM(FW), this method chooses representative Gaussians to match test marginal, instead of choosing individual points.

The chosen Gaussian centres (i.e., tclu) are crucial for KLIEP. In the original paper, the tclu are set to be the test points txju, which can be computationally intensive when one has many test points. Therefore, for large test data, the original authors proposed to randomly select b � 100 test points as Gaussian centres. However, such an approach can significantly reduce the performance of the algorithm, since the 100 chosen centres might not appropriately model the distribution difference, a phenomenon we observe in our experiments (see Section 6.2). If Frank-Wolfe is applied, only one Gaussian is amplified per iteration and the optimal solution can be efficiently attained [Jaggi, 2013], since it tends to be sparse [Sugiyama et al., 2008]. More importantly, we can use all the test points as Gaussian centres in order to utilize their information to capture the distribution difference.

5 Convergence and Complexity

5.1 Convergence Rate

In this section, we compare the convergence and complexity of projected gradient (PG) and Frank-Wolfe methods. For convergence, since both KMM and KLIEP are convex problems, it is well known that the convergence rate of PG is Op1{tq [Bertsekas, 1999]. The convergence rate of Frank-Wolfe for KMM and KLIEP is also Op1{tq without further assumptions [Jaggi, 2013]. However, under some circumstances, Op1{t2 q and even Ope�t q are possible for KMM(FW) [Chen et al., 2010; Beck and Teboulle, 2004]. Moreover, the solutions to KMM and KLIEP tend to be sparse in practice (for example, see Section 6.1), which is a very suitable scenario for the Frank-Wolfe algorithm.

5.2 Complexity Analysis

Next we analyze the time and space complexities of the algorithms. Since both KMM and KLIEP are convex, the computation of the objective is not necessary in each iteration. Instead, one can simply check the gradient of the objective and stop when the norm of the gradient is sufficiently small.

We first focus on KMM. In this case, one can rewrite the objective in (3) in matrix form as

where

and the constant term is ignored. For the projected gradient method, the bottleneck is the computation of the gradient. From (7) it is obvious that Opn2 q multiplications are required to compute the gradient in each iteration. This cannot be simplified because the projection to W creates a nonlinear operation on wt in each iteration. Projected gradient requires additional Opnq space to store the gradient. For Frank-Wolfe using a 1{pt � 1q step size, we can see from (2) that the time complexity is Opnq per iteration. As st�1 must be a single φpxi�q for some i� , we only need to maintain a score list of size n to store txp µt � µp, φpxiqyu for all i. After decomposition, xµp, φpxiqy is simply pkteqi and can be precomputed, while xp µt, φpxiqy can be easily maintained for each t as

This is faster than the projected gradient approach because there is no projection step in Frank-Wolfe, and the score list can be maintained efficiently with only Opnq additional space. Finally, when line search is applied to ρt, the time complexity is still Opnq, since the line search step size is attained in closed form ρt � xµp�pµt, st�1�pµty }st�1�pµt}2 . We can decompose the inner products as before and see that all operations can be performed in Opnq time. While the computation of }p µt�1}2 � wTt�1Ktrwt�1 seems to require Opn2 q multiplications, we can further decompose it, using the fact that wt�1 � p1 � ρtqwt � ρtδwt, as

where pKtrq:i� means the i� th column of Ktr. Therefore, only Opnq multiplications are required since the first term can be used recursively and δwt only has one singular non-zero entry. Similar tricks cannot be applied to projected gradient because δwt is not sparse in that case, and δwTt Ktrδwt still requires Opn2 q computation in general. An additional Opnq storage is still needed for the score list in the line search case.

Next we investigate the complexity of KLIEP. Similar to KMM, the bottleneck is the computation of the gradient (5). The projected gradient method requires 2mb multiplications to compute (5) for all l, while Frank-Wolfe only requires 2m�mb multiplications, since the denominator is linear in α and can be updated efficiently due to (6). This can lead to different run time in practice (see Section 6). Line search is less efficient for KLIEP(FW) compared to KMM(FW) as there is no closed form solution to ρt. Note that Fpαt�1q � Gpρtq is a concave function of ρt. Moreover, the gradient G1 p�q is decreasing in (0,1) with G1 p0q ¡ 0. If G1 p1q ¥ 0, a step size of 1 is used, otherwise G1 p1q � 0 and we use binary search on the interval r0, 1s to find ρt such that G1 pρtq � 0. Such a binary search requires Opmq multiplications. All three versions require Opm�bq space to compute and store the gradient. Table 1 summarizes the results of this section.

6 Experiments

In this section, we demonstrate the results of the proposed approach on both synthetic and some benchmark datasets. In particular, we compare the proposed Frank-Wolfe method on KMM and KLIEP with the projected gradient (PG) method. In the experiments, a Gaussian kernel is applied to KMM where the kernel width is chosen to be the median of the pairwise distances over the training set. For KLIEP, the width is chosen according to the criterion of Sugiyama et al. [2008].

6.1 Synthetic Datasets

To investigate the efficiency of the different methods, we compare their runtime to reach the same accuracy. Synthetic data is generated from y � x3 � x � 1 � ϵ, where ϵ � Np0, 0.12 q. Training points are generated from Np0.5, 0.52 q while test points are generated from Np0, 0.32 q [Shimodaira, 2000]. In the experiment, n � m. We first run the projected gradient method to convergence, then run the Frank-Wolfe methods (with two different step sizes: ρt � 1{pt�1q and line search) until it reaches the same objective. Figure 1 shows the runtime in log scale with varied sample sizes. The FrankWolfe methods are consistently faster than the projected gradient method. In the case of KLIEP, although Frank-Wolfe and projected gradient have the same theoretical complexity and convergence as shown in Section 5, their actual runtime can differ in practice, as the constant factor in the big-O notation can be different. Also, note that line search has better efficiency in the case of KLIEP, since the final solution is usually extremely sparse, meaning that line search can find it in a few iterations while 1{pt � 1q takes some more time.

We next check the solution quality of the different methods on the same problems, to determine the effectiveness of Frank-Wolfe. We generated n � m � 128 data points and show the final training weights in Figures 2 and 3. On one hand, the weights of the three KMM methods differ from each other significantly, but their objective values remain very close. This implies that although the objective of KMM (3) is convex, there are many "optimal" solutions (in terms of objective values) that minimize the discrepancy between the means in the RKHS. This also leads to different behaviours in the underling task (regression or classification, etc). No further criteria is proposed in the literature to distinguish which solution is preferable for KMM. For example, we include a red line and a blue line in each graph, showing the linear models learned by least squares from the unweighted (i.e., with uniform weights) and reweighted training set, respectively. The blue lines of the three KMMmethods are slightly different from each other. More results can be found on benchmark datasets in Section 6.2. On the other hand, although KLIEP also demonstrates multiple solutions with close objective values, the final weights produced by the three KLIEP methods are very similar, as shown in Figure 3. Note that Frank-Wolfe tends to find sparser solutions than projected gradient in general, which suggests the effectiveness of Frank-Wolfe in finding sparse solutions if they exist.

6.2 Benchmark Datasets

Next we applied the reweighting methods on some benchmark datasets from the libsvm 2 and delve 3 libraries to show their performance in correcting covariate shift on reweighted learning. All methods are run until convergence.

For each dataset, we first normalize the input features and the output to r�1, 1s if it is significantly out of scale. We introduce a covariate shift by the following scheme. First we compute the first (with largest eigenvalue) principle component of the dataset. Let zmax and zmin be the maximum and minimum projected values along the principle direction and σz be the standard deviation of the projected values. In each trial, we first draw m � 2000 test points without replacement from the pool, with probability proportional to N pzmax, σ2 z {4q. After eliminating these test points from the dataset, we then draw n � 5000 training points without replacement from the pool, with probability proportional to N pzmin, σ2 z {4q. After computing the weights for the training points, we learn a linear model from the reweighted set and evaluate the model performance on the test set, using hinge loss for classification and squared loss for regression. We compare the performance of the reweighting methods with the unweighted method (i.e., with uniform weights for training points). For KMM(FW) we use 1{pt � 1q as the step size, while for KLIEP(FW) we use the line search step size, since these choices are faster in practice. In addition to the projected gradient solver KLIEP(PG), we also compare the KLIEP(FW) method with the implementation of the original authors KLIEP(100), where the "b � 100" trick is used for projected gradient when there are too many test instances

2 www.csie.ntu.edu.tw/�cjlin/libsvm

3 www.cs.toronto.edu/�delve/data/datasets.html

(only randomly select 100 Gaussian centres; see Section 4.2 for more details).

Table 2 shows the test loss and Table 3 shows the runtime over 50 trials. First, note that the performance of KMM is not very reliable because 1) the kernel width used in the experiment is difficult to tune, as observed by others [Cortes et al., 2008] and 2) the solution to KMM tends to be sensitive to random initialization - i.e., different solutions may have similar objective values - which leads to unstable behaviour on the underlying task. Next, we can see that KLIEP(FW) is effective in correcting covariate shift. In most cases its test loss is smaller than unweighted learning as well as KLIEP(PG) on the test set. One possible reason is that KLIEP(PG) may fail to improve the objective on a step and stop prematurely. Comparing KLIEP(PG) and KLIEP(100), we can see that using all information from the test set will probably reduce test loss, although it may suffer from increased runtime. Finally, judging from Table 3, Frank-Wolfe is a very efficient algorithm compared to projected gradient in the current setting where sparse solutions are possible. KLIEP(100) is extremely fast but its task performance is so restrictive that it becomes impractical. We may apply values other than b � 100, but the best value of b is dataset-dependent and could be hard to identify. Instead, FW does not require parameter tuning and can utilize all Gaussians without computational issues.

7 Conclusion

We have proposed an efficient method that can be incorporated into two well-known covariate shift correction techniques, Kernel Mean Matching (KMM) and Kullback-Leibler Importance Estimation Procedure (KLIEP). Our approach is inspired by noticing a connection between kernel herding and KMM. By exploiting sparse solution structure, we apply the Frank-Wolfe algorithm to KMM and KLIEP. Since Frank-Wolfe is projection-free, its time complexity per iteration is in general smaller than the traditional projected gradient method. Our convergence and complexity analysis show that KMM(FW) has an advantage over the projected gradient method, while KLIEP(FW) has runtime performance comparable to the projected gradient method. Our empirical studies show that Frank-Wolfe is very suitable and practical for finding optimal solutions and correcting covariate shift as the sample size grows. Whether Frank-Wolfe can be applied to other covariate shift correction techniques remains an open question. We are also exploring ways to incorporate covariate shift correction to the underlying task for better performance.