Матричное дифференцирование

Сначала про соглашение

Прежде чем считать, нужно договориться, как располагать результат. Есть два соглашения, и они дают транспонированные друг другу ответы.

• Numerator layout (он же якобианов): индекс выхода нумерует строки. Тогда $\partial \mathbf{y}/\partial \mathbf{x}$ для $\mathbf{y} \in \R^m$, $\mathbf{x} \in \R^n$ имеет форму $m \times n$ — это якобиан из прошлого урока.
• Denominator layout (он же градиентный): индекс входа нумерует строки. Та же производная получает форму $n \times m$.

numerator layout denominator layout

производная	формы операндов	форма результата
∂f/∂x	f скаляр, x из Rⁿ	1 × n (строка)
∂y/∂x	y из Rᵐ, x из Rⁿ	m × n
∂(Ax)/∂x	A размера m × n	A
∂f/∂W	f скаляр, W размера m × n	n × m

Практическое правило: в машинном обучении почти всегда denominator layout, потому что градиент должен иметь ту же форму, что параметр — иначе шаг $\mathbf{w} \leftarrow \mathbf{w} - \eta\nabla f$ не записать. В учебниках по анализу и в теоретических статьях чаще numerator.

Как жить с этим: не запоминать соглашение, а проверять формы. Если в вашей выкладке градиент по $W$ имеет форму $W$ — вы в denominator layout и всё сойдётся с кодом. Дальше в этом уроке используется он.

Формула первая: линейное отображение

$\frac{\partial (A\mathbf{x})}{\partial \mathbf{x}} = A^\top$

Проверяется по элементам: $(A\mathbf{x})_i = \sum_j A_{ij}x_j$, значит $\partial (A\mathbf{x})_i/\partial x_j = A_{ij}$. В denominator layout строки нумеруются входом $j$, поэтому на месте $(j, i)$ стоит $A_{ij}$ — это и есть $A^\top$.

Это тот же факт, что «якобиан линейного отображения равен самому отображению», записанный в другом соглашении.

Формула вторая: квадратичная форма

$$\frac{\partial}{\partial \mathbf{x}}\htmlData{k=q}{\mathbf{x}^\top A \mathbf{x}} = \htmlData{k=g}{(A + A^\top)\mathbf{x}}$$

Для симметричной $A$ это $2A\mathbf{x}$ — прямой аналог $(x^2)' = 2x$. Гессианы и ковариационные матрицы симметричны, поэтому на практике встречается почти всегда эта форма.

Вывод стоит проделать один раз по индексам: $\mathbf{x}^\top A\mathbf{x} = \sum_{i,j}A_{ij}x_ix_j$, и при дифференцировании по $x_k$ переменная $x_k$ встречается дважды — в роли $i$ и в роли $j$. Отсюда и сумма двух слагаемых. Это ровно та «сумма по путям» из урока про правило цепочки.

Формула третья: наименьшие квадраты

$\frac{\partial}{\partial \mathbf{x}}\|A\mathbf{x} - \mathbf{b}\|^2 = 2A^\top(A\mathbf{x} - \mathbf{b})$

Самая используемая формула блока. Вывод — цепочка из двух звеньев: внешнее $\|\mathbf{r}\|^2$ даёт $2\mathbf{r}$, внутреннее $\mathbf{r} = A\mathbf{x}-\mathbf{b}$ даёт $A^\top$ (первая формула).

Приравняв нулю, получаем $A^\top A\mathbf{x} = A^\top \mathbf{b}$ — нормальное уравнение из практики блока 1. Тогда мы вывели его геометрически, через ортогональность остатка; теперь он же получается дифференцированием. Два разных пути к одному ответу — хороший признак, что вы поняли оба.

Формула четвёртая: производная по матрице

Для слоя $\mathbf{y} = W\mathbf{x}$ и скалярного лосса $L$:

$\frac{\partial L}{\partial W} = \frac{\partial L}{\partial \mathbf{y}}\,\mathbf{x}^\top$

Внешнее произведение вектора градиента на вектор входа. Проверка форм: если $\partial L/\partial \mathbf{y}$ это $m \times 1$, а $\mathbf{x}^\top$ это $1 \times n$, произведение даёт $m \times n$ — форму самой $W$, как и требуется.

И симметрично, для передачи градиента ниже по сети:

$\frac{\partial L}{\partial \mathbf{x}} = W^\top \frac{\partial L}{\partial \mathbf{y}}$

Эти две строки — целиком backward для линейного слоя. Всё, что делает nn.Linear.backward, здесь написано.

Как проверять выкладки

Три приёма, которые ловят почти все ошибки:

1. Формы. Каждый промежуточный результат должен иметь осмысленную форму, а финальный градиент — форму параметра. Транспонирование, поставленное не туда, почти всегда ломает размерности.
2. Скалярный случай. Подставьте $n = m = 1$: матрицы станут числами, и формула должна превратиться в знакомое одномерное правило. $2A^\top(A x - b)$ становится $2a(ax-b)$ — верно.
3. Численная проверка. Сравните с центральной разностью. Это то, что делают в реальности, и то, чем мы закроем блок.

Из этих трёх второй недооценивают чаще всего, а он самый быстрый: скалярный случай проверяется в голове.