🔙

# LayerNormの誤差逆伝播法

2023/09/01に公開

LayerNorm[1]の逆伝播計算式備忘録

となるから、ヤコビ行列は

\begin{align*} \frac{\partial y_i}{\partial x_j}&= \frac{\delta_{ij}-1/H}{\sigma} - \frac{x_i - \mu}{\sigma^2}\frac{\partial\sigma}{\partial x_j}\\ &=\frac{\delta_{ij}-1/H}{\sigma} - \frac{x_i - \mu}{2\sigma^3}\frac{\partial}{\partial x_j}\frac1H\sum_{h=1}^{H}(x_h-\mu)^2\\ &=\frac{1}{\sigma}\left(\delta_{ij} - \frac1H - \frac{(x_i - \mu)(x_j - \mu)}{H\sigma^2}\right)\\ &=\frac{1}{\sigma}\left(\delta_{ij} - \frac{1 + y_iy_j}{H}\right)\\ \end{align*}

となる。勾配は

\begin{align*} \frac{\partial \mathcal{L}}{\partial x_j} &=\sum_i\frac{\partial y_i}{\partial x_j}\frac{\partial \mathcal{L}}{\partial y_i}\\ &=\frac1\sigma\left(\frac{\partial \mathcal{L}}{\partial y_j} - \frac1H\left(\sum_i\frac{\partial \mathcal{L}}{\partial y_i}+\left(\bm{y}\cdot\frac{\partial\mathcal{L}}{\partial\bm{y}}\right)y_j\right)\right) \end{align*}

となる。
Pythonでの実装を以下に示す。

import numpy as np

N, H = 10, 20

# forward
x = np.random.randn(N, H)
mu = x.mean(axis=-1, keepdims=True)
sigma = x.std(axis=-1, keepdims=True)
y = (x - mu) / sigma

# backward
dy = np.random.randn(N, H)
a = dy.sum(axis=-1, keepdims=True)
b = np.einsum('...i,...i->...', y, dy)[..., None]
dx = (dy - (a + b * y) / H) / sigma

assert x.shape == dx.shape
{"found_math": true, "script_math_tex": 0, "script_math_asciimath": 0, "math_annotations": 0, "math_alttext": 0, "mathml": 0, "mathjax_tag": 0, "mathjax_inline_tex": 0, "mathjax_display_tex": 0, "mathjax_asciimath": 1, "img_math": 0, "codecogs_latex": 0, "wp_latex": 0, "mimetex.cgi": 0, "/images/math/codecogs": 0, "mathtex.cgi": 0, "katex": 0, "math-container": 0, "wp-katex-eq": 0, "align": 0, "equation": 0, "x-ck12": 0, "texerror": 0, "url": "https://zenn.dev/taro137/articles/2c639b39a32166", "fetch_time": 1726347339000000000, "content_mime_type": "text/html", "warc_filename": "crawl-data/CC-MAIN-2024-38/segments/1725700651580.74/warc/CC-MAIN-20240914193334-20240914223334-00435.warc.gz", "warc_record_offset": 1006338001, "warc_record_length": 12113, "token_count": 535, "char_count": 1281, "score": 3.984375, "int_score": 4, "crawl": "CC-MAIN-2024-38", "snapshot_type": "latest", "language": "en", "language_score": 0.23616671562194824, "file_path": "/sailhome/yjruan/project/latent-lm/data/processed/finemath-4plus/finemath-4plus/train-00003-of-00064.parquet"}