LSTM и GRU

Обычная рекуррентная ячейка из предыдущей заметки на каждом шаге переписывает всё своё состояние через умножение на $W_h$, и градиент через $k$ шагов — произведение $k$ таких множителей. Исправление, предложенное Хохрайтером и Шмидхубером (1997) и упрощённое Чо и соавторами (2014), архитектурное: дать ячейке путь, вдоль которого к состоянию прибавляют, а не преобразуют его, и позволить сети выучить — для каждой координаты и на каждом шаге, — сколько сохранить, записать и прочитать.

Обозначения. Как и раньше, $x_t \in \mathbb R^d$ — вход, $h_t \in \mathbb R^H$ — скрытое состояние. LSTM добавляет состояние ячейки (cell state) $c_t \in \mathbb R^H$. $\sigma$ — логистическая функция с областью значений $(0, 1)$; у $\tanh$ область значений $(-1, 1)$; $\odot$ — поэлементное умножение. У каждого гейта (gate) свои веса $W_{\cdot}$ (при $x_t$), $U_{\cdot}$ (при $h_{t-1}$) и смещение $b_{\cdot}$.

Отправная точка: одно состояние, переписываемое на каждом шаге#

$h_t = \tanh(W_x x_t + W_h h_{t-1} + b)$. Состояние вычисляется заново на каждом шаге; ничто в ячейке не может сказать «не трогай координату 7». Переживёт ли старая информация шаг, зависит от $W_h$, общей для всех шагов и всех координат, а множитель градиента на шаг равен $\mathrm{diag}(1 - h_t^2) W_h$.

LSTM#

По $(x_t, h_{t-1})$ вычисляются четыре величины: три гейта в $(0, 1)$ и одно значение-кандидат в $(-1, 1)$:

$$\begin{aligned} f_t &= \sigma(W_f x_t + U_f h_{t-1} + b_f) && \text{гейт забывания: сколько от } c_{t-1} \text{ сохранить},\\ i_t &= \sigma(W_i x_t + U_i h_{t-1} + b_i) && \text{входной гейт: сколько от кандидата записать},\\ \tilde c_t &= \tanh(W_c x_t + U_c h_{t-1} + b_c) && \text{кандидат: что записать},\\ o_t &= \sigma(W_o x_t + U_o h_{t-1} + b_o) && \text{выходной гейт: сколько памяти показать}. \end{aligned}$$

Затем обновляется память, и из неё читается рабочее состояние:

$$c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t, \qquad h_t = o_t \odot \tanh(c_t).$$

Каждая операция покоординатная. Координата $c$ с $f \approx 1$ и $i \approx 0$ копируется без изменений столько шагов, сколько гейты остаются такими; координата с $f \approx 0$ и $i \approx 1$ переписывается. Область значений сигмоиды делает гейты долями, а область значений $\tanh$ ограничивает кандидата и показываемое значение, поэтому $c_t$ может расти только повторным сложением ограниченных величин, а не геометрически.

Разобранный пример#

Одна координата, один шаг. Предыдущая память $c_{t-1} = 2$, предыдущее состояние $h_{t-1} = 1$, вход $x_t = 1$. Возьмём скалярные веса (условные значения), так что каждая предактивация — небольшая сумма:

Гейт	Предактивация	Значение
забывания $f$	$1.63 \cdot 1 + 2.70 \cdot 1 + 1.62 = 5.95$	$\sigma(5.95) = 0.997$
входной $i$	$1.65 \cdot 1 + 2.00 \cdot 1 + 0.62 = 4.27$	$\sigma(4.27) = 0.986$
кандидат $\tilde c$	$0.94 \cdot 1 + 1.41 \cdot 1 - 0.32 = 2.03$	$\tanh(2.03) = 0.966$
выходной $o$	$-0.19 \cdot 1 + 4.38 \cdot 1 + 0.59 = 4.78$	$\sigma(4.78) = 0.992$

Обновление памяти: $c_t = 0.997 \times 2 + 0.986 \times 0.966 = 1.995 + 0.953 = 2.948$. Показываемое состояние: $h_t = 0.992 \times \tanh(2.948) = 0.992 \times 0.9945 = 0.986$.

Старая память почти полностью сохранилась (гейт забывания около 1), и к ней прибавилось новое значение около $0.95$. Если бы гейт забывания оказался равен $0.1$, память упала бы до $0.2 + 0.95$: тот же вход, другое решение, которое принимают веса, видящие $x_t$ и $h_{t-1}$.

Шаг за шагом#

Рисунок

Шаг LSTM: что куда идёт

По номерам на рисунке:

1. Гейт забывания. $f = \sigma(1.63 \cdot 1 + 2.70 \cdot 1 + 1.62) = 0.997$: какую долю старой памяти сохранить, покоординатно.
2. Входной гейт и кандидат. $i = 0.986$ говорит, сколько записать, $\tilde c = 0.966$ — что записать. Кандидат — это обновление обычной RNN из предыдущей заметки со своими весами.
3. Сохранить. Старая память умножается на гейт забывания: $0.997 \times 2 = 1.995$.
4. Записать. Прибавляется кандидат, умноженный на гейт: $0.986 \times 0.966 = 0.953$, и $c_t = 2.948$. Шаги 3 и 4 — единственные операции на линии памяти: умножение на число из $(0, 1)$ и сложение.
5. Выходной гейт. $o = 0.992$: какую долю памяти показать наружу.
6. Прочитать. $h_t = o \cdot \tanh(c_t) = 0.992 \times 0.9945 = 0.986$. И $c_t$, и $h_t$ уходят на шаг $t + 1$: гейты следующего шага читают $h_t$, а линия памяти несёт $c_t$.

По сравнению с рисунком RNN из предыдущей заметки входы и рабочее состояние $h_t$ нарисованы так же; новое — линия памяти сверху в янтарной рамке и три гейта, которые на неё действуют.

Почему градиент выживает#

Продифференцируем обновление памяти по предыдущей памяти:

$$\frac{\partial c_t}{\partial c_{t-1}} = \mathrm{diag}(f_t) + (\text{члены через } f_t, i_t, \tilde c_t \text{, зависящие от } h_{t-1}).$$

Первое слагаемое — сам гейт забывания. Вдоль линии памяти градиент через $k$ шагов содержит произведение $\prod f_s$, и сеть может держать его около 1, держа гейты забывания около 1 там, где нужна память, тогда как множитель обычной RNN $\mathrm{diag}(1 - h^2) W_h$ — степень фиксированной матрицы. Это не доказывает, что LSTM выучивает сколь угодно длинные зависимости (остальные слагаемые по-прежнему затухают или взрываются, а гейты насыщаются), но убирает геометрическое затухание, из-за которого обычная ячейка проваливалась на десятках шагов; на практике LSTM выучивает зависимости на сотни токенов. Смещение $b_f$ часто инициализируют единицей, чтобы гейты забывания начинали открытыми (Jozefowicz, Zaremba and Sutskever, 2015).

Цена#

У каждого из четырёх блоков есть матрица $H \times d$, матрица $H \times H$ и смещение, поэтому у ячейки LSTM $4(Hd + H^2 + H)$ параметров — вчетверо больше, чем у обычной ячейки (PyTorch хранит по два вектора смещений на блок, что даёт $4(Hd + H^2 + 2H)$). Модель перевода из заметки о seq2seq использует $d = H = 128$: $4(128 \cdot 128 + 128 \cdot 128 + 2 \cdot 128) = 132\,096$ параметров на направление. Вычисление на шаге — тоже четыре матричных произведения вместо одного, и оно по-прежнему строго последовательно.

GRU#

Рекуррентная ячейка с гейтами GRU (gated recurrent unit; Cho et al., 2014) объединяет память и рабочее состояние и использует два гейта:

$$\begin{aligned} z_t &= \sigma(W_z x_t + U_z h_{t-1} + b_z) && \text{гейт обновления},\\ r_t &= \sigma(W_r x_t + U_r h_{t-1} + b_r) && \text{гейт сброса},\\ \tilde h_t &= \tanh\big(W_h x_t + U_h (r_t \odot h_{t-1}) + b_h\big) && \text{кандидат},\\ h_t &= (1 - z_t) \odot h_{t-1} + z_t \odot \tilde h_t . \end{aligned}$$

Гейт обновления (update gate) играет роли и $f$, и $i$, связанные как $f = 1 - z$: что не сохранено, то переписано. Гейт сброса (reset gate) управляет тем, сколько предыдущего состояния может прочитать кандидат. Выходного гейта нет; $h_t$ и есть память. Число параметров — $3(Hd + H^2 + H)$, три четверти LSTM того же размера, и на шаге на одно матричное произведение меньше. Классификатор AG News из предыдущей заметки — это GRU с $d = 100$ и $H = 128$.

Соглашения. В уравнениях выше $z_t$ — доля состояния, которая переписывается. Исходная статья (Cho et al., 2014, уравнения 5–8) и nn.GRU в PyTorch интерполируют в обратную сторону, $h_t = z_t \odot h_{t-1} + (1 - z_t) \odot \tilde h_t$, так что их $z_t$ — это наше $1 - z_t$. Кроме того, PyTorch применяет сброс после рекуррентного линейного отображения, $\tilde h_t = \tanh\big(W_h x_t + b_h + r_t \odot (U_h h_{t-1} + b'_h)\big)$, тогда как уравнения выше, как и у Cho et al., сбрасывают $h_{t-1}$ до умножения на $U_h$; в общем случае кандидаты получаются разными. Со вторым вектором смещений на каждый гейт число параметров в PyTorch — $3(Hd + H^2 + 2H)$. Оба варианта — стандартные GRU, но формула и реализация должны использовать один и тот же.

	Обычная RNN	LSTM	GRU
Состояния	$h$	$h, c$	$h$
Гейты	0	3 ($f, i, o$)	2 ($z, r$)
Обновление памяти	переписывание	$f \odot c + i \odot \tilde c$	$(1 - z) \odot h + z \odot \tilde h$
Параметры	$Hd + H^2 + H$	$4\times$	$3\times$
Множитель градиента вдоль памяти	$\mathrm{diag}(1-h^2)W_h$	$\approx \mathrm{diag}(f)$	$\approx \mathrm{diag}(1 - z)$

Какая ячейка лучше — вопрос эмпирический и зависит от задачи. Прямые сравнения находят их близкими (Chung et al., 2014; Jozefowicz, Zaremba and Sutskever, 2015), а большое исследование вариантов LSTM показало, что связывание входного гейта и гейта забывания, как в GRU, почти ничего не стоит (Greff et al., 2017).

Чтение в обе стороны#

Рекуррентное состояние в позиции $t$ подытоживает $x_{\le t}$ и ничего не знает о $x_{> t}$. Для классификатора или энкодера, который видит весь вход сразу, можно запустить вторую ячейку справа налево и конкатенировать два состояния в каждой позиции: двунаправленная RNN даёт $[\overrightarrow{h}_t; \overleftarrow{h}_t] \in \mathbb R^{2H}$ — представление каждого токена с учётом всего контекста — ценой удвоенных вычислений. Энкодер из следующей заметки — двунаправленная LSTM. Языковая модель не может быть двунаправленной, потому что не должна видеть токен, который предсказывает; заметка о BERT находит другой способ использовать оба направления.

Что изменилось#

	Обычная RNN	RNN с гейтами
Может оставить координату без изменений	Только если так вышло у $W_h$	Да, гейтами
Длинные зависимости	Десятки токенов	Сотни, на практике
Параметров на ячейку	$1\times$	$3\times$ (GRU), $4\times$ (LSTM)
Последовательная	Да	Да
Состояние фиксированного размера	Да	Да

Ограничения#

• По-прежнему последовательна. Шаг $t$ зависит от шага $t - 1$. Время обучения и инференса растёт с $T$, и измерение последовательности нельзя распараллелить.
• По-прежнему один вектор. Как бы хорошо оно ни помнило, состояние в конце — это $H$ (или $2H$) чисел, которые должны вместить всё, что нужно следующему этапу. Заметка о seq2seq делает этот провал видимым на переводе.
• Гейты насыщаются. У сигмоид около 0 или 1 малые производные; гейт, который научился закрываться, медленно открывается снова.
• Ёмкость стоит денег. Четыре матрицы на шаг; у двухслойной двунаправленной LSTM с $H = 512$ над 512-мерными входами 10.5 млн параметров ещё до всяких эмбеддингов.

Что нужно запомнить#

• LSTM: $c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t$, $h_t = o_t \odot \tanh(c_t)$; гейты в $(0, 1)$ от $\sigma$, кандидат в $(-1, 1)$ от $\tanh$.
• Аддитивное обновление памяти делает множитель градиента вдоль $c$ примерно равным гейту забывания, который сеть может держать около 1.
• GRU: два гейта, одно состояние, $h_t = (1 - z_t)\odot h_{t-1} + z_t \odot \tilde h_t$; три четверти параметров LSTM.
• Двунаправленные ячейки дают каждой позиции полный контекст; для предсказания следующего токена их использовать нельзя.
• Гейты исправляют градиент, но не последовательность вычислений и не сводку в один вектор.

Ноутбук#

week04_lstm_attention_transformer.ipynb. Номера ячеек считают все ячейки, включая текстовые, начиная с 1. Приведённые числа взяты из прогона, сохранённого в курсе; у ячеек, изменённых с тех пор под текущие версии библиотек, сохранённого выхода нет.

• Ячейка 39 определяет Seq2SeqLSTM с nn.LSTM(128, 128, bidirectional=True) в роли энкодера и однонаправленной nn.LSTM(128, 128) в роли декодера; два последних скрытых состояния и два состояния ячейки энкодера (h_enc, c_enc формы [2, B, 128]) конкатенируются в [1, B, 256] и проецируются обратно в 128 слоями h_proj и c_proj, чтобы инициализировать декодер.
• Ячейка 44 печатает всего 8 380 474 параметра и 492 602 вне матриц эмбеддингов и выхода; последнее — это $2 \times 132\,096$ (двунаправленный энкодер) $+ 132\,096$ (декодер) $+ 65\,792$ (две проекции) $+ 30\,522$ (смещение выходного слоя, которое вычитание в ноутбуке оставляет).

week02_NN_text_classification.ipynb, ячейка 40, — упомянутый выше классификатор на GRU. Он использует nn.GRU, так что его гейты следуют соглашению PyTorch из раздела о GRU: гейт обновления — доля сохраняемого, а сброс действует после рекуррентной матрицы.

Дальше#

С ячейкой, которая умеет помнить, курс берётся за первую задачу порождения с двумя последовательностями: перевести немецкое предложение на английский. Одна рекуррентная сеть читает источник, другая пишет перевод. Следующая заметка строит этот энкодер–декодер, измеряет его метрикой BLEU и обнаруживает, что единственный вектор, передаваемый от энкодера декодеру, теперь и есть ограничивающий фактор.

Литература#

• S. Hochreiter, J. Schmidhuber. Long Short-Term Memory. Neural Computation, 9(8), 1997.
• K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, Y. Bengio. Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation. EMNLP, 2014. GRU.
• J. Chung, C. Gulcehre, K. Cho, Y. Bengio. Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. NIPS Deep Learning Workshop, 2014. Сравнение LSTM и GRU на моделировании музыки и речи.
• R. Jozefowicz, W. Zaremba, I. Sutskever. An Empirical Exploration of Recurrent Network Architectures. ICML, 2015. Инициализация смещения гейта забывания; сравнение LSTM и GRU.
• K. Greff, R. K. Srivastava, J. Koutník, B. R. Steunebrink, J. Schmidhuber. LSTM: A Search Space Odyssey. IEEE TNNLS, 28(10), 2017.
• C. Olah. Understanding LSTM Networks. 2015. Разбор по гейтам, которому следует разобранный пример.
• M. Schuster, K. K. Paliwal. Bidirectional recurrent neural networks. IEEE Transactions on Signal Processing, 45(11), 1997.