Source-linked AI summary
Adaptive Wing Loss for Robust Face Alignment via Heatmap Regression
Xinyao Wang, Liefeng Bo, Li Fuxin
TL;DR
The paper addresses underexplored loss design and foreground-background imbalance in heatmap regression for difficult face alignment cases. It introduces Adaptive Wing loss with Weighted Loss Map, boundary prediction, and CoordConv, and reports significant gains across multiple benchmarks and metrics.
Problem
Heatmap-regression losses were rarely studied, while MSE can underemphasize small foreground errors and equal pixel weighting allows background pixels to dominate.
Method
The approach combines Adaptive Wing loss, Weighted Loss Map, jointly trained boundary prediction, and CoordConv with boundary coordinates.
Results
The method outperforms state-of-the-art approaches on multiple datasets and metrics, reducing COFW failure rate at 10% NME from 3.73% to 0.99%.
Takeaways & Limitations
Adaptive Wing loss is reported to improve heatmap regression for face alignment and another human pose estimation task.
Takeaways & Limitations
The original Wing loss is not applicable to heatmap regression because its gradient is discontinuous at zero and overly emphasizes small background errors.
Abstract
from arXiv · showhide
Heatmap regression with a deep network has become one of the mainstream approaches to localize facial landmarks. However, the loss function for heatmap regression is rarely studied. In this paper, we analyze the ideal loss function properties for heatmap regression in face alignment problems. Then we propose a novel loss function, named Adaptive Wing loss, that is able to adapt its shape to different types of ground truth heatmap pixels. This adaptability penalizes loss more on foreground pixels while less on background pixels. To address the imbalance between foreground and background pixels, we also propose Weighted Loss Map, which assigns high weights on foreground and difficult background pixels to help training process focus more on pixels that are crucial to landmark localization. To further improve face alignment accuracy, we introduce boundary prediction and CoordConv with boundary coordinates. Extensive experiments on different benchmarks, including COFW, 300W and WFLW, show our approach outperforms the state-of-the-art by a significant margin on various evaluation metrics. Besides, the Adaptive Wing loss also helps other heatmap regression tasks. Code will be made publicly available at https://github.com/protossw512/AdaptiveWingLoss.
1. Introduction
The paper identifies weaknesses in MSE-based heatmap regression and proposes adaptive pixel weighting, boundary prediction, and coordinate encoding for more accurate face alignment.
- Motivation: MSE is insufficiently sensitive to small foreground errors, while equal weighting lets background pixels dominate heatmap training.These issues can produce blurry, dilated heatmaps and inaccurate landmark locations.
- Contributions: The model encodes full and boundary-specific coordinate information with CoordConv and jointly predicts facial boundaries.Boundary prediction adds an auxiliary channel to the heatmap representation.
- Contributions: Adaptive Wing loss adapts its curvature to ground-truth pixel values, emphasizing foreground accuracy while tolerating small background errors.Weighted Loss Map further focuses training on foreground and difficult background pixels.
- Results: The approach outperforms prior state-of-the-art methods on COFW, 300W, and WFLW, with particular gains in failure rate and AUC.The authors also report benefits for another heatmap regression task.
2. Related Work
Prior work established strong CNN architectures and coordinate or boundary enhancements, but heatmap-regression loss functions remained comparatively underexplored.
- CNN heatmap regression: Stacked Hourglass and related multi-scale CNN architectures achieved state-of-the-art face-alignment performance.Other approaches include quantized densely connected U-Nets, fully convolutional networks, and regression-tree hybrids.
- Loss functions: Heatmap-regression loss functions were rarely studied compared with network architectures.GoDP introduced distance-aware penalties, while Wing loss targeted direct landmark-coordinate regression.
- Loss functions: Original Wing loss is unsuitable for heatmap regression because it overemphasizes small background errors and has a discontinuous gradient at zero.Adaptive Wing loss is designed to retain foreground sensitivity while tolerating background errors.
- Boundary and coordinate information: Boundary information and CoordConv provide complementary ways to incorporate facial structure and coordinate information into alignment models.LAB previously used boundary maps in a two-stage design, while CoordConv adds coordinate channels before convolution.
3. Our Model
The model builds on stacked Hourglass heatmap regression, adds boundary prediction as an auxiliary task, and injects coordinate encodings throughout the network.
- Heatmap prediction: Each stacked Hourglass module is supervised using the ground-truth heatmap for landmark prediction.The architecture follows a stacked HG design derived from prior work.
- Boundary prediction: Boundary prediction is added as an additional heatmap channel and trained jointly with landmark channels.This makes facial boundary estimation an auxiliary sub-task of the model.
- Coordinate encoding: Coordinate encoding is inserted before the first convolution and before the first convolution block of every Hourglass module.The complete model is summarized in Figure 3.
4. Adaptive Wing Loss for Face Alignment
The section analyzes why fixed losses struggle in heatmap regression and introduces Adaptive Wing loss with Weighted Loss Map to emphasize pixels most important for landmark localization.
- Loss function rationale: MSE gives small-error pixels weak influence, producing blurry heatmaps and low-intensity foreground predictions near difficult landmarks.L1 avoids weak small-error influence but can destabilize training because its gradient is discontinuous at zero.
- Loss function rationale: Wing loss amplifies small errors but its large gradient around zero gives background pixels disproportionate influence, preventing heatmap-regression training from converging.The paper reports that Wing loss works for coordinate regression but is not applicable to heatmap regression.
- Adaptive Wing loss: Adaptive Wing loss adapts its curvature to ground-truth intensity, increasing small-error influence near Gaussian modes while behaving more like MSE for low-intensity background pixels.Its threshold θ switches between nonlinear and linear regions, with continuity and smoothness enforced at |y − ŷ| = θ.
- Adaptive Wing loss: The method uses α = 2.1, ω = 14, ϵ = 1, θ = 0.5, and W = 10 in its experiments.These settings were selected empirically, with parameter ablations reported separately.
- Weighted loss map: Weighted Loss Map addresses foreground-background imbalance by assigning higher weights to foreground and difficult background pixels during training.For a typical 64 × 64 heatmap with a 7 × 7 Gaussian, foreground pixels comprise only 1.2% of pixels; the mask is produced using 3 × 3 dilation and thresholding.
5. Boundary Information
The model adds boundary prediction as an auxiliary heatmap channel to capture global human-face information and feed boundary cues into the network.
- 5. Boundary Information: Boundary prediction uses one additional channel combining all boundary lines rather than splitting boundaries into separate parts.The boundary information is aggregated through convolution during the forward pass and is also used subsequently in the model.
6. Coordinate aggregation
The model integrates CoordConv with boundary-derived coordinates so convolutional layers can use explicit spatial information for face alignment.
- 6. Coordinate aggregation: Boundary coordinate encoding augments standard coordinate channels with X and Y coordinates generated only at predicted boundaries.The coordinate channels are created at runtime and concatenated with the original input before regular convolution.
7. Experiments
The experiments evaluate the approach across COFW, 300W, and WFLW using NME, FR, and AUC-related protocols, with standardized cropping, training, augmentation, and inference settings. On COFW, the method substantially improves failure rate and NME, indicating robustness to pose and occlusion.
- Datasets: The approach is evaluated on COFW, 300W, 300W private test, and WFLW datasets.WFLW is described as the most difficult dataset among those evaluated.
- Evaluation metrics: NME measures normalized landmark-coordinate error, while FR marks predictions exceeding a dataset-specific NME threshold.The experiments also use CED curves and AUC, where larger AUC indicates a larger well-predicted portion of the test set.
- Preprocessing: Faces are cropped from dataset bounding boxes, with WFLW boxes enlarged by 10% to preserve landmarks and ground-truth cropping used for 300W private test.The longer side of each bounding box defines the square crop, except for the private-test setting without official boxes.
- Training and inference: The network uses 256 × 256 inputs, four stacked Hourglass modules, 240 training epochs, and augmentation with rotation, translation, flipping, rescaling, blur, noise, and occlusion.Inference shifts the maximum-response location toward the second-highest nearby pixel, and boundary lines are generated from landmarks.
- COFW results: 0.99% is the method’s failure rate at 10% NME on COFW, reduced from 3.73%, while its NME is reported as better than the 5.60% human reference.The COFW result is presented as evidence of robustness against large pose and heavy occlusion.
7.4. Evaluation on 300W
On 300W and its private test set, the method reaches state-of-the-art performance across reported metrics, including NME, AUC, and failure rate. The supplied passages also report best performance on WFLW, including every subset, with visualizations provided for the test set.
- 300W testing dataset: The method achieves state-of-the-art performance on the 300W testing dataset.The result is reported in Table 3.
- 300W private test: On the 300W private test dataset, the method outperforms previous state-of-the-art results across NME, AUC, and FR at both 8% and 10% NME thresholds.The failure rate at the 8% threshold is more than halved versus the next-best baseline, reaching 0.83%.
- WFLW: On WFLW, the method outperforms previous state-of-the-art approaches by a significant margin on every subset.The WFLW dataset is described as significantly more difficult than COFW and 300W.
- WFLW: On WFLW, the method reduces failure rate and increases AUC, while failing on only 2.84% of images.The passage characterizes this as more than a two-times improvement compared with the referenced baseline.
- WFLW: Figure 7 provides visualizations on the WFLW test dataset.The supplied figure passage identifies the dataset but does not state a visual comparison outcome.
7.6. Ablation study
The ablations show that Adaptive Wing loss contributes most to localization gains, while Weighted Loss Map, boundary prediction, and coordinate encoding add further improvements. Additional experiments examine parameter sensitivity and transfer to human pose heatmap regression.
- Adaptive Wing parameters: The parameter search selects ω = 14 and ϵ = 1 for Adaptive Wing loss after evaluating combinations on WFLW.The experiments report that performance is not very sensitive to ω and ϵ, although the search space was reduced because resources were limited.
- Module ablation: Adaptive Wing loss provides the major ablation gain, improving the WFLW benchmark by 0.74%.The baseline model is trained with MSE and underperforms the state-of-the-art.
- Module ablation: Weighted Loss Map improves localization by 0.25% and outperforms a naive weight mask by focusing on difficult background pixels.The naive baseline uses WM_base = ˆH_W + 1 with W = 10.
- Module ablation: Boundary prediction and coordinate encoding contribute another 0.09% to localization performance.The evaluated coordinate variant includes CoordConv with boundary coordinates.
- Human pose estimation: On LSP human pose estimation, Adaptive Wing loss significantly boosts performance compared with MSE.The experiment uses original labels and low-resolution images, with PCK and torso dimension as the normalization factor, and targets applicability rather than state-of-the-art performance.
8. Conclusion
The paper addresses two MSE limitations in heatmap regression with Adaptive Wing loss and Weighted Loss Map, then adds boundary prediction and boundary-aware CoordConv. Experiments report significant improvements across multiple datasets and metrics, particularly for difficult scenarios.
- 8. Conclusion: Adaptive Wing loss and Weighted Loss Map target accurate facial-landmark localization by addressing two identified MSE limitations.The proposed loss and weighting focus training on foreground and difficult background pixels.
- 8. Conclusion: Boundary prediction and CoordConv with boundary coordinates are introduced to further improve localization results.
- 8. Conclusion: The approach outperforms state-of-the-art methods on multiple datasets and evaluation metrics, especially failure rate and AUC.The reported pattern indicates greater robustness to difficult scenarios.
10. Supplementary Material
The supplementary material details boundary-aware coordinate encoding, datasets, evaluation, loss behavior, noise experiments, model-size comparisons, and visualization results. It reports faster Adaptive Wing convergence, robustness across challenging datasets, and strong performance with fewer Hourglass stacks.
- Boundary-aware CoordConv: Boundary coordinates are generated by masking X and Y coordinate channels with thresholded boundary predictions from the previous Hourglass module.The boundary mask uses a threshold of 0.05.
- Datasets Used in Our Experiments: WFLW contains 7,500 training images, 2,500 testing images, and 98 landmarks, with attributes covering pose, expression, illumination, makeup, occlusion, and blur.The dataset supports separate analysis across six subsets with different properties.
- AFLW Results: On AFLW, the method surpasses state-of-the-art methods despite using limited boundary information from 19 landmarks.The AFLW experiment targets faces with large poses and uses boundaries for selected facial lines.
- Adaptive Wing Loss Analysis: Adaptive Wing loss reduces pixel-wise average MSE by almost 30% and foreground-pixel MSE by more than 23% within 50 epochs.The comparison evaluates MSE for models trained with original stacked Hourglass architecture without a weight map.
- Number of Hourglass Stacks: With only one Hourglass block, the approach remains competitive across datasets and runs at 120 FPS on an Nvidia GTX 1080Ti.It outperforms previous state-of-the-art results except on the common and full subsets of 300W.
- Result Visualization: Figures visualize localization results for AFLW, COFW, 300W, 300W private, and WFLW datasets.