Were RNNs All We Needed?
https://arxiv.org/abs/2410.01201

Revisits RNNs and shows that by removing the hidden states from input, forget, and update gates RNNs can be efficiently trained in parallel.

This is possible because with this change architectures like LSTMs and GRUs no longer require backpropagate through time (BPTT).

They introduce minLSTMs and minGRUs that are 175x faster for a 512 sequence length.

Paper quote: "We show that these stripped-down versions of decade-old RNNs match the empirical performance of recent sequence models."

Ист. X

Проверил что выдаёт необученный модуль внимания для последнего токена. Находит ли он закономерности в промпте, если их туда специально положить.

Промпт: " ABCD ABCDABCE"

Здесь повторяющаяся два раза подстрока ABCD и частично схожая с ним ABCE. E - последний токен, перед которым находится подстрока-ключ ABC.

Разумно предположить, что из-за совпадающего ключа вектор последнего токена после обработки блоком внимания должен быть схож с векторами токенов D.

X - входная матрица векторов токенов, n x d, n - длина контекста/промпта (15), d - размерность векторов (1000). Исходные векторы токенов полуортогональны.

Так как матрица внимания (key-query матрица) не обучена, и никакого её обучения не предполагается, то она отобразит входные полуортогональные векторы в такие же полуортогональные, а softmax для них, соответственно, вернёт почти равные скоры.

Поэтому выбрасываем эту матрицу, а коэффициенты софтмакса заменяем на константное значение равное 1 / m, m - длина контекста для текущего токена (размер окна в треугольной маске).

Т.е. внимание в урезанном виде считает средний вектор от входа.

torch.sum(X[0:i], dim=0) / i)

И тут делаем ещё шаг от ванильного внимания, делаем его рекуррентным, но без RNN. Посчитали вектор внимания-контекста для текущего токена, через резконнект прибавили его к своему входу и не только передали в следующий слой, но и заменили им своё значение. На каждой итерации одного слоя вдоль промпта происходит обновление вектора текущего токена.

X[0] = X[0]
X[1] = X[1] + X[0] / 1
X[2] = X[2] + torch.sum(X[0:2], dim=0) / 2
X[3] = X[3] + torch.sum(X[0:3], dim=0) / 3
X[4] = X[4] + torch.sum(X[0:4], dim=0) / 4

Делаем полтора десятка таких слоёв и послойно смотрим схожесть последнего токена с предыдущими.

for p in range(15):
  for i in range(1, len(X)):
    X[i] = X[i] + torch.sum(X[0:i], dim=0) / i

Получается картинка градиента.

Ранее я говорил, что резконнект напоминает мне узел сравнения в автомате. А то, что возвращает внимание - это предикт от контекста по обратной отрицательной связи. И был бы это автомат, то суммировав эти "сигналы", получили бы ошибку предсказания или отклонение.

Вычитаем выход внимания и смотрим как меняется картинка.

for p in range(15):
  for i in range(1, len(X)):
    X[i] = X[i] - torch.sum(X[0:i], dim=0) / i

Нарисовалась какая-то структура. Эмерджентно появилась подсветка токенов D, то, что и требовалось.

А теперь шлифанём выход нормализацией, не зря же её пихают везде и регистрируют улучшение. Возьмём популярный сейчас RMS.

for p in range(15):
  for i in range(1, len(X)):
    X[i] = RMSnorm(X[i] - torch.sum(X[0:i], dim=0) / i)

Идеально!

Даже странный выброс на A вытерся. Повторяющаяся структура в данных "самозакодировалась" в эмбеддингах.

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

Таким образом здесь в следующий слой отправляется нормализованная ошибка предсказания. Рассогласование между фактическим текущим токеном-вектором и токеном-вектором, предсказанным по контексту.

G1 от Unitree

What Matters in Transformers? Not All Attention is Needed
https://arxiv.org/abs/2406.15786

Из Llama 2 70B выкинули половину слоёв внимания, ускорили вывод в два раза и потеряли в качестве 2,4%.

А вот с FFN такое не прошло.

- Removing entire transformer blocks leads to significant performance degradation.
- Removing MLP layers results in significant performance degradation.
- Removing attention layers causes almost no performance degradation!

ист X

Думаю, основную фичу трансформера я уловил.

Вчера подумал, что преобразования вектора X умножением на матрицу W можно заменить на его сложение с вектором f(X).

X•W = X + f(X)

Это резконнект с функцией слоя f(X).

Значит матрицу внимания W (см. рисунок) можно заменить на сумму входа X и некоторую функцию f(X) от этого входа.

Функция возвращает вектор, который складывается с входным (от текущего последнего токена) и становится вектором запроса q. Дальше без изменений.

q = X + f(X)

см. картинку

f(X) возвращает вектор ассоциаций для входа X. Самый простой вариант её реализации - ассоциативная память из NTM (нейронной машины). Матрица ключей и матрица значений. Сравниваем вектор запроса с ключами, получаем оценки схожести и взвешенным суммированием собираем из матрицы значений ответ. Это то же самое внимание.

Для небольшого эксперимента ключи и значения можно просто задать вручную. Мы в них формируем граф отношений между токенами.

Собрал эту схему.

По мотивам недавней статьи про местоимения придумал тестовую задачу.

Два варианта:

МАША ЕСТ КАШУ ПОТОМУ ЧТО ОНА ВКУСНАЯ

МАША ЕСТ КАШУ ПОТОМУ ЧТО ОНА ГОЛОДНАЯ

Местоимение ОНА в первом случае должно быть релевантно КАШЕ, во втором - МАШЕ.

Граф заданных отношений:

ОНА = МАША + КАШУ
ВКУСНАЯ = КАШУ
ГОЛОДНАЯ = МАША
МАША = ЕСТ
ЕСТ = КАШУ
ПОТОМУ = ЧТО

Чтобы найти логиты для ОНА, добавил её в конец промпта. Два слоя, измерение для зондового токена ОНА было сделано на третьем.

Логиты для первого предложения:

МАША ЕСТ КАШУ ПОТОМУ ЧТО ОНА ВКУСНАЯ *ОНА

МАША      0.42
ЕСТ      -0.1
КАШУ      1.74
ПОТОМУ    0.1
ЧТО       0.01
ОНА      -0.03
ВКУСНАЯ   0.98
ГОЛОДНАЯ -0.13

ОНА ассоциируется с КАШУ (ВКУСНАЯ).


Логиты для второго предложения:

МАША ЕСТ КАШУ ПОТОМУ ЧТО ОНА ГОЛОДНАЯ *ОНА

МАША      1.88
ЕСТ      -0.1
КАШУ      0.06
ПОТОМУ   -0.07
ЧТО       0.05
ОНА       0.75
ВКУСНАЯ   0.11
ГОЛОДНАЯ  0.17

ОНА ассоциируется с МАША (ОНА).

Вот ещё часть инфы по второму предложению.

Аттеншен для ГОЛОДНАЯ во втором слое, это МАША и ОНА.

ATT [0.75 0. 0. 0. 0. 0.25]

И аттеншен для зонда *ОНА в третьем слое. Всё внимание на ОНА ГОЛОДНАЯ.

ATT [0. 0. 0. 0. 0. 0.61 0.39]

Внимание - это билатеральный фильтр

Обычный билатеральный фильтр (БФ) берёт взвешенную сумму вокруг центрального, текущего пикселя, опираясь на близость соседних ("контекстных") пикселей по цвету. БФ сравнивает скалярные значения цвета пикселя, а внимание делает сравнение запроса через скалярное произведение с ключами. В графике у БФ двумерное симметричное ядро, у внимания одномерное и ассиметричное.

началось

Походка всё лучше.

Записал функции внимание и MLP через вызовы одной единственной функции - чтение из ассоциативной памяти - mem(q, K, V).
* напоминаю, что это не та какава (QKV), что в формулах трансформера

Сама функция чтения mem, как говорил, - это функция внимания, она же ассоциативная память нейронной машины.

def mem(q, K, V, temp): 
  return softmax(q @ K.T / temp) @ V

Возвращаемся к рисунку схемы внимания.

Начальное значение вектора запроса q равно вектору последнего токена (входа).

q = X[-1]

Запрос отправляется в постоянную ассоциативную память внимания. Результат чтения из памяти r_ass дополняет через skip connection запрос q.

r_ass = mem(q, K, V)

K и V - матрицы ключей и значений релевантных/ассоциативных связей между токенами.

q = q + r_ass

Расширенный ассоциациями запрос теперь отправляется в оперативную, динамическую контекстную память, которой является контекст/промпт.

У контекстной памяти ключи и значения равны входу.

Аналогично делаем запрос к памяти, ответ прибавляем к запросу.

r_ctx = mem(q, X[0:-1], X[0:-1])

q = q + r_сtx

В этом месте q это выход слоя внимания. Он поступает на вход MLP, который реализован опять же через чтение постоянной ассоциативной памяти, но с рокировкой ключей и значений (см. выше для чего это сделано).

r_mlp = mem(q, V, K)

q = q + r_mlp

Всё, прошли оба слоя: attention и MLP. Далее по необходимости повторяем.

Всё вместе:

q = X[-1]
q += mem(q, K, V)
q += mem(q, X[0:-1], X[0:-1])
q += mem(q, V, K)

Только чтение из двух типов/банков памяти и ничего более.