group-telegram.com/gonzo_ML/1753
Last Update:
Retentive Network: A Successor to Transformer for Large Language Models
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei
Статья: https://arxiv.org/abs/2307.08621
Код: https://github.com/microsoft/unilm/tree/master/retnet (https://github.com/microsoft/torchscale/blob/main/examples/fairseq/models/retnet.py)
Очередные новости гибридизации в нашем вестнике сельского хозяйства.
Microsoft Research совместно с Tsinghua University предложили новую архитектуру под названием Retentive Network (RetNet).
Все хотят эффективный параллелизм при обучении, O(1) инференс и, конечно, хороший перформанс. Выберите любые два: у рекуррентных сетей традиционно не было параллелизма, у классических трансформеров дешёвого инференса, а у линейных трансформеров -- хорошего качества. Это всё, конечно, с поправкой на современные модели типа S4, RWKV, LRU, но авторы считают, что они все таки где-то не дотягивают и однозначного победителя трансформеров нету. Но теперь типа его придумали.
В чём суть?
RetNet состоит из стека L блоков с residual connection и pre-LayerNorm, как и трансформер. Внутри каждого RetNet блока есть блочок Multi-Scale Retention (MSR) и блочок FFN. Вычисления выглядят классически для трансформера:
Y^l = MSR(LN(X^l)) + X^l
X^{l+1} = FFN(LN(Y^l)) + Y^l,
где FFN(X) = gelu(XW_1)W_2
То есть MSR это замена MHSA (Multi-head Self Attention).
Вход x=x_1, …, x_n RetNet обрабатывает авторегрессионно. Входные векторы x сначала эмбеддятся в X^0 размерности |x|×d_model, где d_model -- это hidden dimension, а затем в каждом слое l из L всего происходит вычисление контекстуализированных репрезентаций X^l = RetNet_l(X^{l−1}). На этом уровне от трансформера отличий нет, все отличия внутри MSR.
Собственно на смену механизму Attention приходит механизм Retention. Жду продолжения рифм. Механизм Retention имеет форму как параллельную, так и рекуррентную, то есть можно обучать в параллельной, а исполнять в рекуррентной.
Входная последовательность X (размерности |x|×d_model) проецируется в v_n = X_n · w_V, а моделирование последовательности является отображением входа v_n в выход o_n через скрытые состояния s_n. В итоге маппинг можно описать рекуррентностью:
s_n = As_{n−1} + K^⊺_n v_n
o_n = Q_n s_n = sum_{m=1}^{n} Q_n A^{n−m} K^⊺_m v_m
где A -- матрица d×d, K и Q -- векторы 1×d.
Проекции Q и K контекстно-зависимы Q = XW_Q, K = XW_K, где W_Q, W_K -- обучаемые матрицы размерности d×d.
Матрица A диагонализируется (снова через комплексные числа как в LRU, https://www.group-telegram.com/ru/gonzo_ML.com/1734):
A = Λ(γe^{iθ})Λ^{−1} и выражение для o_n переписывается так, что Λ отправляются в матрицы W_Q, W_K и после преобразований получается сумма входов, взвешенных с относительными позиционными эмбеддингами. Формулы лучше смотреть на картинке, чем тут текстом парсить.
В итоге в параллельной формулировке механизм Retention выглядит так:
Q = (XW_Q) ⊙ Θ
K = (XW_K) ⊙ conjugate(Θ)
V = XW_V
Θ_n = e^{inθ} (позиционные эмбеддинги типа xPos из Lex Transformer, https://arxiv.org/abs/2212.10554)
/γ^{n−m}, n ≥ m
D_{nm} = { (causal masking and exponential decay)
\0, n < m
Retention(X) = (QK^⊺ ⊙ D)V
Ну то есть в целом весьма похоже на обычное внимание. Ушёл softmax, добавили xPos, появилась рекуррентная формулировка.
В рекуррентной формулировке это записывается как
S_n = γS_{n−1} + K^⊺ V_n
Retention(X_n) = Q_n S_n, n = 1, · · · , |x|
Есть ещё гибридная форма Chunkwise Recurrent Representation для длинных последовательностей, когда они разбиваются на чанки.
Это был одиночный Retention. Далее идёт Gated Multi-Scale Retention, это аналог многоголовости трансформера, когда каждая голова Retention работает по своему кусочку пространства размерности d из полного d_model. У каждой головы свои матрицы W_Q, W_K, W_V и у каждой головы свой параметр γ, который про экспоненциальное затухание. В работе эти параметры выставляли одинаковым образом у разных слоёв.
Итоговый механизм выглядит так:
BY gonzo-обзоры ML статей
Share with your friend now:
group-telegram.com/gonzo_ML/1753