Residual Scaling (normalization strategy)

[PortillaS-Predictia]

Is your feature request related to a problem? Please describe.

Following Watt-Meyer et al., 2023, 2024; ACE & ACE2. Variables are normalized using a residual scaling approach such that predicting outputs equal to input would result in each variable contributing equally to the loss function.

This should be straightforward, since anemoi-datasets already computes the tendencies statistics needed for this normalization strategy. However, the reference formulas (Appendix H) use the standard deviation of the mean-std normalized fields (not the unnormalized fields, which is what anemoi-datasets actually computes these statistics for).

Describe the solution you'd like

See our proposed solution. This is an implementation of the reference formulas (Appendix H), using the tendencies statistics computed by anemoi-datasets. We reworked the reference formulas so that the statistics of the unnormalized fields are used instead:

Let $a$ be the target field, and $a_{\textup{ff}}$ the mean-std normalized (or full-field normalized, using their terminology) image $a_{\textup{ff}}=\frac{a-\mu(a)}{\sigma(a)}$. Then, the residual scaling is $a_{\textup{res}}=\frac{a_{\textup{ff}}}{\sigma_{\textup{res}}(a)}$, where $\sigma_{\textup{res}}(a) = \frac{\sigma(a\prime_\textup{ff})}{\eta_{a\in\mathbf{T}}\left(\sigma(a\prime_\textup{ff})\right)}$ is the standard deviation of the tendency (of the mean-std normalized field, not the unnormalized field), divided by the geometric mean $\eta$ (of this quantity) of all targeted fields $\mathbf{T}$ (hereafter, we will omit the $a\in\mathbf{T}$ for clarity). Notice that these are just the reference equations from the paper (using $\eta$ for the geometric mean instead). Now, we want to rewrite these equations in terms of the statistics of the unnormalized field $a$ (instead of the normalized field $a_\textup{ff}$). Thus, we have $\sigma(a\prime_\textup{ff})=\sigma(a_\textup{ff}(t+1)-a_\textup{ff}(t))=\sigma\left(\frac{a(t+1) - \mu(a)}{\sigma(a)} - \frac{a(t) - \mu(a)}{\sigma(a)}\right)=\frac{\sigma(a(t+1)-a(t))}{\sigma(a)}= \frac{\sigma(a\prime)}{\sigma(a)}$, which depends only on the unnormalized field $a$. Then, we rewrite the residual scaling as $a_{\textup{res}} = \frac{a_{\textup{ff}}}{\sigma_{\textup{res}}(a)} = \frac{(a-\mu(a))/\sigma(a)}{\sigma(a\prime_{\textup{ff}})/\eta\left(\sigma(a\prime_\textup{ff})\right)} = \eta\left(\sigma(a\prime_\textup{ff})\right) \cdot \frac{a-\mu(a)}{\sigma(a)\cdot\sigma(a\prime_{\textup{ff}})}$. Now, using the relation $\sigma(a\prime_\textup{ff})=\frac{\sigma(a\prime)}{\sigma(a)}$, we have $\eta\left(\sigma(a\prime_\textup{ff})\right)=\eta\left(\frac{\sigma(a\prime)}{\sigma(a)}\right)$ and $\frac{a-\mu(a)}{\sigma(a)\cdot\sigma(a\prime_{\textup{ff}})}=\frac{a-\mu(a)}{\sigma(a)\cdot\frac{\sigma(a\prime)}{\sigma(a)}}=\frac{a-\mu(a)}{\sigma(a\prime)}$. Thus, $a_{\textup{res}} = \eta\left(\frac{\sigma(a\prime)}{\sigma(a)}\right)\cdot\frac{a-\mu(a)}{\sigma(a\prime)}$, which depends only on the unnormalized field $a$ (which is what anemoi-datasets actually computes the tendencies statistics for). Thus, the residual scaling consists of adding $\mathbf{add}=-\frac{\mu(a)\cdot\eta\left(\frac{\sigma(a\prime)}{\sigma(a)}\right)}{\sigma(a\prime)}$ and multiplying by $\mathbf{mul}=\frac{\eta\left(\frac{\sigma(a\prime)}{\sigma(a)}\right)}{\sigma(a\prime)}$.

In our implementation, the geometric mean $\eta\left(\frac{\sigma(a\prime)}{\sigma(a)}\right)$ is computed iteratively within the main for loop, and multiplied afterwards to both $\mathbf{add}$ and $\mathbf{mul}$ (notice that we use the logarithmic definition of geometric mean). If the tendencies' stdev doesn't exist in the statistics dictionary (because the tendencies statistics weren't computed during the creation of the dataset), then the code fallbacks to using the stdev instead, in which case the formulae above reduces to a mean-std normalization.

Additionally, we added the tendencies' stdev (and its ratio to the stdev) in the inspect command in anemoi-datasets.

Describe alternatives you've considered

No response

Additional context

No response

Organisation

Predictia Intelligent Data Solutions - DestinationEarth393

[jakob-schloer]

Thanks for the idea and clear description. We have recently included the option to use residual scaling (called tendency scaling) of the loss function (#52).
Instead of scaling the variables $x$ as you described, we scale the loss function, which should have the same effect:

$MSE_{res} \approx |\hat{x}_{res} - x_{res}|^2 = |\hat{x}/\sigma_{res} - x/\sigma_{res}|^2 = 1/\sigma_{res} |\hat{x} - x|^2$

Do you have a use case where scaling the loss function is not enough?

[PortillaS-Predictia]

Thanks Jakob,

I'm inclined to think both approaches are equivalent, like you propose. However, I'm not confident $\hat{x}_{\text{res}} = \hat{x} / \sigma_{\text{res}}$ is unconditionally true (for a given model $M_{\theta}$, if $M_\theta(x / \sigma_\text{res}) = M_\theta(x) / \sigma_\text{res}$ for any $\sigma_\text{res}$, then $M_\theta$ would be linear function — homogeneous, at least). Moreover, Anemoi's models work residually on the prognostic variables $\hat{x}(t+1) = x(t) + \delta M_\theta(x(t))$ (there is a residual connection around the model). Thus, maybe $x(t+1) - x(t)$ would be the better target for normalization (regardless of the loss scaling). Indeed, any weight initialization of $\theta$ would be better suited for a residually scaled target $x$ (see images below, the PlotHistogram at the very beginning of training, top image with residual scaling and bottom one without — though the effect isn't as noticeable for all variables and it heavely depends on the model's recipe), I don't think it's possible to mimic this effect with any scaling of the loss function. Indeed, we have observed that weight initialization is somehow very important for the long-term (autorregresive) stability of the model (two runs, with different initializations but identical otherwise, may have different long-term stability) (we have observed noticeable stability improvements when using residual scaling, though we haven't tested the other alternative).

However, I agree both approaches are likely interchangeable for most use cases.

[jakob-schloer]

Just to clarify, you want to use the residual scaled variables as input, i.e. $M(x(t) / \sigma_{res})$? Is that your use case?
In this case, we have not considered that yet. My intuition, however, would be that it is beneficial for the model to get the absolute values of the input fields. Since we typically pass two fields, $x(t), x(t-1) $, the difference between the input states is accessible to the model.

This is a different story for the output. As you mentioned, we have a skip connection in the model, so gradients are naturally larger for fast-varying variables than for slow-varying variables. The residual scaling of the loss should account for that and scale the loss accordingly.

I don't quite understand how the weight initialization of the NN is connected to that. It could well be that we don't optimally initialize the weights of the NN, but could this not be addressed by explicitly setting the weight initialization?
Maybe I'm missing something. I'm happy to be convinced :)

[PortillaS-Predictia]

Thanks Jakob,

Indeed, we want to use the residual-scaled variables, following Watt-Meyer et al., 2023, 2024; ACE & ACE2. Notice that this isn't too different from other common scalings (such as mean-std scaling) (it's just a linear transformation $x\rightarrow \mathbf{mul}\cdot x + \mathbf{add}$). I'm afraid I don't understand what you mean by it is beneficial for the model to get the absolute values of the input fields. Just to make it clear — residual scaling does not mean the model will use $x(t)-x(t-1)$ as input.

I still believe $\hat{x}_\text{res}(t+1)=\hat{x}(t+1) / \sigma_\text{res}$ isn't true, unless I misunderstood your notation. If $\hat{x}_\text{res}(t+1)$ is the model's output when fed with the residual-scaled input $x_\text{res}(t, t-1, ...)$, and $\hat{x}(t+1)$ is the model's output when fed with the unscaled input $x(t, t-1, ...)$, I don't see how it follows that $\hat{x}_\text{res}(t+1)=\hat{x}(t+1) / \sigma_\text{res}$. I agree with the statement the residual scaling of the loss should account for that and scale the loss accordingly.

How does this relate to weight initialization? Consider a slowly varying field ($x$ such that $x\prime\simeq 0$). Then, for a given $t_0$, we would have $x(t_0+1) \simeq x(t_0)$ (equivalently, $x(t_0+1) - x(t_0) \simeq 0$, or simply $x\prime(t_0)\simeq 0$). However, for most sensible initializations of $\theta$, we would have $|M_\theta(x(t_0, t_0 - 1, ...))|$ > 0 (simply put, the nn's output isn't zero). Thus, the forecast for $x\prime(t_0)$ (which is what $M_\theta(x(t_0, t_0 - 1, ...))$ is actually forecasting) is far off — it's greater than zero but it should be $\simeq$ zero (regardless of the loss' scaling). Notice that the condition $x\prime\simeq0$ would never happen with a residual-scaled variable. I agree with the statement this [could] be addressed by explicitly setting the weight initialization.

Hopefully, we will soon update our Anemoi version and have the chance to compare both approaches.

[timothyas]

Hi folks, I'll just chime in since I've been very interested to add the residual scaling in Anemoi myself, so I'm very glad to see this discussion.

Just a comment on using residual scaling for only outputs vs inputs & outputs, the authors of the ACE paper weigh in on this in the last 2 sentences of Appendix H:

Residual normalization is more justifiable for prognostic variables (which are both inputs and outputs) than for diagnostic variables. Hence, we attempted to apply residual scaling only to prognostic variables, but this leads to slightly larger time-mean biases than also rescaling diagnostic variables.

From the way that's worded, it sounds to me like it's not entirely understood why normalizing all time-varying variables via the residual scaling works... but I guess it seems to for their application. In any case it sounds like it would be nice to have both options:

1. residual (or tendency, etc) scaling only on the outputs and "mean-std" normalization on the inputs. This is what's done in GraphCast and is what @jakob-schloer mentioned (it sounds like this was implemented in fix!: Rework Loss Scalings to provide better modularity #52, is that right?)
2. residual scaling of all time varying variables as in ACE, including the geometric mean factor, etc, as @PortillaS-Predictia has been advocating for

[jakob-schloer]

Hi @timothyas ,
Thanks for the nice summary. I don't quite understand the difference since the means should cancel in the loss.

[jakob-schloer]

Indeed, we want to use the residual-scaled variables, following Watt-Meyer et al., 2023, 2024; ACE & ACE2. Notice that this isn't too different from other common scalings (such as mean-std scaling) (it's just a linear transformation x → mul ⋅ x + add ). I'm afraid I don't understand what you mean by it is beneficial for the model to get the absolute values of the input fields. Just to make it clear — residual scaling does not mean the model will use x ( t ) − x ( t − 1 ) as input.

I still believe x ^ res ( t + 1 ) = x ^ ( t + 1 ) / σ res isn't true, unless I misunderstood your notation. If x ^ res ( t + 1 ) is the model's output when fed with the residual-scaled input x res ( t , t − 1 , . . . ) , and x ^ ( t + 1 ) is the model's output when fed with the unscaled input x ( t , t − 1 , . . . ) , I don't see how it follows that x ^ res ( t + 1 ) = x ^ ( t + 1 ) / σ res . I agree with the statement the residual scaling of the loss should account for that and scale the loss accordingly.

Thanks for the clarification @PortillaS-Predictia and sorry for my sloppy notation of $\hat{x}_{res}$. Here is a more explicit formulation of what we do, taking a similar notation to ACE but neglecting the ff part since these can be factored out.

We define the residual of a field as $x'(t+1) = x(t+1) - x(t) $ and the normalized residual as $\tilde{x}'(t+1) = (x(t+1) - x(t)) / \sigma_{res}$ where $\sigma_{res} = \sqrt{ E[x'^2] - (E[x'])^2}$ is the standard deviation of the residual. Our model predicts the residual, i.e. $ \hat{x}(t+1) = x(t) + M(x(t), x(t-1)) => \hat{x'}(t+1) = M(x(t), x(t-1))$.

When we normalize the loss by $\sigma_{res}$ we get:

$MSE \propto 1/\sigma^2_{res} | \hat{x}(t+1) - x(t+1) |^2 = |x(t) + M(x(t), x(t-1)) - x(t+1) | = | M(x(t), x(t-1))/\sigma_{res} - (x(t+1) - x(t))/\sigma_{res} |$

This way, the loss is computed over the normalized residuals at the respective timestep.
Can you explain how this is different from ACE?

How does this relate to weight initialization? Consider a slowly varying field ( x such that x ′ ≃ 0 ). Then, for a given t 0 , we would have x ( t 0 + 1 ) ≃ x ( t 0 ) (equivalently, x ( t 0 + 1 ) − x ( t 0 ) ≃ 0 , or simply x ′ ( t 0 ) ≃ 0 ). However, for most sensible initializations of θ , we would have | M θ ( x ( t 0 , t 0 − 1 , . . . ) ) | > 0 (simply put, the nn's output isn't zero). Thus, the forecast for x ′ ( t 0 ) (which is what M θ ( x ( t 0 , t 0 − 1 , . . . ) ) is actually forecasting) is far off — it's greater than zero but it should be ≃ zero (regardless of the loss' scaling). Notice that the condition x ′ ≃ 0 would never happen with a residual-scaled variable. I agree with the statement this [could] be addressed by explicitly setting the weight initialization.

That's a nice observation/argument. Did you try to set the weight initialization to a very small range?

[timothyas]

Hey @jakob-schloer 👋

So the difference between 1 and 2 is that in 2 the input fields to the network are normalized by residuals, isn't it?

That's my understanding, yes. They don't come out and say that explicitly as far as I can tell, but the just say in the Data Normalization section "Variables are normalized using the "residual scaling" approach ...". Additionally, they report the values of the residual statistics for all variables, including forcings (input only).

[PortillaS-Predictia]

Thanks Jakob,

From your last equation (hereafter, Eq. 1), we have $\text{MSE}_\text{res} \propto |\frac{M(x(t), x(t-1))}{\sigma_\text{res}} - \frac{x(t+1)-x(t)}{ \sigma_\text{res}}\,|^2$.

Now, I'll try to deduce the same formula using ACE's approach. Let $x_\text{res}=x / \sigma_\text{res}$ and $\hat{x}_\text{res}(t+1) = x_\text{res}(t) + M(x_\text{res}(t), x_\text{res}(t-1))$. Thus, we have (hereafter, Eq. 2):

$MSE_\text{ace} \propto |\hat{x}_\text{res}(t+1) - x_\text{res}(t+1)\,|^2 = |M(x_\text{res}(t), x_\text{res}(t-1)) - (x_\text{res}(t+1) - x_\text{res}(t))\,|^2 = |M(\frac{x(t)}{\sigma_\text{res}}, \frac{x(t-1)}{\sigma_\text{res}}) - \frac{x(t+1) - x(t)}{\sigma_\text{res}}\,|^2$

Now, it's clear that Eq. 1 and Eq. 2 are equivalent if and only if $M(\frac{x(t)}{\sigma_\text{res}}, \frac{x(t-1)}{\sigma_\text{res}}) = \frac{M(x(t), x(t-1))}{\sigma_\text{res}}$, which is generally not true.

So, what's the difference, really? From Eq. 1, it's clear that $\text{MSE}_\text{res} = 0$ (i.e., perfect forecast) if and only if $M(x(t), x(t-1)) = x(t+1)-x(t)$. Now, consider a model with two variables $x_1$ and $x_2$, such that $x_1(t+1)-x_1(t)\sim0$ and $x_2(t+1)-x_2(t)\sim1$ (i.e., $x_1$ is slow-varying and $x_2$ is fast-varying). Then, we would have $\text{MSE}_\text{res} = 0$ only if $M(x(t), x(t-1))[1]\sim0$ and $M(x(t), x(t-1))[2]\sim1$. Thus, for a perfect forecast the model's outputs must live on completely different scales (for instance, 1e-3 and 1). However, this isn't the case when following ACE's approach, since $\frac{x(t+1) - x(t)}{\sigma_\text{res}}\sim1$ always (for both slow-varying and fast-varying variables), so the model's outputs (for a perfect forecast) are always $\text{Normal}(0, 1)$. Now, how relevant is this? That I don't know, but I'm inclined to think that the model's outputs should live on similar scales, especially when all channels are mixed in the deeper layers of the network.

Also, thanks Timothy for chiming in.

[jakob-schloer]

Thanks @PortillaS-Predictia for the clear explanation.

Now, I'll try to deduce the same formula using ACE's approach. Let x res = x / σ res and x ^ res ( t + 1 ) = x res ( t ) + M ( x res ( t ) , x res ( t − 1 ) ) . Thus, we have (hereafter, Eq. 2):

M S E ace ∝ | x ^ res ( t + 1 ) − x res ( t + 1 ) | 2 = | M ( x res ( t ) , x res ( t − 1 ) ) − ( x res ( t + 1 ) − x res ( t ) ) | 2 = | M ( x ( t ) σ res , x ( t − 1 ) σ res ) − x ( t + 1 ) − x ( t ) σ res | 2

The main difference between GraphCast and ACE residual normalization is that in ACE, the input to the network is normalized by the residual, i.e., $x_{res} = x / \sigma_{res}$.
While it's still not fully clear to me why one would do that, since $x(t)$ and $(x(t) - x(t-1))$ have quite different distributions, it seems to work well for ACE and so in my opinion worth to explore.

Before bringing the ACE normalization strategy into anemoi, I think we should explore how well it works for our architecture and what the difference in performance is to the GraphCast normalization.

[clessig]

To add on to @jakob-schloer message: @PortillaS-Predictia could you rebase to the current anemoi and then do a clean comparison of the two approaches? This would be very helpful and probably interesting for others as well (such as @timothyas )!