🌟 Возвращение RNN: LSTM и GRU — все, что нам было нужно?

Архитектура Transformer доминирует в моделировании последовательностей уже несколько лет, демонстрируя отличные результаты в задачах NLP, машинного перевода и генерации текста. Главный недостаток Transformer — они долго считают длинные последовательности. А если вычислительных ресурсов мало, то реализация занимает либо много времени, либо требует их увеличения.

Авторы исследования предлагают вернуться к RNN, ведь они быстрее считают и параллельно учитывают контекст. Чтобы отвязаться от обратного распространения ошибки (BPTT), которая требует линейного времени обучения, применяется алгоритм параллельного сканирования за счет устранения зависимости от срытых состояний из гейтов LSTM и GRU.

В предлагаемом методе представлены "уменьшенные" LTSM и GRU - minLSTM и minGRU. Они не только обучаются параллельно, но и используют значительно меньше параметров, чем их старшие аналоги.

Минимализм версий достигается следующим образом:

🟢Устранение зависимостей скрытых состояний из гейтов.
В minLSTM и minGRU input, forget и update gate зависят только от входных данных, а не от предыдущих скрытых состояний.

🟢Отказ от ограничения диапазона candidate hidden state.
В традиционных LSTM и GRU функция гиперболического тангенса используется для ограничения диапазона значений скрытых состояний. В minLSTM и minGRU это ограничение снимается.

🟢Неизменность масштаба выходных данных во времени (только для minLSTM).
Для minLSTM выполняется нормализация forget и input гейтов, чтобы гарантировать, что масштаб состояния ячейки не зависит от времени.

Результаты экспериментов:

🟠Время выполнения: minLSTM и minGRU скорость обучения по сравнению с LSTM и GRU, больше в 1361 раз для последовательности длиной 4096;

🟠Задача выборочного копирования: minLSTM и minGRU успешно справились, в отличие от S4, H3 и Hyena;

🟠Обучение с подкреплением на датасете D4RL: minLSTM и minGRU обошли Decision S4 и показали производительность, сопоставимую с Decision Transformer, Aaren и Mamba;

🟠Языковое моделирование: minLSTM, minGRU, Mamba и Transformer показывают одинаковые результаты, но Transformer требует значительно большего количества шагов обучения.

Прикладная реализация численно-устойчивой в логарифмическом пространстве версии метода minGRU на Pytorch представлена в репозитории на Github.

▶️ Локальная установка и запуск minGRU в последовательном и параллельном режиме :

# Install miniGRU-pytorch
pip install minGRU-pytorch

# Usage
import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(512)
x = torch.randn(2, 1024, 512)
out = min_gru(x)
assert x.shape == out.shape

# Sanity check
import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(dim = 512, expansion_factor = 1.5)
x = torch.randn(1, 2048, 512)

# parallel

parallel_out = min_gru(x)[:, -1:]

# sequential

prev_hidden = None
for token in x.unbind(dim = 1):
    sequential_out, prev_hidden = min_gru(token[:, None, :], prev_hidden, return_next_prev_hidden = True)
assert torch.allclose(parallel_out, sequential_out, atol = 1e-4)

📌Лицензирование : MIT License.


🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #RNN #miniGRU

🌟 FlashRNN: оптимизация RNN на современном оборудовании.

FlashRNN - библиотека, которая реализует традиционные RNN, такие как LSTM, GRU и сети Элмана, а также новейшую архитектуру sLSTM в CUDA и Triton.

В отличие от распространенных современных моделей архитектуры Transformers, RNN обладают возможностями отслеживания состояния, оставаясь актуальными для решения задач моделирования временных рядов и логического мышления.

FlashRNN предлагает два варианта оптимизации: чередующийся и объединенный.

🟢Чередующийся позволяет обрабатывать данные с большим размером скрытых состояний и значительно превосходит по скорости базовую реализацию PyTorch.

🟢Объединенный вариант агрегирует операции умножения матриц и вычисления функций в одно ядро, снижая количество обращений к памяти и позволяет хранить рекуррентные матрицы весов непосредственно в регистрах GPU.

За автоматизацию настройки параметров FlashRNN отвечает библиотека ConstrINT, которая решает задачи целочисленного удовлетворения ограничений, моделируя аппаратные ограничения в виде равенств, неравенств и ограничений делимости.

Эксперименты с FlashRNN показали существенное увеличение скорости работы: до 50 раз по сравнению с PyTorch. FlashRNN также позволяет использовать большие размеры скрытых состояний, чем нативная реализация Triton.

▶️ Локальная установка и пример запуска FlashRNN:

# Install FlashRNN
pip install flashrnn


# FlashRNN employs a functional structure, none of the parameters are tied to the `flashrnn` function:

import torch
from flashrnn import flashrnn

device = torch.device('cuda')
dtype = torch.bfloat16
B = 8        # batch size
T = 1024     # sequence length
N = 3        # number of heads
D = 256      # head dimension
G = 4        # number of gates / pre-activations for LSTM example
S = 2        # number of states

Wx = torch.randn([B, T, G, N, D], device=device, dtype=dtype, requires_grad=True)
R = torch.randn([G, N, D, D], device=device, dtype=dtype, requires_grad=True)
b = torch.randn([G, N, D], device=device, dtype=dtype, requires_grad=True)
states_initial = torch.randn([S, B, 1, N, D], device=device, dtype=dtype, requires_grad=True)

# available functions
# lstm, gru, elman, slstm

# available backend
# cuda_fused, cuda, triton and vanilla

states, last_states = flashrnn(Wx, R, b, states=states_initial, function="lstm", backend="cuda_fused")

# for LSTM the hidden h state is the first of [h, c]
# [S, B, T, N, D]
hidden_state = states[0]

📌Лицензирование: NXAI Community License:

🟠бесплатное использование в некоммерческих целях с маркировкой при публикации в отрытых источниках;

🟠получение коммерческой лицензии при годовом доходе свыше 100 млн.евро


🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #RNN #FlashRNN

📌MesaNet: оптимальная адаптация весов в реальном времени.

MesaNet — это новое поколение RNN-архитектур от команды Google Paradigms of Intelligence, созданное для эффективного моделирования длинных последовательностей (например, текста или временных рядов).

Ключевая задача MesaNet - преодолеть главный недостаток трансформеров: линейный рост вычислительных затрат и потребления памяти при увеличении длины последовательности во время инференса. В отличие от классических подходов, MesaNet достигает постоянной сложности на токен за счёт «оптимального обучения в реальном времени» — динамической подстройки внутренних весов под контекст прямо во время работы модели, без хранения всей истории токенов.

Архитектурно, MesaNet построен как стек чередующихся блоков: Mesa-слои (для смешивания информации вдоль последовательности) и MLP (для обработки признаков внутри токена).

Mesa-слой - это сердце системы. Вместо стандартного обновления весов через градиентный спуск (как в Mamba или DeltaNet), он решает оптимизационную задачу для каждого нового токена: ищет матрицу весов, минимизирующую квадратичную ошибку предсказания на всей текущей последовательности.

Для этого используется метод сопряженных градиентов (Conjugate Gradient, CG), который эффективно решает линейную систему из накопленной ковариация ключей, регуляризатора и оптимизированного запроса. Состояние слоя хранится в двух матрицах, которые обновляются через «забывающие» и «входные» гейты, зависящие от данных.

Еще одна, не менее важная опция — динамическое распределение вычислений. Число шагов сопряженного градиента не фиксировано: сложные последовательности требуют больше итераций для сходимости. Это позволяет гибко балансировать точность и скорость.

Сравнение с трансформерами (MHA) и современными RNN (Mamba2, xLSTM, DeltaNet) на синтетике (MAD, RegBench) и языке (SlimPajama) показало: MesaNet сопоставим с трансформерами по perplexity, но выигрывает у других RNN на задачах, требующих длинного контекста. При этом он сохраняет преимущество RNN — постоянные память/вычисления на токен при инференсе.

Интересный паттерн выявили во время тестов: MesaNet, да и просто RNN, точнее предсказывают ранние токены последовательности, а трансформеры - поздние. На длинной экстраполяции (до 32k токенов) MesaNet обошла Mamba2 и xLSTM, но уступила трансформеру.


🔜 Посмотреть видео с докладом про работу.


@ai_machinelearning_big_data

#AI #ML #RNN #MesaNet