pytorch中的函数定义如下：
torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)函数参数说明如如下：
normalized_shape: 进行LayerNorm的维度定义，对于一个多维矩阵[N, C, H, W]来说，这里的normalized_shape定义都是要和矩阵最后几个维度保持一致的，这里就是[C, H, W]。对比数学公式，其中的 和
的维度都是[C, H, W]，
和
的维度都是[N, C, H, W]。
eps：为了防止计算公式中的分母为0，加上一个极小的数，默认值: 1e-5
elementwise_affine：设为True的时候，进行elementwise的仿射变换, 和
才会生效，在训练过程中做为参数会被学习更新，为False的话不生效。
所有元素初始为1，
所有元素初始为0的。
在代码实现中对应
,
在代码实现中对应
。
LayerNorm的数学公式定义如下：
pytorch使用示例，给定一个[N, C, H, W]的矩阵，在[C, H, W]维上进行LayerNorm操作：
>>> # Image Example
>>> N, C, H, W = 20, 5, 10, 10
>>> input = torch.randn(N, C, H, W)
>>> # Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
>>> # as shown in the image below
>>> layer_norm = nn.LayerNorm([C, H, W])
>>> output = layer_norm(input)
2. LayerNorm反向推导公式
为了方便推导，eps先忽略，输入为一维矩阵。对应LayerNorm的数学公式定义如下, 其中是由
组成的一维向量,
是输出向量，维度跟
一样;
是期望，简写为
;
是方差【
】; 标准差【
】简写为
。
这里有三个地方需要求梯度(即需要进行求导)，分别是对参数gamma和beta
, 以及输入x的求导, 即
、
、
。同时在计算
时会用到
，
，
。
这里 与
是一一对应的, 所以不用累加；但对于
参与了所有
的计算，反向的时候计算梯度也需要对涉及的所有的
3. 源码实现
代码仓版本：https://github.com/pytorch/pytorch/tree/v2.0.1
3.1 前向计算
在aten/src/ATen/native/native_functions.yaml中的定义如下：
- func: native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
dispatch:
CPU: layer_norm_cpu
CUDA: layer_norm_cuda
MPS: layer_norm_mps
CompositeExplicitAutograd: math_native_layer_norm
NestedTensorCPU, NestedTensorCUDA: nested_layer_norm
autogen: native_layer_norm.out
tags: core这里以layer_norm_cpu的实现为例，layer_norm_cpu定义在aten/src/ATen/native/layer_norm.cpp中。
在layer_norm_cpu的前向函数中，会根据input和normalized_shape进行shape的转换计算，从多维矩阵转为的二维矩阵，比如input的shape是[2, 3, 4, 5]，normalized_shape是[4, 5], 那么M=2*3=6, N=4*5=20；同时还会进行weight(对应
)和bias(对应
)矩阵的初始化。
std::tuple<Tensor, Tensor, Tensor> layer_norm_cpu(
const Tensor& input,
IntArrayRef normalized_shape, const c10::optional<Tensor>& weight_opt /* optional */, const c10::optional<Tensor>& bias_opt /* optional */,
double eps) {
// weight和bias初始化
c10::MaybeOwned<Tensor> weight_maybe_owned = at::borrow_from_optional_tensor(weight_opt);
const Tensor& weight = *weight_maybe_owned;
c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
const Tensor& bias = *bias_maybe_owned;
// 计算M和N
auto M_N = _check_layer_norm_inputs(input, normalized_shape, weight, bias);
auto M = M_N.first;
auto N = M_N.second;
auto X = input.expect_contiguous();
auto gamma = weight.expect_contiguous();
auto beta = bias.expect_contiguous();
// 初始化mean/rstd，维度是M个，每N个input的元素会计算一个mean和rstd
Tensor mean = at::empty({M}, X->options().dtype(dtype));
Tensor rstd = at::empty({M}, X->options().dtype(dtype));
// layer_norm_with_mean_rstd_out中会调用前向kernel(LayerNormKernel)
layer_norm_with_mean_rstd_out(Y, mean, rstd, *X, normalized_shape, *gamma, *beta, eps, M, N);
return std::make_tuple(std::move(Y), std::move(mean), std::move(rstd));
}LayerNormKernel定义在aten/src/ATen/native/cpu/layer_norm_kernel.cpp中，实际的实现是LayerNormKernelImplInternal, 定义如下：
template <typename T, typename T_ACC>
void LayerNormKernelImplInternal(
const Tensor& X,
const Tensor& gamma,
const Tensor& beta,
int64_t M,
int64_t N,
T_ACC eps,
Tensor* Y,
Tensor* mean,
Tensor* rstd) {
...
}在LayerNormKernelImplInternal首先了解at::parallel_for函数的使用，它的基本作用是对输入先进行分块，然后通过多线程进行并行处理，如下函数的定义是对[0, M]分成多段，分别调用匿名函数。
at::parallel_for(0, M, 1, [&](int64_t start, int64_t end) {...})回顾下前向计算过程：
匿名函数逻辑中，对于 M * N的矩阵，每次处理N个元素进行LayerNorm操作。mean对应, rstd_val和scale对应
, bias对应
, 因此,
for (const auto i : c10::irange(start, end)) {
const T* X_ptr = X_data + i * N;
T* Y_ptr = Y_data + i * N;
T mean_val;
T rstd_val;
// 1. 计算mean_val和rstd_val
std::tie(mean_val, rstd_val) = RowwiseMoments(X_ptr, N);
rstd_val = T(1) / std::sqrt(rstd_val + eps);
const T scale = rstd_val;
const T bias = -rstd_val * mean_val;
if (gamma_null || beta_null) {
for (const auto j : c10::irange(N)) {
const T gamma_v = gamma_null ? T(1) : gamma_data[j];
const T beta_v = beta_null ? T(0) : beta_data[j];
Y_ptr[j] = (X_ptr[j] * scale + bias) * gamma_v + beta_v;
}
} else {
// 2. 计算layer norm的前向公式
vec::map3<T>(
[scale, bias](Vec x, Vec gamma, Vec beta) {
return (x * Vec(scale) + Vec(bias)) * gamma + beta;
},
Y_ptr,
X_ptr,
gamma_data,
beta_data,
N);
}
if (!mean_null) {
mean_data[i] = mean_val;
}
if (!rstd_null) {
rstd_data[i] = rstd_val;
}
}
}3.2 反向计算
对于多维矩阵求反向，可以看成是M个大小为N的向量，以一个5维向量为例，向量维度为，layer_norm的维度是
，对应的
,
在aten/src/ATen/native/native_functions.yaml中的定义如下：
- func: native_layer_norm_backward(Tensor grad_out, Tensor input, SymInt[] normalized_shape, Tensor mean, Tensor rstd, Tensor? weight, Tensor? bias, bool[3] output_mask) -> (Tensor, Tensor, Tensor)
dispatch:
CPU: layer_norm_backward_cpu
CUDA: layer_norm_backward_cuda
MPS: layer_norm_backward_mps
autogen: native_layer_norm_backward.out
tags: core这里以layer_norm_backward_cpu的实现为例，layer_norm_backward_cpu定义在aten/src/ATen/native/layer_norm.cpp中。跟layer_norm_cpu类似，在backward中初始化相关tensor，和进行kernel的调用。
std::tuple<Tensor, Tensor, Tensor> layer_norm_backward_cpu(
const Tensor& dY,
const Tensor& input,
IntArrayRef normalized_shape,
const Tensor& mean,
const Tensor& rstd,
const c10::optional<Tensor>& weight_opt /* optional */,
const c10::optional<Tensor>& bias_opt /* optional */,
std::array<bool, 3> grad_input_mask) {
......
if (M > 0) {
LayerNormBackwardKernel(
kCPU, dY, *X, mean, rstd, *gamma, M, N, &dX, &dgamma, &dbeta);
}
return std::make_tuple(std::move(dX), std::move(dgamma), std::move(dbeta));
}为了方便和后续pytorch源码实现中对应，对上面推导公式的最后结果中做下相应的展开，展开如下：
kernel的实现在是aten/src/ATen/native/cpu/layer_norm_kernel.cpp文件的LayerNormBackwardKernelImplInternal函数中，实现分为两个阶段：
- 初始化一个shape大小为{2, max_threads, N}的buffer矩阵，对应其中的buffer[0]用于dgamma_buffer, buffer[1]用于dbeta_buffer。多线程分别计算dY和X。
- 对dgamma/dbeta的值进行累加操作，复用X[i]和dY[i]
对于代码实现是通过两层嵌套进行的，对于第一步来说，最外面是对 的矩阵按行进行多线程并行，每个线程处理
个元素；第二步是按N列进行元素的累加。layer_norm_backward_frame函数中包含了主要的计算逻辑, 后面进一步分析。
template <typename T>
void LayerNormBackwardKernelImplInternal(
const Tensor& dY,
const Tensor& X,
const Tensor& mean,
const Tensor& rstd,
const Tensor& gamma,
int64_t M,
int64_t N,
Tensor* dX,
Tensor* dgamma,
Tensor* dbeta) {
......
// 第一步：计算dgamma/dbeta and dX
at::parallel_for(0, M, 1, [&](int64_t start, int64_t end) {
int tid = at::get_thread_num();
TORCH_CHECK(
tid < num_threads,
"expect thread id smaller than ",
num_threads,
", got thread id ",
tid);
T* dgamma_buffer_ptr = dgamma_null ? nullptr : buffer_data + tid * N;
T* dbeta_buffer_ptr =
dbeta_null ? nullptr : buffer_data + num_threads * N + tid * N;
for (const auto i : c10::irange(start, end)) {
layer_norm_backward_frame<T, T2, T_ACC>(dY_data, X_data, mean_data, rstd_data, gamma_data, dX_data, dgamma_buffer_ptr, dbeta_buffer_ptr, scale, gamma_null, dX_null, dgamma_null, dbeta_null, N, i);
}
});
// 第二步：计算dgamma/dbeta的累加
if (buffer_data != nullptr) {
parallel_for(0, N, 1, [&](int64_t start, int64_t end) {
for (const auto j : c10::irange(start, end)) {
T_ACC dgamma_v = T_ACC(0);
T_ACC dbeta_v = T_ACC(0);
for (const auto i : c10::irange(num_threads)) {
dgamma_v += buffer_data[i * N + j];
dbeta_v += buffer_data[num_threads * N + i * N + j];
}
if (!dgamma_null) {
// NOLINTNEXTLINE(clang-analyzer-core.NullDereference)
dgamma_data[j] = dgamma_v;
}
if (!dbeta_null) {
// NOLINTNEXTLINE(clang-analyzer-core.NullDereference)
dbeta_data[j] = dbeta_v;
}
}
});
}
......
}layer_norm_backward_frame函数中计算dgamma的逻辑如下，对应公式：, 其中
,
。
if (!dgamma_null) {
const T_ACC a = rstd_data[i];
const T_ACC b = -a * mean_data[i];
// Scalar math:
// for (const auto j : c10::irange(N)) {
// dgamma_data[j] += dY_ptr[j] * (a * X_ptr[j] + b);
// }
vec::map3<T>(
[a, b](Vec dgamma, Vec dy, Vec x) {
return dgamma + dy * (Vec(a) * x + Vec(b));
},
dgamma_buffer_ptr,
dgamma_buffer_ptr,
dY_ptr,
X_ptr,
N);
}layer_norm_backward_frame函数中计算dbeta的逻辑如下，对应公式：。
if (!dbeta_null) {
// Scalar math:
// for (const auto j : c10::irange(N)) {
// dbeta_data[j] += dY_ptr[j];
// }
vec::map2<T>(
[](Vec dbeta, Vec dy) { return dbeta + dy; },
dbeta_buffer_ptr,
dbeta_buffer_ptr,
dY_ptr,
N);
}layer_norm_backward_frame函数中计算dx的逻辑如下，对应公式:
layer_norm_backward_frame函数核心代码实现如下：
if (gamma_null) {
......
} else {
ds = vec::map3_reduce_all<T>(
[](Vec x, Vec y, Vec z) { return x * y * z; },
[](Vec x, Vec y) { return x + y; },
dY_ptr,
X_ptr,
gamma_data,
N);
db = vec::map2_reduce_all<T>(
[](Vec x, Vec y) { return x * y; },
[](Vec x, Vec y) { return x + y; },
dY_ptr,
gamma_data,
N);
}
const T_ACC a = rstd_data[i];
const T_ACC b = (db * mean_data[i] - ds) * a * a * a * scale;
const T_ACC c = -b * mean_data[i] - db * a * scale;
if (gamma_null) {
......
} else {
vec::map3<T>(
[a, b, c](Vec dy, Vec gamma, Vec x) {
return Vec(a) * dy * gamma + Vec(b) * x + Vec(c);
},
dX_ptr,
dY_ptr,
gamma_data,
X_ptr,
N);
}
}代码中变量与公式对应关系如下：
-
ds对应
-
db对应
-
a对应
-
scale对应
-
b对应
-
c对应
- 最终结果：dx = Vec(a) * dy * gamma + Vec(b) * x + Vec(c)
4. 参考资料
- Vector, Matrix, and Tensor Derivatives
- 手推公式之“层归一化”梯度
- 矩阵求导浅析（一）
- 矩阵求导术（上）
- 矩阵求导术（下）
- 道理我都懂，但是神经网络反向传播时的梯度到底怎么求？
- 神经网络反向传播的数学原理
- layernorm 反向传播推导及代码
- pytorch-LAYERNORM
- Layer Normalization-paper
