layer_norm_backward_f32

void layer_norm_backward_f32(const float * x, const float * gamma, const float * saved_mean, const float * saved_rstd, const float * g, float * dx, float * dgamma, float * dbeta, int outer, int N)

source edit

Single-precision LayerNorm backward — produces input, weight, and bias gradients from a saved-tensor forward pass.

Uses the closed-form LayerNorm gradient that avoids recomputing the row statistics. dgamma and dbeta are accumulated into the supplied buffers; callers reducing over multiple sub-batches must zero them once at the start. dx is written fully (not accumulated).

Math

Let $\hat{x}_i = (x_i - \mu)\,\text{rstd}$ and $\widetilde{g}_i = \gamma_i g_i$ . Then

\frac{\text{rstd}}{N} \!\left(N \widetilde{g}_i - \sum_j \widetilde{g}_j - \hat{x}_i \sum_j \widetilde{g}_j \hat{x}_j\right),$$ $$\frac{\partial L}{\partial \gamma_i} = \sum_{\text{rows}} g_i \hat{x}_i, \qquad \frac{\partial L}{\partial \beta_i} = \sum_{\text{rows}} g_i.$$

Parameters

xconst float*

Original forward input of shape (outer, N).

gammaconst float*

Affine weight of shape (N,).

saved_mean, saved_rstdconst float*

Per-row statistics produced by layer_norm_forward_f32.

gconst float*

Upstream gradient

\partial L / \partial y

of shape (outer, N).

dxfloat*

Output gradient w.r.t. x; shape (outer, N); overwritten.

dgamma, dbetafloat*

Output gradients w.r.t. the affine parameters; shape (N,); accumulated (caller must zero on first call).

outer, Nstd::size_t

Layout extents.