SlowGuess/NeurIPS_2023_OCR huggingface dataset:
FlatMatch: Bridging Labeled Data and Unlabeled Data with Cross-Sharpsness for Semi-Supervised Learning
Zhuo Huang $^{1}$ , Li Shen $^{2}$ , Jun Yu $^{3}$ , Bo Han $^{4}$ , Tongliang Liu $^{1,*}$
$^{1}$ Sydney AI Centre, The University of Sydney;  $^{2}$ JD Explore Academy;
$^{3}$ Department of Automation, University of Science and Technology of China;
$^{4}$ Department of Computer Science, Hong Kong Baptist University
zhua0420@uni.sydney.edu.au, mathshenli@gmail.com, harryjun@ustc.edu.cn, bhanml@comp.hkbu.edu.hk, tongliang.liu@sydney.edu.au
Semi-Supervised Learning (SSL) has been an effective way to leverage abundant unlabeled data with extremely scarce labeled data. However, most SSL methods are commonly based on instance-wise consistency between different data transformations. Therefore, the label guidance on labeled data is hard to be propagated to unlabeled data. Consequently, the learning process on labeled data is much faster than on unlabeled data which is likely to fall into a local minima that does not favor unlabeled data, leading to sub-optimal generalization performance. In this paper, we propose FlatMatch which minimizes a cross-sharpness measure to ensure consistent learning performance between the two datasets. Specifically, we increase the empirical risk on labeled data to obtain a worst-case model which is a failure case that needs to be enhanced. Then, by leveraging the richness of unlabeled data, we penalize the prediction difference (i.e., cross-sharpness) between the worst-case model and the original model so that the learning direction is beneficial to generalization on unlabeled data. Therefore, we can calibrate the learning process without being limited to insufficient label information. As a result, the mismatched learning performance can be mitigated, further enabling the effective exploitation of unlabeled data and improving SSL performance. Through comprehensive validation, we show FlatMatch achieves state-of-the-art results in many SSL settings. Our code is available at https://github.com/tmllab/2023_NeurIPS_FLatMatch.
Training deep models [31, 45] not only requires countless data but also hungers for label supervision which commonly consumes huge human laboring and unaffordable monetary costs. To ease the dependency on labeled data, semi-supervised learning (SSL) [6, 12, 27] has been one of the most effective strategies to exploit abundant unlabeled data by leveraging scarce labeled data. Traditional SSL commonly leverages unlabeled data by analyzing their manifold information. For example, semi-supervised support vector machine [6, 13] finds a classifier that crosses the low-density region based on large-margin theorem and label-propagation [60, 84] computes an affinity matrix in the input space to propagate the labels to their neighbor unlabeled data. Due to the computational burden of exploring manifold knowledge, recent advances in SSL benefit from sophisticated data augmentation techniques [4, 19, 20, 74, 80], they usually enforce predictive consistency between the original inputs and their augmented variants [75, 8, 7, 57, 66, 71, 76, 79], meanwhile using pseudo labels [46, 52, 53, 2] to guide the learning on unlabeled data.
*Corresponding to Tongliang Liu.
Figure 1: Loss landscapes of labeled data and unlabeled data obtained simultaneously from training using FixMatch [57] on CIFAR10 with 250 labels per class. The first row and second row show the results obtained from epoch 60 and epoch 150, respectively. The first column and second column show the 2D loss contours of labeled data and unlabeled data, respectively, and the last column shows the 1D loss curves.
However, as most of the existing methods only apply instance-wise consistency on each example with its own transformation, labeled data and unlabeled data are disconnected during training. Hence, the label guidance cannot be sufficiently propagated from the labeled set to the unlabeled set, causing the learning process on labeled data is much faster than on unlabeled data. As a result, the learning model could be misled to a local minima that cannot favor unlabeled data, leading to a non-negligible generalization mismatch. As generalization performance is closely related to the flatness of loss landscape [14, 24, 41, 51], we plot the loss landscapes of labeled data and unlabeled data using FixMatch [57] in Fig.  $1^{2}$ . We have two major findings: 1) The loss landscape of labeled data is extremely sharp, but the loss curve of unlabeled data is quite flat. 2) Moreover, the optimal loss value of labeled data is much lower than that of unlabeled data. Such two intriguing discoveries reveal the essence of SSL: The learning on scarce labeled data converges faster with lower errors than on unlabeled data, but it is vulnerable to perturbations and has an unstable loss curvature which rapidly increases as parameters slightly vary. Therefore, the abundant unlabeled data are leveraged so that SSL models are fitted to a wider space, thus producing a flatter loss landscape and generalizing better than labeled data. Based on the analysis, existing methods that use instance-wise consistency with mismatched generalization performance could have two non-trivial pitfalls:
- The label supervision might not be sufficiently leveraged to guide the learning process of unlabeled data.
- The richness brought by unlabeled data remains to be fully exploited to enhance the generalization performance of labeled data.
As seeking the connection in input space is quite prohibitive, we resort to exploring the parameter space and propose FlatMatch which targets to encourage consistent learning performance between labeled data and unlabeled data, such that SSL models can obtain strong generalization ability without being limited by insufficient label information. Specifically, as the loss landscape on labeled data is quite unstable, we aim to mitigate this problem through applying an adversarial perturbation [24, 43, 69, 82, 83] to the model parameters so that the worst failure case can be found. Further, we enforce the worst-case model and the original model to achieve an agreement on unlabeled data via computing their prediction differences which are dubbed cross-sharpness. By minimizing such a cross-sharpness measure, the richness of unlabeled data can be fully utilized to calibrate the learning direction. In turn, the label information can be activated for producing accurate pseudo labels, thus successfully bridging the two datasets to enable improved SSL performance.
Note that the loss curves cannot be directly generated using the raw data from training, because networks can fit all examples perfectly without showing any generalization evidence. Therefore, we load the data for plotting using slight data augmentation, i.e., Random Rotation with  $90^{\circ}$ , which is very commonly used in SSL.
Furthermore, the proposed FlatMatch is easy to implement and can be adapted to complement many popular SSL methods [10, 15, 16, 28]. Although computing the proposed cross-sharpness normally requires an additional back-propagation step, we propose an efficient strategy so that our FlatMatch can be processed without extra computational burden. Through extensive experiments and analyses on various benchmark datasets, we show that our FlatMatch achieves huge performance improvement compared to some of the most powerful approaches. Specifically, on CIFAR100 dataset with 2500 labeled examples, our FlatMatch surpasses the foremost competitor so far with  $1.11\%$  improvement. In general, our contribution can be summarized into three points:
- We identify a generalization mismatch of SSL due to the disconnection between labeled data and unlabeled data, which leads to two critical flaws that remain to be solved (Section 1).
- We propose FlatMatch which addresses these problems by penalizing a novel cross-sharpness that helps improve the generalization performance of SSL (Section 4.2).
- We reduce the computational cost of FlatMatch by designing an efficient implementation (Section 4.3).
- Extensive experiments are conducted to fully validate the performance of FlatMatch, which achieves state-of-the-art results in several scenarios (Section 5).
SSL has been developed rapidly since the booming trend of deep learning [31, 45, 25]. We roughly introduce it through three stages, namely traditional SSL, data-augmented SSL, and open-set SSL.
Traditional SSL: One of the first SSL methods is Pseudo Label [46] which has been very effective for training neural networks in a semi-supervised manner, which relies on the gradually improved learning performance to select confident samples [73, 72]. Another classical method is Co-training [9] which utilizes two separate networks to enforce consistent predictions on unlabeled data so that the classification can be more accurate. Inspired by co-training, many SSL approaches are motivated to conduct consistency regularization. Temporal ensembling [44] proposes to leverage historical model predictions to ensemble an accurate learning target that can be used to guide the learning of the current model. Mean Teacher [59] further suggests such an ensemble can be operated on the model weights to produce a teacher model that can provide improved learning targets. Instead of improving the teacher model, Noisy student [75] finds that adding noise to the student model can also improve SSL.
Data-Augmented SSL: Thanks to the development of data augmentation techniques and self-supervised learning [30, 34, 68, 67], further improvements in SSL are achieved by enforcing consistency between the original data and their augmented variants. VAT [49] proposes to add adversarial perturbations [25, 26] to unlabeled data. By using the pseudo-label guidance to leverage such adversarial unlabeled data, the model can be more robust to strong input noises, thus showing better performance on the test set. MixMatch [8] employs the MixUp [80] technique to interpolate between the randomly shuffled examples which largely smooths the input space as well as the label space, further enhancing the generalization result. Moreover, FixMatch [57] is based on strong augmentation such as AutoAugment [19] and RandAugment [20] to utilize the predictions of weakly augmented data to guide the learning of strongly augmented data, which can achieve near supervised learning performance on CIFAR10. Recently, Dash [76], FlexMatch [79], and FreeMatch [66] advanced SSL through investigating the threshold strategy for pseudo labels. They analyze the training dynamic of SSL and design various approaches to enhance the selection of unlabeled data, which effectively boosts the SSL performance to a new level.
Open-Set SSL: Moreover, another branch of SSL studies the realistic scenario where unlabeled data contains out-of-distribution (OOD) data [37, 38, 32, 33, 63, 64, 62]. The common goal is to alleviate the influence of OOD such that the learning process of SSL will not be misled. Several studies propose many techniques that can achieve this goal, such as distillation [17], meta-learning [29], large-margin regularization [10], content-style disentanglement [36], and consistency regularization can also help [56]. But in this paper, we mainly focus on the common assumption that labeled data and unlabeled are sampled independently and identically.
3 Preliminary: Improving Generalization via Penalizing Sharpness
Sharpness-Aware Minimization (SAM) [24, 43, 47, 69, 81] which focuses on optimizing the sharp points from parameter space so that the training model can produce a flat loss landscape. Specifically, given a set of fully-labeled data  $\mathcal{D}^l = \{(x_i^l,y_i^l)\}_{i = 1}^n$  containing  $n$  data points  $(x_{i}^{l},y_{i}^{l})$ , SAM seeks to optimize a model parameterized by  $\theta \in \Theta$  to fit the training dataset  $\mathcal{D}^l$  so that  $\theta$  can generalize well on a test set  $\mathcal{D}^{te}$  drawn independently and identically as  $\mathcal{D}^l$ . Such an optimization process is normally conducted via minimizing the empirical risk  $\mathcal{L}_l(\theta) = \frac{1}{n}\sum_{i = 0}^n\ell_{ce}(\theta ;g_\theta (x_i^l),y_i^l)$  which is realized by the cross-entropy loss  $\ell_{ce}$  computed between the label prediction  $f^{l} = g_{\theta}(x_{i}^{l})$  and the class label  $y_{i}^{l}$ . To find the sharp points that maximally increase the empirical risk, SAM applies a weight perturbation  $\epsilon \in \Theta$  to the model parameters  $\theta$  and conducts the following inner maximization:
$$
\epsilon^ {*} (\theta) := \underset {\| \epsilon \| _ {p} \leq \rho} {\arg \max } \mathcal {L} _ {l} (\theta + \epsilon) \approx \underset {\| \epsilon \| _ {p} \leq \rho} {\arg \max } \epsilon^ {\top} \nabla_ {\theta} \mathcal {L} _ {l} (\theta) ^ {p = 2} \approx \rho \frac {\nabla_ {\theta} \mathcal {L} _ {l} (\theta)}{\| \nabla_ {\theta} \mathcal {L} _ {l} (\theta) \| _ {2}}, \tag {1}
$$
where  $\epsilon^{*}$  denotes the optimal perturbation which can be estimated using the gradient information of mini-batch inputs,  $\rho$  restricts the perturbation magnitude of  $\epsilon$  within a  $\ell_p$ -ball, and the approximation is given by the dual norm problem [24].
By perturbing the original model  $\theta$  with  $\epsilon^{*}$ , we can obtain the sharpness measure from SAM: sharpness :=  $\mathcal{L}_l(\theta + \epsilon^*) - \mathcal{L}_l(\theta)$  which indicates the how quickly loss changes under the worst-case model perturbation  $\epsilon^{*}$ . To combine the sharpness term with empirical risk minimization (ERM), the SAM objective can be differentiated as follows:
$$
\begin{array}{l} \nabla_ {\theta} \mathcal {L} _ {l} ^ {S A M} := \nabla_ {\theta} \left[ \mathcal {L} _ {l} (\theta + \epsilon^ {*} (\theta)) - \mathcal {L} _ {l} (\theta) \right] + \mathcal {L} _ {l} (\theta) \approx \nabla_ {\theta} \mathcal {L} _ {l} (\theta + \epsilon^ {*} (\theta)) \\ = \frac {d \left(\theta + \epsilon^ {*} (\theta)\right)}{d \theta} \nabla_ {\theta} \mathcal {L} _ {l} (\theta) | _ {\theta + \epsilon^ {*} (\theta)} = \nabla_ {\theta} \mathcal {L} _ {l} (\theta) | _ {\theta + \epsilon^ {*} (\theta)} + o (\theta), \tag {2} \\ \end{array}
$$
where the last second-order term  $o(\theta)$  can be discarded without significantly influencing the approximation. For detailed derivation and analysis, we refer to the original paper and related studies [1, 24, 82]. Intuitively, the gradient of SAM is computed on the worst-case point  $\theta + \epsilon^{*}$ , then such a gradient is utilized to update the original parameter  $\theta$ . As a result, the model can produce a more flat loss landscape than ERM which would not change drastically, and being robust along the sharpest direction  $\epsilon^{*}$  among its  $\ell_{p}$ -norm neighbor.
Since the obtained flatness property has been demonstrated to have many benefits for generalization, such as being resistant to distribution shift [11, 39, 61], robustness against adversarial attack [69, 58], and effectiveness under label noise [24, 3, 35, 40, 70, 77], SAM has inspired many research progresses. However, SAM is only conducted under a fully-supervised setting, and whether or how can it be leveraged to improve learning with unlabeled data is still undiscovered. Next, we carefully introduce the proposed FlatMatch which improves the generalization performance of SSL by bridging labeled data and unlabeled data with a novel cross-sharpness.
4 FlatMatch: Semi-Supervised Learning with Cross-Sharpness
In this section, we carefully introduce our FlatMatch. First, we describe the SSL setting in a generalized way. Then, we demonstrate a novel regularization from FlatMatch which is dubbed cross-sharpness. Finally, we explain its optimization and design an efficient implementation.
4.1 General Semi-Supervised Learning
In SSL, we are commonly given a small set of labeled data  $\mathcal{D}^l = \{(x_i^l,y_i^l)\}_{i = 1}^n$  containing  $n$  labeled examples  $(x_{i}^{l},y_{i}^{l})$  and a large unlabeled data  $\mathcal{D}^u = \{x_i^u\}_{i = 1}^m$  containing  $m$  ( $m\gg n$ ) unlabeled examples  $x_{i}^{u}$  which are drawn independently and identically. Similar to previous notations, we aim to optimize a deep model  $\theta \in \Theta$  in a semi-supervised manner so that  $\theta$  can perform well on i.i.d sampled test set  $\mathcal{D}^{te}$ . The general objective for SSL is as follows:
$$
\min  _ {\theta} \mathcal {L} _ {s s l} = \min  _ {\theta} \mathcal {L} _ {l} + \mathcal {L} _ {u} = \min  _ {\theta} \frac {1}{n} \sum_ {i = 0} ^ {n} \ell_ {c e} \left(\theta ; g _ {\theta} \left(x _ {i} ^ {l}\right), y _ {i} ^ {l}\right) + \frac {1}{m} \sum_ {i = 0} ^ {m} \mathbb {I} (\hat {y} _ {i} > \tau) \ell_ {d} \left(\theta ; g _ {\theta} \left(\mathcal {A} \left(x _ {i} ^ {u}\right)\right), \hat {y} _ {i}\right), \tag {3}
$$
Figure 2: Illustration of FlatMatch. The blue arrows and red arrows denote the learning flows related to labeled data and unlabeled data, respectively; and the black arrows indicate the computation of the worst-case model.
where the first term is similar to Eq. 1 and denotes the empirical risk under the label supervision, the second term indicates the semi-supervised regularization for exploiting unlabeled data, the function  $\ell_d(\cdot)$  denotes the loss criterion for unlabeled data, such as KL-divergence or cross-entropy loss,  $\mathcal{A}(\cdot)$  stands for an augmentation function that transforms the original input  $x_{i}$  to its augmented variants,  $\hat{y}_i$  represents the unsupervised learning target which can be a pseudo label or a guiding signal from a teacher model, and an indexing function  $\mathbb{I}(\cdot)$  is commonly used to select the confident unlabeled data if their learning target surpasses a predefined threshold  $\tau$ . Our FlatMatch replaces the second term with a cross-sharpness regularization as described below.
4.2 The Cross-Sharpsness Regularization
The computation of cross-sharpness can be briefly illustrated in Fig. 2. Particularly, to solve the sharp loss landscape problem, as demonstrated in Section 1, we penalize the cross-sharpness regularization on each unlabeled example using a worst-case model derived from maximizing the loss on labeled data. Formally, such optimization is shown as follows:
$$
\min  _ {\theta} \mathcal {R} _ {X - s h a r p} := \min  _ {\theta} \frac {1}{m} \sum_ {x ^ {u} \in \mathcal {D} ^ {u}} \ell_ {d} \left(g _ {\tilde {\theta}} \left(x ^ {u}\right), g _ {\theta} \left(x ^ {u}\right)\right), \text {w h e r e} \tilde {\theta} = \theta + \arg \max  _ {\| \epsilon \| _ {p} \leq \rho} \mathcal {L} _ {l}, \tag {4}
$$
in which  $g_{\theta}(\cdot)$  indicates the forward propagation using the parameter  $\theta$  whose prediction result is  $f = g_{\theta}(\cdot)$ , and  $\tilde{\theta}$  stands for the worst-case model that maximally increases the empirical risk on labeled data. Here we drop the subscript  $i$  from each datum  $x$  for simplicity. By meaning "cross", the sharpness is not obtained by revisiting labeled data as SAM does, instead, we compute the sharpness by leveraging abundant unlabeled data to take full advantage of their generalization potential.
In detail, we leverage the original model  $\theta$  to conduct a first propagation step as  $f^{l} = g_{\theta}(x^{l})$ ;  $f^{u} = g_{\theta}(x^{u})$ , where  $f^{l}$  and  $f^{u}$  are label predictions from labeled data and unlabeled data, respectively. Then, the empirical risk on labeled data is computed as  $\mathcal{L}_l = \sum_{x^l\in \mathcal{D}^l}\ell_{ce}(\theta ;g_\theta (x^l),y^l)$ . Following Eq. 1, we maximize the empirical risk  $\mathcal{L}_l$  to obtain the weight perturbation as  $\epsilon^{*}\coloneqq \rho \frac{\nabla_{\theta}\mathcal{L}_{l}(\theta)}{\|\nabla_{\theta}\mathcal{L}_{l}(\theta)\|_{2}}$ , further having the worst-case model  $\tilde{\theta} = \theta +\epsilon^{*}$ . Moreover, we conduct a second propagation step by passing the unlabeled data to  $\tilde{\theta}$  and have  $\tilde{f}^u = g_{\tilde{\theta}}(x^u)$ . The second unlabeled data prediction  $\tilde{f}^u$  is combined with the first prediction  $f^u$  obtained before to compute the cross-sharpsiness in Eq.  $4^{3}$ .
Why does crossing the sharpness work? As shown by many theoretical signs of progress, the generalization error of SSL is largely dependent on the number of labels [5, 48]. Especially in recently proposed barely-supervised learning [57], the generalization error would be enormous. Although such a result may be pessimistic, unlabeled data can still be helpful. It is shown that when two distinct hypotheses on labeled data are co-regularized to achieve an agreement on unlabeled data, the generalization error can be reduced, as we quote: "The reduction is proportional to the difference between the representations of the labeled data in the two different views, where the measure of difference is determined by the unlabeled data" [54]. Intuitively, disconnecting labeled data and unlabeled data during training may be sub-optimal. To achieve better theoretical performance, the SSL model should focus on learning something from unlabeled data that are not contained by labeled
Note that there is a slight difference between the sharpness from SAM and our cross-sharpness: The former one is measured by the change of cross-entropy loss under model perturbation. But in SSL, the pseudo labels of unlabeled data is not accurate enough compared to labeled data, so we use the difference between label predictions of  $\theta$  and  $\tilde{\theta}$  to measure our cross-sharpness.
Algorithm 1 FlatMatch and FlatMatch-e
Input: Labeled dataset  $\mathcal{D}^l$ , unlabeled dataset  $\mathcal{D}^u$ , model  $\theta$ , memory buffer  $\mathcal{M}$  for storing historical gradients, and EMA factor  $\alpha$  for updating buffer.  
1: for  $t \in 0,1,\ldots,T-1$  do  
2: if Efficient training then ▷ FlatMatch-e  
3: Compute optimal perturbation as  $\epsilon^* = \rho \frac{\mathcal{M}^t}{\|\mathcal{M}^t\|_2}$ ;  
4: else ▷ FlatMatch  
5: Compute gradient as  $\nabla_\theta \mathcal{L}_l(\theta) = \nabla_\theta \sum_{x^l \in \mathcal{D}^l} \ell_{ce}(\theta; g_\theta(x^l), y^l)$ ;  
6: Compute optimal perturbation as  $\epsilon^* := \rho \frac{\nabla_\theta \mathcal{L}_l(\theta)}{\|\nabla_\theta \mathcal{L}_l(\theta)\|_2}$ ;  
7: Obtain worst-case model as  $\tilde{\theta} = \theta + \epsilon^*$ ;  
8: Compute cross-sharpness via Eq. 4;  
9: Optimizing  $\theta$  via Eq. 5, meanwhile save the gradient  $\nabla_\theta \mathcal{L}_l(\theta)$ ;  
10: Update memory buffer as  $\mathcal{M}^{t+1} = (1 - \alpha) \times \mathcal{M}^t + \alpha \times \nabla_\theta \mathcal{L}_l(\theta)$ ;
data. In our FlatMatch, we can achieve this goal by enforcing  $\theta$  and  $\tilde{\theta}$ , which produces the maximal loss difference on labeled data, to be consistent in classifying unlabeled data. Therefore, FlatMatch can be more effective than other instance-wise consistent methods by bridging the knowledge of labeled data and unlabeled data for training.
4.3 Optimization and Efficient Implementation
To this end, we solve the optimization of FlatMatch by substituting  $\mathcal{L}_u$  in Eq. 3 with  $\mathcal{R}_{X\text{-} sharp}$  in Eq. 4:
$$
\theta := \underset {\theta} {\arg \min } \mathcal {L} _ {l} (\theta) + \mathcal {R} _ {X - s h a r p} \left(\theta + \epsilon^ {*} (\theta)\right) \approx \nabla_ {\theta} \mathcal {L} _ {l} + \nabla_ {\theta} \mathcal {R} _ {X - s h a r p} (\theta) | _ {\theta + \rho \frac {\nabla_ {\theta} \mathcal {L} _ {l} (\theta)}{\| \nabla_ {\theta} \mathcal {L} _ {l} (\theta) \| _ {2}}}. \tag {5}
$$
Particularly, the second gradient is obtained on the worst-case point  $\theta + \rho \frac{\nabla_{\theta} \mathcal{L}_l(\theta)}{\|\nabla_{\theta} \mathcal{L}_l(\theta)\|_2}$ , then it is applied on the original model  $\theta$  for gradient descent. As we can see, such optimization requires an additional back-propagation on labeled data, which doubles the computational complexity. Fortunately, our FlatMatch can be implemented efficiently with no extra computational burden than the baseline stochastic gradient descent (SGD) optimizer, which is dubbed FlatMatch-e. The implementation is summarized in Algorithm 1.