CAAL: Confidence-Aware Active Learning for Heteroscedastic Atmospheric Regression

Quantifying the impacts of air pollution on health and climate relies on key atmospheric particle properties such as toxicity and hygroscopicity. However, these properties typically require complex observational techniques or expensive particle-resolved numerical simulations, limiting the availability of labeled data. We therefore estimate these hard-to-measure particle properties from routinely available observations (e.g., air pollutant concentrations and meteorological conditions). Because routine observations only indirectly reflect particle composition and structure, the mapping from routine observations to particle properties is noisy and input-dependent, yielding a heteroscedastic regression setting. With a limited and costly labeling budget, the central challenge is to select which samples to measure or simulate. While active learning is a natural approach, most acquisition strategies rely on predictive uncertainty. Under heteroscedastic noise, this signal conflates reducible epistemic uncertainty with irreducible aleatoric uncertainty, causing limited budgets to be wasted in noise-dominated regions. To address this challenge, we propose a confidence-aware active learning framework (CAAL) for efficient and robust sample selection in heteroscedastic settings. CAAL consists of two components: a decoupled uncertainty-aware training objective that separately optimises the predictive mean and noise level to stabilise uncertainty estimation, and a confidence-aware acquisition function that dynamically weights epistemic uncertainty using predicted aleatoric uncertainty as a reliability signal. Experiments on particle-resolved numerical simulations and real atmospheric observations show that CAAL consistently outperforms standard AL baselines. The proposed framework provides a practical and general solution for the efficient expansion of high-cost atmospheric particle property databases.

💡 Research Summary

The paper tackles a pressing problem in atmospheric science: estimating hard‑to‑measure particle properties such as toxicity and hygroscopicity from cheap, routinely collected observations (e.g., pollutant concentrations, meteorological variables). These target properties are intrinsically noisy and the noise level varies with the input, making the task a heteroscedastic regression problem. Because obtaining ground‑truth labels requires expensive laboratory analyses or particle‑resolved numerical simulations, only a few hundred to a few thousand labeled examples are available, far fewer than the billions of low‑cost observations. This scarcity motivates the use of active learning (AL) to select the most informative samples under a limited labeling budget.

Standard AL strategies rely on predictive uncertainty, but in heteroscedastic settings total uncertainty conflates epistemic (reducible) and aleatoric (irreducible) components. Selecting samples solely on total uncertainty wastes budget on regions dominated by aleatoric noise, where additional labels cannot improve the model. Existing remedies either discard high‑noise regions entirely or ignore noise altogether, both of which are sub‑optimal. Moreover, training a heteroscedastic model with the usual Gaussian negative log‑likelihood (NLL) couples mean and variance learning; the model can reduce loss by inflating predicted variance on difficult points, which weakens the learning signal for the mean, especially when data are scarce.

The authors propose CAAL (Confidence‑Aware Active Learning), a framework designed specifically for heteroscedastic regression under active‑learning constraints. CAAL has two complementary components:

1. Mean‑variance decoupled training – Each ensemble member is trained with two separate losses. The predictive mean is optimized with mean‑squared error (MSE), while the variance head is calibrated using an NLL‑style loss that receives a stop‑gradient on the residual term. This prevents variance gradients from contaminating the mean parameters, stabilising mean learning even when many samples have high noise. A weighting control mechanism regulates how much the variance head influences shared feature layers, avoiding the destabilising effect of variance‑only updates.

2. Confidence‑aware acquisition function – Using deep ensembles, the model obtains an epistemic uncertainty estimate (variance of ensemble means) and an aleatoric uncertainty estimate (average of ensemble variances). The acquisition score multiplies the epistemic term by a confidence weight derived from the predicted aleatoric variance: low aleatoric variance (high confidence) leaves the epistemic term largely unchanged, while high aleatoric variance (low confidence) down‑weights the score. Consequently, the algorithm preferentially queries points where the model is both uncertain and reliable, i.e., where additional labels are expected to reduce epistemic uncertainty.

The methodology is evaluated on two fronts: (a) synthetic particle‑resolved numerical simulations that exhibit strong input‑dependent noise, and (b) real‑world atmospheric datasets with measured toxicity and hygroscopicity. Experiments follow a pool‑based AL protocol: an initial small labeled set, batch queries of size 5–10, and a total budget of a few hundred labels. Baselines include classic uncertainty‑based methods (BALD, MaxEntropy), diversity‑based methods (BADGE, Coreset), and recent heteroscedastic‑aware strategies that either avoid noisy regions or focus solely on epistemic uncertainty.

Results show that CAAL consistently outperforms all baselines. On the primary simulation dataset, CAAL improves R² by 9.6 % while using 45.6 % fewer labels to reach the same performance as the best baseline. In highly heteroscedastic regions, standard methods waste queries on noise‑dominated points, causing performance plateaus; CAAL’s confidence weighting prevents this waste. Ablation studies confirm that (i) the decoupled loss is essential—training with joint NLL leads to unstable mean predictions—and (ii) the confidence weighting yields a measurable gain over using raw epistemic uncertainty alone. Calibration analysis also indicates that the predicted aleatoric variances correlate well with empirical residuals, validating their use as a reliability signal.

In summary, CAAL offers a principled solution to the dual challenges of heteroscedastic noise and limited labeling budgets. By separating mean and variance learning and by modulating acquisition with a data‑driven confidence term, it achieves more efficient label utilization and higher predictive accuracy. The framework is generic and can be transferred to other scientific domains where expensive, noisy labels must be inferred from abundant cheap measurements (e.g., material property prediction, biomedical imaging). Future work may explore richer Bayesian approximations, multi‑task extensions, and integration with cost‑sensitive labeling strategies.