Uncertainty-guided EMA (UEMA)

• UEMA is a principled extension of EMA that integrates loss variance as an uncertainty proxy to fine-tune task weighting in multi-task learning.
• It builds on EMA-based methods like REMA and DWEMA to mitigate negative transfer by normalizing task losses and stabilizing noisy gradients.
• By separating mean loss from volatility, UEMA enables adaptive weighting strategies that improve model performance and convergence.

Searching arXiv for papers related to uncertainty-guided EMA and the cited EMA literature. Uncertainty-guided Exponential Moving Average (UEMA) denotes a family of EMA-based schemes in which the exponential moving average is used not merely as a smoother, but as a mechanism for weighting or averaging according to an uncertainty signal. In the multi-task learning setting, UEMA can be understood as an extension of the EMA-based loss weighting strategies introduced in "Mitigating Negative Transfer in Multi-Task Learning with Exponential Moving Average Loss Weighting Strategies" (Lakkapragada et al., 2022): per-task losses are normalized by their exponential moving averages, and uncertainty is then introduced through additional EMA-tracked statistics such as loss variance. Across the cited literature, EMA appears in several technically distinct roles—as a loss normalizer in multi-task optimization (Lakkapragada et al., 2022), as a parameter-averaging operator with explicit bias–variance consequences in online SGD (Li et al., 19 Feb 2025), as a higher-moment tracker in risk modeling (Gabrielsen et al., 2012), and as a stabilizing coefficient prior in ODE-inspired Transformers (Li et al., 2024). UEMA is therefore best viewed not as a single canonical algorithm, but as a principled uncertainty-aware generalization of EMA.

1. EMA-based foundations in multi-task learning

The most direct origin of UEMA lies in multi-task learning (MTL), where the joint objective is written as

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^K \lambda_k(t)\, L_k(t),$$

with task-specific losses $L_k(t)$ and weights $\lambda_k(t)$. In naïve MTL, fixed or equal weights can cause certain tasks to dominate training, producing negative transfer: some tasks perform better in single-task models than in the multi-task model. The EMA-based remedy is to track the magnitude of each task loss online through

$$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$

and then to define

$\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$

The resulting objective is

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$

This construction does not enforce $\sum_k \lambda_k = 1$; the weights are unnormalized, and the operative principle is relative rescaling rather than simplex normalization (Lakkapragada et al., 2022).

The same work introduces two extensions of vanilla EMA. REMA incorporates the training rate

$r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$

and uses

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} r_k(t)\,\frac{L_k(t)}{\tilde{L}_k(t)}.$$

DWEMA combines DWA-style training-rate weighting with EMA normalization, so that losses are first brought onto a comparable magnitude scale and only then modulated by rate information. In all three cases—EMA, REMA, and DWEMA—the shared structural idea is that loss-scale imbalance should be addressed directly by EMA-based normalization before more elaborate dynamics are added.

The significance of this construction for UEMA is straightforward. EMA already provides an online estimate of the scale of each task’s loss. UEMA preserves that mechanism but augments it with an explicit uncertainty proxy derived from the same loss time series. In that sense, UEMA is not a departure from EMA-based MTL weighting; it is an elaboration of its statistical state.

2. Formalization of uncertainty-guided EMA

A natural UEMA extension is to maintain not only an EMA of the loss mean, but also an EMA of the squared loss. Let

$$m_{k,t} = \beta m_{k,t-1} + (1-\beta)\mathcal{L}_{k,t},$$

$L_k(t)$0

and define the variance proxy

$L_k(t)$1

Here $L_k(t)$2 plays the role of the smoothed loss magnitude, while $L_k(t)$3 measures temporal instability of that loss. In the supplied synthesis, UEMA is then expressed through weights such as

$L_k(t)$4

or, more generally,

$L_k(t)$5

with $L_k(t)$6 and $L_k(t)$7 preventing division by zero (Lakkapragada et al., 2022).

These formulas separate two signals that plain EMA conflates. The mean $L_k(t)$8 captures persistent scale: a task with consistently large losses is downweighted to prevent domination. The standard deviation or variance term captures instability: highly variable losses are treated as uncertain. This suggests a two-axis weighting policy in which large mean and large variability need not be interpreted identically. A stable but difficult task and an erratic noisy task may both exhibit elevated loss, yet UEMA allows them to be handled differently.

This construction is explicitly presented as an extension rather than an established canonical algorithm. The cited MTL paper introduces EMA, REMA, and DWEMA, but not a standalone UEMA method. The uncertainty-guided formulation arises by extending the same EMA machinery from first-order loss statistics to second-order loss statistics (Lakkapragada et al., 2022). A related theoretical paper on EMA in SGD likewise does not explicitly mention uncertainty-guided EMA, but it does analyze general averaging schemes with time-dependent $L_k(t)$9, which provides a formal basis for adaptive EMA schedules (Li et al., 19 Feb 2025).

3. Negative transfer, task dominance, and empirical behavior

The core MTL motivation for UEMA is negative transfer induced by task dominance. When losses or gradients differ sharply in scale, equal weighting causes dominant tasks to steer shared representations disproportionately. EMA-based weighting addresses this by dividing each instantaneous loss by a smoothed estimate of its own magnitude. If a task has persistently large loss, its EMA is also large, and its effective coefficient decreases; if a task’s loss is smaller, its coefficient increases correspondingly. This rescaling is designed to keep effective task contributions on a similar scale over time (Lakkapragada et al., 2022).

The empirical results reported for EMA-based methods supply the main evidence for this perspective. On CelebA, the overall F1 values are Baseline $\lambda_k(t)$0, GradNorm $\lambda_k(t)$1, UW $\lambda_k(t)$2, DWA $\lambda_k(t)$3, EMA $\lambda_k(t)$4, DWEMA $\lambda_k(t)$5, and REMA $\lambda_k(t)$6. On AffWild2, the overall values are Baseline $\lambda_k(t)$7, GradNorm $\lambda_k(t)$8, UW $\lambda_k(t)$9, DWA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$0, REMA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$1, DWEMA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$2, and EMA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$3. On AffectNet, the overall values are Baseline $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$4, GradNorm $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$5, UW $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$6, DWA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$7, REMA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$8, DWEMA $$\tilde{L}_k(t) = \beta L_k(t) + (1-\beta)\tilde{L}_k(t-1),$$9, and EMA $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$0 (Lakkapragada et al., 2022).

The appendix of the same work attributes an important part of this behavior to the decay parameter $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$1. Smaller $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$2 values, such as $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$3 to $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$4, produce smoother loss-coefficient trajectories and better performance, whereas $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$5 yields sharper spikes in $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$6 and subsequent spikes in the overall loss. This is particularly relevant for UEMA because the uncertainty proxy $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$7 is itself an EMA-derived statistic. A plausible implication is that uncertainty guidance inherits the same stability–responsiveness trade-off already observed for first-order EMA weighting: overly aggressive tracking can make the uncertainty signal as volatile as the raw loss it was meant to regularize.

UEMA also clarifies a common misconception. The objective is not necessarily to suppress difficult tasks. In the proposed formulations, large mean loss and large variance are separable signals. This suggests that UEMA can distinguish stable-but-hard tasks from unstable-noisy tasks, rather than uniformly downweighting all large-loss tasks.

4. Relation to uncertainty weighting and EMA theory

UEMA is closely related to the uncertainty weighting literature discussed in the MTL paper. In homoscedastic uncertainty weighting, each task has a parameter $\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$8, and the multi-task objective takes the form

$\lambda_k(t) = \frac{1}{\tilde{L}_k(t)}.$9

The supplied synthesis notes that, in practice, the learned coefficients often converge to something close to inverse loss magnitude, $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$0, but that the method can exhibit volatile spikes in weights. EMA-based weighting can therefore be interpreted as a direct, smoothed approximation to inverse-loss weighting: $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$1 with the smoothing parameter $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$2 acting as the stabilizer (Lakkapragada et al., 2022).

A complementary theoretical perspective comes from "Understanding SGD with Exponential Moving Average: A Case Study in Linear Regression" (Li et al., 19 Feb 2025). That paper studies EMA over parameter iterates,

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$3

and shows two facts that are conceptually important for UEMA: the variance error of SGD with EMA is always smaller than that of SGD without averaging, and the bias error of SGD with EMA decays exponentially in every eigen-subspace of the data covariance matrix. The excess risk is decomposed into bias and variance terms, and the per-eigenvalue bias factor is

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$4

The same paper also develops a unified averaging framework with time-dependent $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$5, showing that EMA, no averaging, iterate averaging, and tail averaging all fit into a common formalism (Li et al., 19 Feb 2025).

This theory does not analyze UEMA directly, but it suggests a rigorous interpretation of uncertainty-guided averaging. If uncertainty is encoded through adaptive $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$6, then UEMA becomes an averaging scheme that dynamically reallocates bias correction and variance reduction over time or across coordinates. High uncertainty can be mapped to one averaging regime and low uncertainty to another. The theory does not prove such a rule optimal, but it provides the operator-level machinery in which it can, in principle, be analyzed.

5. Broader interpretations of EMA relevant to UEMA

The UEMA concept is reinforced by work outside MTL because several papers use EMA to track or privilege statistically meaningful structure beyond simple recency. In "Forecasting Value-at-Risk with Time-Varying Variance, Skewness and Kurtosis in an Exponential Weighted Moving Average Framework" (Gabrielsen et al., 2012), variance, skewness, and kurtosis are all updated by separate EWMA recursions: $$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$7

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$8

$$\mathcal{L}_{\text{MTL}(t)} = \sum_{k=1}^{K} \frac{L_k(t)}{\tilde{L}_k(t)}.$$9

Those higher moments are then embedded into a modified Gram–Charlier density and Cornish–Fisher VaR. Empirically estimated decay factors differ across moments—for example, for the S&P 500, $\sum_k \lambda_k = 1$0, $\sum_k \lambda_k = 1$1, and $\sum_k \lambda_k = 1$2—which means the model already implements uncertainty-sensitive EMA through different memory scales for volatility, skewness, and kurtosis (Gabrielsen et al., 2012). This is not UEMA in the MTL sense, but it is an explicit example in which EMA tracks multiple uncertainty descriptors jointly.

In "Predictor-Corrector Enhanced Transformers with Exponential Moving Average Coefficient Learning" (Li et al., 2024), EMA appears as a structural prior over coefficients rather than as a temporal smoother over losses or weights. The predictor is written as

$\sum_k \lambda_k = 1$3

and the corrector uses a geometrically decaying combination of recent derivatives: $\sum_k \lambda_k = 1$4 The paper does not explicitly refer to uncertainty, but it interprets later Runge–Kutta stages as more accurate approximations and assigns them larger EMA-shaped weights. This suggests an implicit uncertainty prior: more refined approximations are treated as more trustworthy (Li et al., 2024).

A further complement comes from "Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits" (Morales-Brotons et al., 2024), which studies EMA over SGD weights and reports improved generalization, robustness to noisy labels, prediction consistency, calibration, and transfer learning. The EMA update is

$\sum_k \lambda_k = 1$5

The paper reports, for example, that on CIFAR-100 with ResNet-18 the baseline accuracy is $\sum_k \lambda_k = 1$6 and the EMA model selected by best accuracy reaches $\sum_k \lambda_k = 1$7, and that prediction churn drops from $\sum_k \lambda_k = 1$8 to $\sum_k \lambda_k = 1$9 while JS divergence drops from $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$0 to $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$1 (Morales-Brotons et al., 2024). These observations do not define UEMA, but they reinforce the idea that EMA already improves several quantities commonly associated with predictive uncertainty. UEMA can therefore be read as making explicit what fixed-decay EMA only captures implicitly.

6. Limitations, misconceptions, and open directions

The principal limitation of UEMA is that its canonical form remains inferential rather than standardized. The cited MTL paper introduces EMA, REMA, and DWEMA, but not a named UEMA algorithm; the uncertainty-guided formulation is a technically natural extension of its EMA machinery (Lakkapragada et al., 2022). Likewise, the linear-regression theory paper analyzes constant or time-dependent EMA, but not an explicit uncertainty-driven policy (Li et al., 19 Feb 2025). UEMA is therefore an overview grounded in published mechanisms rather than a universally fixed method.

Several practical limitations are already visible in the underlying EMA literature. First, performance depends on smoothing hyperparameters. In MTL, $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$2 that is too high causes instability; too low slows adaptation (Lakkapragada et al., 2022). In weight EMA, BatchNorm statistics can become the limiting factor for very slow averages, so recomputation of BN statistics is needed for large averaging windows (Morales-Brotons et al., 2024). Second, EMA-based normalization assumes that bringing tasks to similar effective scales is beneficial; that assumption can fail if low-loss tasks are disproportionately important (Lakkapragada et al., 2022). Third, uncertainty proxies derived from loss variance are not equivalent to Bayesian uncertainty or predictive posterior variance. The linear-regression paper explicitly notes that its variance analysis concerns gradient-noise variance rather than Bayesian uncertainty (Li et al., 19 Feb 2025).

A related misconception is that UEMA necessarily downweights uncertainty. The proposed formulations do not require that interpretation. The supplied synthesis explicitly notes two possibilities: highly uncertain tasks may be downweighted to stabilize learning, or they may be temporarily upweighted if the supervision is reliable but the task is intrinsically hard (Lakkapragada et al., 2022). UEMA is therefore a framework for encoding policies over $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$3, not a single monotone rule.

The main research directions follow directly from the source material. One direction is adaptive uncertainty policies that map $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$4 to weights, potentially via meta-learning or reinforcement learning (Lakkapragada et al., 2022). Another is coordinate-wise or layer-wise adaptive averaging, motivated by the general $r_k(t) = \frac{L_k(t-1)}{L_k(t-2)}$5 framework in the SGD analysis (Li et al., 19 Feb 2025). A third is integration with other balancing methods, such as using variance-weighted EMA to inform desired gradient norms in GradNorm-style procedures (Lakkapragada et al., 2022). More broadly, the financial EWMA paper suggests that uncertainty need not be limited to variance alone: higher-order moments can also be tracked by EMA, implying that a more general UEMA could incorporate skewness, kurtosis, or other structured uncertainty indicators when the application demands them (Gabrielsen et al., 2012).

In that broader sense, UEMA is best understood as a design pattern: use EMA to maintain online statistics, interpret those statistics as uncertainty-relevant state, and let that state govern either task weights, model averages, or solver coefficients. The published literature establishes each component separately; UEMA unifies them under an uncertainty-aware EMA perspective.