Что является результатом перемножения ключей и запросов в трансформерах

Что проверяют

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

Как отвечать

Результат это матрица оценок совместимости, своя на каждую голову внимания. Если запросов n, ключей m, а размерность головы d, то Q имеет форму n на d, K такую же, а произведение Q на транспонированную K даёт матрицу n на m. Элемент на пересечении строки i и столбца j это скалярное произведение запроса i и ключа j, то есть ненормированная мера того, насколько эта пара подходит друг другу.

Дальше идут три шага, и пропускать их нельзя:

• Деление на корень из размерности головы. С ростом размерности растёт и разброс скалярного произведения, softmax уходит в насыщение, градиенты почти обнуляются.
• Маска. Каузальная закрывает будущие позиции, маска заполнения закрывает добивку до общей длины. Применяется до softmax, обычно подстановкой большого отрицательного значения.
• Softmax по оси ключей. Только после него получаются веса: неотрицательные и дающие в сумме единицу по строке.

Потом веса умножаются на значения, и выходит взвешенная смесь векторов V.

Полезное следствие: матрица n на m на каждую голову и каждый слой это и есть квадратичная память по длине контекста, из-за которой появились экономные реализации внимания и оконные схемы.

Где спотыкаются

Называют произведение весами внимания. Это логиты: среди них есть отрицательные числа, и сумма по строке ничему не равна, пока не применён softmax.

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

Ставят маску после нормировки. Занулённые веса тогда не складываются в единицу, и на коротких примерах в батче модель учится по искажённому распределению.