Нео-РНН или Make RNNs great again
Когда в 2017 году появились трансформеры, популярные до этого RNN обрели слишком серьезного конкурента и отошли на второй план. Трансформеры допускали распараллеливание, а значит — ускоренное обучение, поэтому быстро захватили NLP. Преимущества трансформеров понятны, но с моделированием длинных последовательностей возникают проблемы даже у них. Для RNN это тоже непростая задача из-за исчезающих или взрывающихся градиентов. Но RNN с их линейной зависимостью от масштаба выглядят гораздо привлекательнее квадратичной сложности трансформеров. Идеальным вариантом было бы совместить преимущества тех и других.
RWKV
Ровно это призвана сделать архитектура Receptance-Weighted Key-Value (RWKV). Уже в названии авторы заявляют, что заново изобретают RNN чтобы соответствовать эре трансформеров. RWKV допускает распараллеливание и хорошую масштабируемость. Она состоит из четырех модулей: входной (receptance) вектор принимает прошлую информацию, обучаемый вектор затухания весов (weight), key и value — векторные аналоги K и V в стандартном модуле внимания. На каждом слое скрытое состояние используется чтобы рассчитывать скрытое состояние следующего токена на этом же слое. Получается, что состояния для токена считаются частично параллельно, точнее — по каскаду.
В модели чередуются два слоя — смешения каналов и смешения по времени. На первом токен смешивается с предыдущем выходом предыдущего слоя и не может прочитать собственное состояние с предыдущей итерации. Это аналог кратковременной памяти. На слое смешивания по времени происходит то же самое, но часть предыдущих состояний сохраняется (аналог долговременной памяти). RWKV собрало вокруг себя довольно активное сообщество, которое пробует RWKV на разных задачах.
S4
Другой подход к неоРНН пришел из теории динамических систем. Заключается он в использовании пространства состояний (оно же — фазовое пространство). Вход, выход и скрытое внутреннее состояние описывается системой из двух ОДУ первого порядка. В них входят четыре матрицы — системы, управления, выхода и прямой связи. Матрицы, вообще говоря, находятся градиентным спуском, но таким способом ничего хорошего не получается (возможно потому, что решения ОДУ первого порядка экспоненциальны по характеру). Чтобы решить эту проблему, матрицу системы заменяют на нижнетругольную матрицу специального вида HiPPO. Сама матрица связывает латентное состояние с входом, а ее треугольность обеспечивает, что оно будет помнить историю входа. Затем уравнения переводятся в дискретный эквивалент, в котором текущее внутреннее состояние выражается рекуррентно. Можно провести и свертку. Получается, что считать такую систему можно как РНН. Теоретически выглядит хорошо, но в жизнь не воплотимо — слишком высокая вычислительная сложность из-за вида матрицы. Авторы S4 представили HiPPO матрицу в виде суммы нормальной матрицы и низкоранговой. Результаты получились неплохими на Long Range Arena, даже удалось решить задачу Path-X, с которой не справлялись некоторые трансформеры. В более свежей статье авторы сделали матрицы зависящими от входных данных и таким образом реализовали избирательность модели. Очень подробно об этом написал Гриша Сапунов. Списывать не будем, но с удовольствием посоветуем к прочтению.
HGRN
Авторы из OpenNLPLab вспомнили, что всё новое это хорошо забытое старое и обратили внимание на гейты забывания, которые в свое время сыграли ключевую роль в LSTM. Чтобы сохранять long-term dependencies нужно удерживать значения на выходных гейтах близкими к единице. Но в слишком контрастном режиме (когда значения очень близки к нулю или единице) возникает проблема исчезновения градиентов. Если же все значения близки к единице, то не забудется ненужная информация. Чтобы решить эту проблему авторы HGRN предложили красивый ход. Добавили еще один обучаемый параметр, нижнюю границу (LB). На нижних уровнях LB мал и монотонно возрастает по мере продвижения по уровням, приближаясь к единице. В результате получается, что на нижних уровнях забывается ненужное, а на верхних — запоминается нужное. Звучит изящно даже на гуманитарном уровне.
Ключевой момент в том, что и S4, и RWKV с разных сторон, но по сути реализуют экспоненциальное скользящее среднее, где скорость затухания постоянна, то есть не зависит от состояния. Возврат к гейтам забывания в HGRN делает скорость забывания зависящей от данных.
Результат убедительный — в бенчмарке Long Range Arena NGRN победил и S4 и RWKV:
Больше обзоров AI-статей на канале Pro AI. Подписывайтесь)