Source-linked AI summary
Puzzle Mix: Exploiting Saliency and Local Statistics for Optimal Mixup
Jang-Hyun Kim, Wonho Choo, Hyun Oh Song
TL;DR
Existing mixup methods can provide misleading supervisory signals because they create virtual examples while disregarding saliency and local data statistics. Puzzle Mix jointly optimizes saliency-aware masks and transport plans, and experiments report improved generalization and adversarial robustness across CIFAR-100, Tiny-ImageNet, and ImageNet.
Problem
Mixup methods may create misleading supervisory signals by disregarding regional saliency and local statistics, while deep networks remain vulnerable to overfitting and adversarial attacks.
Method
Puzzle Mix jointly optimizes a multi-label mixing mask and saliency-aware transport plans while respecting local statistics, with adversarial training incorporated without computation overhead.
Results
Puzzle Mix outperforms other mixup baselines in generalization and robustness against adversarial perturbations and data corruption on CIFAR-100, Tiny-ImageNet, and ImageNet.
Takeaways & Limitations
Saliency-aware masking and transport are both supported as contributors to Puzzle Mix's generalization performance.
Takeaways & Limitations
Puzzle Mix requires an additional forward-backward network evaluation per mini-batch to calculate the saliency signal.
Abstract
from arXiv · showhide
While deep neural networks achieve great performance on fitting the training distribution, the learned networks are prone to overfitting and are susceptible to adversarial attacks. In this regard, a number of mixup based augmentation methods have been recently proposed. However, these approaches mainly focus on creating previously unseen virtual examples and can sometimes provide misleading supervisory signal to the network. To this end, we propose Puzzle Mix, a mixup method for explicitly utilizing the saliency information and the underlying statistics of the natural examples. This leads to an interesting optimization problem alternating between the multi-label objective for optimal mixing mask and saliency discounted optimal transport objective. Our experiments show Puzzle Mix achieves the state of the art generalization and the adversarial robustness results compared to other mixup methods on CIFAR-100, Tiny-ImageNet, and ImageNet datasets. The source code is available at https://github.com/snu-mllab/PuzzleMix.
1. Introduction
Puzzle Mix addresses overfitting and distribution-shift vulnerability by incorporating saliency and local statistics into mixup, jointly optimizing masks and transport plans. It also supports adversarial training and reports improved generalization and robustness across three image datasets.
- Deep networks can memorize training data, produce overconfident predictions, and suffer degraded generalization under slight distribution shifts.
- Ignoring regional saliency and local regularity can create mixup examples that misguide training and undermine generalization.
- Puzzle Mix jointly optimizes a regional mixing mask and transport plan to expose saliency while leveraging local statistics.
- Puzzle Mix incorporates adversarial training without additional computation overhead by stochastically including adversarially perturbed examples with random restarts.
- Puzzle Mix significantly improves generalization and adversarial robustness over existing mixup methods on CIFAR-100, Tiny-ImageNet, and ImageNet.
2. Related Works
Related work covers conventional data augmentation, mixup variants, saliency estimation, and optimal transport. These strands motivate Puzzle Mix's use of structured saliency-aware mixing.
- Data augmentation: Data augmentation regularizes models against overfitting by generating virtual training examples near the training dataset.
- Mixup: Input mixup linearly interpolates inputs and one-hot labels, while manifold mixup interpolates hidden representations and CutMix replaces a random rectangular region.
- Saliency: Saliency methods estimate informative regions using gradients, with alternatives trading model modification or saliency resolution against computation cost.
- Optimal transport: Optimal transport moves one distribution to another at minimal cost and can be formulated over discrete domains as a linear program or assignment problem.
3. Preliminaries
The preliminaries define mixup as optimizing model loss through a data-mixing function and mixing distribution. They summarize input, manifold, and CutMix formulations alongside Puzzle Mix's mask and transport notation.
- Mixup-based augmentation optimizes the model loss using a data mixup function h and mixing distribution q over inputs and labels.
- Label mixup uses g(y0, y1) = (1 − λ)y0 + λy1, while input and manifold mixup interpolate inputs or hidden representations.
- CutMix uses a binary rectangular mask 1B to replace a region covering λ proportion of one input with another.
- Puzzle Mix represents z_i as mask values with mixing ratio λ = 1/i z_i and uses transportation plans Π0 and Π1 to move data locations.
- Table 1 summarizes the mixup functions introduced in the preliminaries.
4. Methods
Puzzle Mix jointly optimizes a regional mixing mask and transport plans to preserve salient information while respecting local data structure. The method alternates these optimizations to generate mixed examples efficiently, with graph-based optimization and convergence guarantees for key subproblems.
- Mask and transport objectives: Puzzle Mix optimizes a mask that reveals more salient data from either transported input in each region.
- Mask and transport objectives: The transport plans move input regions to maximize the saliency of the portions revealed by the mask.
- Local statistics: Spatial regularization controls mask and output smoothness, while transport costs and binary plans preserve local structural integrity.
- Mask optimization: The mask objective combines regional saliency, label and data smoothness, and a mixing-weight prior in a multi-label energy minimization problem.
- Optimization procedure: The mask problem is solved with α-β swap graph cuts under submodularity conditions, while the transport update uses iterative column conflict resolution.
- Optimization procedure: Algorithm 1 converges to a local minimum in at most n(n−1)/2+1 steps.
- Adversarial training: Puzzle Mix can incorporate adversarial training without additional computation cost because saliency computation already uses loss gradients with respect to the inputs.
5. Implementation Details
Implementation uses down-sampled, randomly sized mask graphs and mini-batch computation, while retaining gradient-based saliency for adversarial training. An additional clean-gradient regularizer improves generalization on Tiny-ImageNet and ImageNet.
- Optimization and sampling: The mask graph is down-sampled and randomly sampled from sizes 2 × 2, 4 × 4, 8 × 8, and 16 × 16.
- Optimization and sampling: Saliency maps are normalized, regional distances use boundary differences, and λ is sampled from Beta(α, α) per mini-batch.
- Adversarial training: The method reuses backpropagated loss gradients for saliency without computational overhead and regularizes mixup gradients with clean-example gradients.
- Adversarial training: The additional gradient regularization improves generalization performance on Tiny-ImageNet and ImageNet.
6. Experiments
Experiments across CIFAR-100, Tiny-ImageNet, ImageNet, and corruption benchmarks show Puzzle Mix improves generalization and robustness over mixup baselines, including under matched evaluation budgets. Ablations associate gains with multi-label masking, transport, and stochastic adversarial-training choices.
- CIFAR-100: Puzzle Mix outperforms mixup baselines in generalization and adversarial robustness on CIFAR-100 with both WRN28-10 and PreActResNet18.With WRN28-10, it improves Top-1 test error over the best baseline by 1.45%, while Puzzle Mix (adv) lowers FGSM error by 8.41% versus AugMix.
- Tiny-ImageNet: 36.52% Top-1 test error is achieved on Tiny-ImageNet with equal epochs, 5.47% lower than the strongest baseline; the half-budget variant improves by 4.35%.Puzzle Mix (adv) also achieves 3.44% lower Top-1 error and 5.12% lower FGSM error than the best baseline.
- ImageNet: 0.43% and 0.24% improvements in ImageNet Top-1 and Top-5 error rates, respectively, make Puzzle Mix the best-performing method in that experiment.The comparison uses ResNet-50 under a 100-epoch training protocol.
- Robustness Against Corruption: Puzzle Mix applied to AugMix images improves generalization and corruption robustness over AugMix with the Jensen-Shannon-divergence objective.On CIFAR-100-C, the improvements over AugMix are 3.95% for generalization and 2.31% for corruption accuracy, while requiring one rather than two additional evaluations per mini-batch.
- Ablation Study: 1.44% higher Top-1 error for binary labels and 0.43% higher error without transport support multi-label masking and transport as effective components.These ablations were conducted on CIFAR-100 with WRN28-10.
- Ablation Study: 4.02% lower FGSM error without linear decay comes with a 0.41% increase in Top-1 error, while independent adversarial random variables improve robustness.The stochastic adversarial-training ablation uses p = 0.1, FGSM at 8/255, and seven-step PGD with a 2/255 step size.
7. Conclusion
Puzzle Mix leverages saliency information and local data statistics in mixup augmentation. The paper reports improved generalization and robustness against adversarial perturbations and data corruption across multiple datasets.
- Puzzle Mix is presented as a mixup augmentation method that optimally leverages saliency information while respecting local data statistics.
- The method efficiently generates mixup examples in a mini-batch stochastic gradient descent setting.
- Puzzle Mix outperforms other mixup baselines in generalization and robustness against adversarial perturbations and data corruption on CIFAR-100, Tiny-ImageNet, and ImageNet.
- Lemma 1 states that a nonnegative weighted sum of pairwise-submodular functions remains pairwise submodular.
A.2. Proof of Proposition 5
The proof establishes that Algorithm 1 reaches a local minimum under its update rule within a bounded number of steps. The bound is obtained by inductively tracking permanently fixed assignments.
- Algorithm 1 converges to a local minimum with respect to its update rule in at most n(n−1)/2+1 steps.
- The proof inductively selects minimum-cost assignments while excluding rows and columns already fixed by earlier updates.
- Once an assignment is fixed, subsequent permutations retain it, allowing the number of updates to be bounded by the cumulative step increments.
- The induction concludes that no further updates occur after n(n−1)/2+1 steps.
B.1. Comparison Experiments for Algorithm 1
Algorithm 1 is evaluated against the exact Hungarian algorithm for transport optimization and against alternative training configurations. The experiments also examine one-cycle optimization and training protocols across datasets.
- At n = 1024, Algorithm 1 is about 8.6 times faster than the exact algorithm in simulations using random cost matrices.
- Figure 6 compares average execution time on a log-log scale between the exact solver and Algorithm 1 using means over 100 trials.
- Figure 7 compares relative objective-function errors for Algorithm 1 and random assignment using errors relative to the exact solution.
- Across ten random seeds, the mean Top-1 accuracy difference between the Hungarian algorithm and Algorithm 1 is −0.025, with standard deviation 0.239 and paired-test P-value 0.919.
- Most optimal masks remain unchanged after the first optimization cycle when optimizing 1,000 CIFAR100 image pairs with varying region counts.
- Puzzle Mix uses alternating mask and transport optimization and is evaluated with specified SGD schedules, adversarial-training settings, and dataset-specific protocols.
C.4. Hyperparameter Sensitivity
Hyperparameter sensitivity is examined on CIFAR-100 with WRN28-10 by sweeping parameters individually. The reported main setting outperforms the cited Manifold Mixup baseline, while adversarial training exposes a clean-accuracy trade-off.
- The main setting (β, γ, η, ξ) = (1.2, 0.5, 0.2, 0.8) achieves 15.95% Top-1 test error on CIFAR-100.
- Hyperparameters are swept individually while the others remain fixed, and mean Top-1 error rates with standard deviations are reported.
- A well of hyperparameters performs better than the Manifold Mixup baseline of 17.40%.
- Adversarial training introduces an inverse relationship between clean error and FGSM error through the adversarial probability parameter.
D.2. Robustness against PGD Attack
Puzzle Mix is evaluated for adversarial robustness against PGD on CIFAR-100, with comparisons against existing mixup methods and adversarial-training settings. The reported results show greater robustness than existing mixup methods.
- Experimental setup: The experiment trains PreActResNet18 on CIFAR-100 and evaluates mixup methods under a PGD attack.The attack uses a 4/255 l∞ epsilon-ball and a 2/255 step size.
- Robustness result: Puzzle Mix is more robust against the PGD attack than existing mixup methods.
- Robustness result: Figure 8 presents adversarial robustness results for various mixup methods.
E.2. More Samples
This section examines Puzzle Mix across different resolutions for its optimal mask and transport. Figure 11 visualizes the resulting Puzzle Mix examples alongside their inputs.
- More Samples: The experiments evaluate Puzzle Mix with various resolutions of the optimal mask and transport.
- More Samples: Figure 11 visualizes Puzzle Mix results together with the given inputs.
- More Samples: The visualization compares Puzzle Mix outputs against their corresponding input examples.