FlexMatch: Boosting Semi-Supervised Learning with Curriculum Pseudo Labeling

In plain words

Semi-supervised learning trains an image classifier on a few labeled pictures and many unlabeled ones. A common trick is to let the model label the unlabeled pictures itself, but only keep the labels it is confident about. The cut-off is usually one fixed value for every class.

That is unfair to hard classes. A class the model has not learned well rarely reaches the cut-off, so its pictures are mostly ignored, and the model stays weak on it. A class it has learned well gets the same strict cut-off, so there is little reason to relax it.

FlexMatch gives each class its own cut-off and changes it during training. A class the model is still struggling with gets a lower cut-off, so more of its unlabeled pictures are used. A class that is already well learned gets a cut-off close to the original strict value. The method keeps no extra model and needs no extra pass over the data, so it adds almost no cost.

The name refers to curriculum learning: the thresholds change as training progresses, so classes are admitted in stages.

TL;DR

• FlexMatch applies Curriculum Pseudo Labeling (CPL) to FixMatch: the fixed confidence threshold becomes a per-class threshold that moves with the model's learning status §3.1.
• The learning status of a class is the number of unlabeled samples currently marked with that class, which the training loop already produces. No extra parameters or forward passes are needed §3.1.
• Author-reported: on CIFAR-100 with 400 labels, error falls from 46.42% (FixMatch) to 39.94%; on STL-10 with 40 labels, from 35.97% to 29.15% tab. 1.
• It is not uniformly better: on SVHN and at 250 labels on CIFAR-10 the authors' own table shows FixMatch as good or better tab. 1 §4.1.

Problem

Consistency and pseudo-labeling methods use unlabeled data only when the model's prediction is confident enough, which is controlled by a threshold τ. FixMatch and related methods (Pseudo-Labeling, UDA) use one constant τ for all classes §1.

A constant τ has two costs. Early in training, few unlabeled samples pass τ, so much unlabeled data is unused §1. Also, all classes are treated equally, although classes differ in difficulty and so in how quickly their predictions become confident §1 §3.1.

Using a per-class threshold with a labeled validation set is impractical in the semi-supervised setting: the labeled set is small and evaluating every class at every step slows training. The paper rules this out and instead estimates learning status from the training predictions §3.1.

Method

FlexMatch keeps the FixMatch training loop and changes only the threshold applied to the unsupervised loss §2 §3.1. Each unlabeled sample gets two views: a weak one used to produce a pseudo label, and a strong one (RandAugment) on which the loss is computed §2 §4.

1. Per-sample marks. Every unlabeled sample n stores a mark: the class it was last predicted as with confidence above the fixed cut-off τ, or -1 if it has never passed τ. Marks are updated during training, not recomputed each step §3.1 alg. 1.
2. Learning effect. For each class c, σ_t(c) is the number of samples whose mark equals c eq. 5 §3.1.
3. Normalisation. β_t(c) = σ_t(c) / max_c σ_t(c) rescales the effects to [0, 1], so the best-learned class gets β = 1 eq. 6 §3.1.
4. Warm-up. Early in training most samples are unused and σ is small and noisy. The denominator is therefore replaced by the larger of max_c σ_t(c) and the number of unused samples, N − Σ_c σ_t(c) eq. 11 §3.2.
5. Mapping. A convex map M(x) = x / (2 − x) turns β into a flexible threshold T_t(c) = M(β_t(c)) · τ. Here M maps [0, 1] onto [0, 1], so thresholds range from 0 to τ and the mapping adds no hyperparameter eq. 12 §3.3.
6. Mask and unsupervised loss. A weak-view sample contributes only if its maximum probability exceeds T_t of its predicted class. The loss is the cross-entropy between the pseudo label and the strong-view prediction, averaged over the whole unlabeled batch, masked samples included as zeros eq. 8 §3.1.
7. Total loss. L = L_s + λ · L_u with λ = 1 eq. 9 eq. 10 app. A.1.

The procedure is summarised in the paper's algorithm, which the steps below follow (with the mapping of step 5 from the text, see Pitfalls) alg. 1.

# FlexMatch, one training iteration  [alg. 1] [eq. 5-12]
# Input:  labeled batch (x, y); unlabeled batch u (B * mu samples)
# State:  marks[n] per unlabeled sample n, -1 = unused  [alg. 1]
# Output: scalar loss

def flexmatch_loss(labeled_x, labeled_y, unlabeled_u, marks, model):
    # 1. Learning effect per class: samples currently marked c  [eq. 5]
    counts = count_per_class(marks)           # (C,)
    unused = count(marks == -1)               # N - sum(counts)
    # 2. Normalise; warm-up denominator while unused dominates  [eq. 6, 11]
    denominator = max(max(counts), unused)    # drop unused for eq. 6
    beta = counts / denominator               # (C,), in [0, 1]
    # 3. Flexible threshold per class: T(c) = M(beta(c)) * tau  [eq. 12]
    thresholds = (beta / (2 - beta)) * tau    # (C,)
    # 4. Weak view: confidence and hard pseudo label, no gradient  [eq. 8]
    probs_weak = softmax(detach(model(unlabeled_u.weak)))  # (B*mu, C)
    confidence = max(probs_weak, over classes)             # (B*mu,)
    pseudo = argmax(probs_weak, over classes)              # (B*mu,)
    # 5. Keep a sample if it beats the threshold of its class  [eq. 8]
    keep = confidence > thresholds[pseudo]                 # (B*mu,)
    # 6. Store the class of samples above tau; the mark persists  [alg. 1]
    confident = confidence > tau                           # (B*mu,)
    marks[indices[confident]] = pseudo[confident]
    # 7. Strong view: masked cross-entropy to the pseudo label  [eq. 8]
    logits_strong = model(unlabeled_u.strong)              # (B*mu, C)
    loss_u = mean(keep * cross_entropy(logits_strong, pseudo))
    # 8. Supervised loss on the labeled batch  [eq. 10]
    loss_s = mean(cross_entropy(model(labeled_x), labeled_y))
    # 9. Weighted total  [eq. 9]
    return loss_s + lambda_u * loss_u

Key equations, as written in the paper:

\mathcal{L}_{u,t} = \frac{1}{\mu B} \sum_{b=1}^{\mu B} \mathbb{1}\big(\max(q_b) > \mathcal{T}_t(\arg\max q_b)\big)\, H(\hat{q}_b, p_m(y \mid \Omega(u_b))) eq. 8

\beta_t(c) = \frac{\sigma_t(c)}{\max\{\max_c \sigma_t,\; N - \sum_c \sigma_t(c)\}} eq. 11

\mathcal{T}_t(c) = \mathcal{M}(\beta_t(c)) \cdot \tau, \qquad \mathcal{M}(x) = \frac{x}{2 - x} eq. 12 §3.3

Where the paper leaves a symbol open, see Pitfalls: the pseudo label type in Eq. 8 is not defined in the text.

Hyperparameters

Values for the main experiments §4 and the appendix tables app. A.1. Algorithm-specific values are for FlexMatch, which is FixMatch with CPL.

Results

Author-reported. Setup: CIFAR-10, CIFAR-100, SVHN and STL-10 with WRN backbones; baselines Pseudo-Labeling, UDA and FixMatch, re-implemented in the authors' PyTorch codebase §4 app. B. Metric: best error rate over all checkpoints, mean ± std over three seeds. The median of the last 20 checkpoints is in the appendix tab. 6 §4. The paper does not state which split selects the best checkpoint.

Error rate (%) tab. 1:

Fully supervised, all labels, one row per dataset tab. 1: CIFAR-10 4.62±0.05; CIFAR-100 19.30±0.09; SVHN 2.13±0.02; STL-10 not available (its unlabeled set has no labels).

Relative error reductions over FixMatch, as the paper reports them (they match tab. 1) §1 §4.1: CIFAR-100 400 labels, 13.96%; STL-10 40 labels, 18.96%; STL-10 250 labels, 16.11%; STL-10 1000 labels, 7.68%.

ImageNet-1K, ResNet-50, 100 labels per class (100K labeled images), after 2^20 steps §4.2 tab. 2. Seeds are not stated. The authors note the models had not fully converged and that hyperparameters were not tuned.

Convergence §4.3 fig. 3: on CIFAR-100 with 400 labels, FlexMatch's loss drops faster and more smoothly. After 50K steps it already exceeds FixMatch's final accuracy §4.3. On CIFAR-10 with 40 labels, at 200K steps FixMatch reaches 56.35% overall accuracy and FlexMatch 94.29% §4.3 fig. 3.

Other results: the same CPL rule applied to Pseudo-Labeling and UDA reduces error in most cells of tab. 1 §4.1. On STL-10 with 40 labels, UDA goes from 37.42% to 29.53% with CPL tab. 1. Per-iteration running time is essentially unchanged §4.1 fig. 2. A class-balancing loss added to FixMatch gives 7.16% error on CIFAR-10 with 40 labels, against 4.97% for FlexMatch §4.4.

Ablations §4.4 fig. 4: on CIFAR-10 with 40 labels, the authors find the best τ around 0.95. The convex mapping gives the best error and the concave mapping the worst. The text gives warm-up gains of about 0.2 points on CIFAR-10 and about 1 point on CIFAR-100 §4.4 fig. 4.

Paper vs code

Code: the authors' TorchSSL repository at f26e1d42967cec7f7c8a00c2e7ff9219d8ab7c92 (MIT) LICENSE. It is the repository linked in the abstract. The FlexMatch model and training loop are in models/flexmatch/ and flexmatch.py.

Consistent with the paper:

• Threshold form, class-wise scaling and the mapping M(x) = x/(2−x) flexmatch_utils.py eq. 12.
• Learning-effect counting from stored marks, with warm-up when the flag is on: unused samples enter the denominator flexmatch.py eq. 11.
• Cosine learning rate with cycle 7/16 over K steps train_utils.py §4.
• Masked cross-entropy averaged over the unlabeled batch flexmatch_utils.py eq. 8.
• Evaluation with EMA weights, EMA momentum 0.999 flexmatch.py.

Differences:

• ImageNet: the paper's appendix gives μ = 1, τ = 0.7 and batch 128 for ImageNet app. A.1 tab. 3 tab. 4. The shipped ImageNet config uses μ = 7, τ = 0.95 and batch 32 flexmatch_imagenet_100000_0.yaml flexmatch_imagenet_100000_0.yaml flexmatch_imagenet_100000_0.yaml. It is not stated which settings produced Table 2 tab. 2.
• ImageNet steps: Table 2 is reported after 2^20 steps tab. 2; the ImageNet config runs 2,000,000 steps flexmatch_imagenet_100000_0.yaml.
• STL-10 backbone: the paper says WRN-37-2 app. A.1 tab. 4; the STL-10 config uses depth 28 with width 2 flexmatch_stl10_40_0.yaml. The model builder defaults to depth 28 wrn_var.py.
• Cut-off comparison: the paper uses a strict comparison ("greater than") in Eq. 8 and in the algorithm eq. 8 alg. 1. The code uses "greater than or equal to" for both the masking and the marking flexmatch_utils.py flexmatch_utils.py.
• Mapping in the algorithm: the algorithm computes thresholds with the plain β (Eq. 7), without M alg. 1. The text and the code use M §3.3 flexmatch_utils.py.
• Pseudo label type: the code's default is hard labels flexmatch_cifar10_40_0.yaml. The paper does not say which one it used for FlexMatch.
• Optimizer: the code's default is Nesterov momentum train_utils.py, which the paper does not mention. The paper describes plain SGD with momentum 0.9 and cites a reference for it §4.
• Weight decay on BN and bias parameters: the code excludes them by default train_utils.py. The paper does not mention this.
• BN controller: the appendix describes a BN controller for separate labeled and unlabeled forward passes app. B.1. The FlexMatch model creates one flexmatch.py, but the training step runs labeled, weak and strong views as one concatenated batch and does not call it flexmatch.py.
• Weak augmentation and RandAugment parameters are not stated in the paper; the code uses flip and crop, and RandAugment with N = 3, M = 5 ssl_dataset.py ssl_dataset.py.

Searched: the paper's abstract and text link the TorchSSL repository, so no further search was needed.

Pitfalls

• Marks are persistent state, not a per-batch count. σ_t(c) counts the samples whose last above-τ prediction was c. A sample below τ keeps its old mark, so counts can go down when a sample changes class §3.1 alg. 1. Store one integer per unlabeled sample, initialised to -1. The paper calls this marking "bonus" work done during the loss computation §3.1.
• Thresholds come from the state before this batch. The algorithm computes thresholds before it updates the marks alg. 1; the code does the same flexmatch.py flexmatch.py. Reversing the order changes which thresholds the batch is masked with.
• Mapping M is missing from the algorithm. Eq. 7 has no M, but the text's FlexMatch uses M eq. 12 §3.3. Implement M, and check against the code.
• Strict versus non-strict cut-off. The paper writes a strict comparison; the code uses ≥ flexmatch_utils.py. The two differ only for samples exactly at a threshold.
• Pseudo label type is open. Eq. 8 writes the pseudo label as a hat-q with no definition eq. 8. The background says hard and soft labels are both used by others §2. The code's default is hard flexmatch_cifar10_40_0.yaml. Argmax is a reasonable default, and it is what the code does.
• Warm-up is load-bearing early in training. Without it, the denominator is small and thresholds fluctuate strongly, and many classes sit near τ and filter out most data §3.2. The code's default is on flexmatch.py flexmatch.py.
• Class imbalance can hurt. On SVHN, FlexMatch is worse than FixMatch in tab. 1 (8.19% against 3.81% with 40 labels). The authors attribute this to a simple but imbalanced dataset, where rare classes never reach β close to 1 and keep low thresholds §4.1. Expect the same on imbalanced data.
• Not uniformly best. At 250 labels on CIFAR-10, FixMatch (4.86%) beats FlexMatch (4.98%) in tab. 1. The paper claims state of the art on most benchmarks, which the table supports.
• STL-10 numbers in the introduction are mislabeled. The introduction gives the STL-10 relative gains for 400, 2500 and 10000 labels §1. The percentages match STL-10 with 40, 250 and 1000 labels in tab. 1, so the label counts in the text are wrong.
• Best-checkpoint numbers are optimistic. The headline metric is the best error across checkpoints §4. The code picks the best checkpoint by evaluation accuracy on the evaluation loader flexmatch.py; the paper does not state which split that is. Use the median of the last 20 checkpoints when comparing methods that converge at different speeds; the authors say that median is unfair when convergence speeds differ and report it only in the appendix §4 tab. 6.
• Long training overfits. FlexMatch's accuracy drops after about 800K steps, and FixMatch's after about 900K §4.3. Cap or checkpoint accordingly.
• Cost claim is about the GPU. The paper says CPL adds no forward or backward passes and that per-iteration time is unchanged §3.1 fig. 2. In the code, the step rebuilds a Python Counter over all unlabeled marks every iteration flexmatch.py. That is CPU work linear in the unlabeled set size, which the paper does not measure.
• Start of training. With all marks at -1, every count is zero. Algorithm 1 then takes the warm-up branch, so the normalisation never divides by zero, and all thresholds are zero alg. 1 eq. 11. The code keeps its initial zero thresholds until some class gets a mark flexmatch.py; almost every weak-view sample passes the mask until then.
• Equation numbers are counted from the LaTeX source. This card cites equations by their order in the source. Check them against the compiled PDF before quoting in print.

Lineage

• Builds on: FixMatch, the base algorithm whose fixed threshold is replaced §1 §2 §3.1. The weak/strong split and the masked consistency loss are FixMatch's own design, as the paper describes it §2.
• Applied to other methods: CPL is also applied to Pseudo-Labeling (Flex-PL) and UDA (Flex-UDA), which the paper names as methods that use a fixed threshold §1 §4.1 tab. 1.
• Compared against: Pseudo-Labeling, UDA and FixMatch in the main table §4 tab. 1; the appendix also reports Π-Model, VAT, Mean Teacher, MixMatch and ReMixMatch, which are not reproduced here app. B.2.
• Alternative considered: a class-balancing objective added to FixMatch, reported in §4.4.
• Related work §5 covers curriculum learning and dynamic thresholds in other fields; those works are not covered by a card here.