В этой статье мы рассмотрим требования к памяти Transformers и связанные с этим проблемы.

Фон

Архитектура Transformer привела в действие последние достижения в области обработки естественного языка (NLP), включая BERT, серию GPT. Для обучения этих чрезмерно параметризованных и требовательных к данным моделей требуются огромные объемы памяти и вычислительные ресурсы. GPT3 имеет параметры 175 млрд и обучается с помощью токенов 300 млрд. Новые исследования продолжают раздвигать границы аппаратного обеспечения. Это приводит к появлению новых платформ ускорителей, таких как DeepSpeed, которые используют параллелизм модели, данных и конвейера для обучения на кластерах из нескольких графических процессоров. Даже в этом случае для простой установки GPT3 в память потребуется 8 DGX-2 (128 GPU) [1]. Поэтому повышение эффективности самой архитектуры Transformer остается активной областью исследований [2]. Многие из этих работ используют разреженность, в частности разреженное внимание для борьбы с квадратичными требованиями к памяти модуля внутреннего внимания. В этой статье мы рассмотрим архитектуру, чтобы понять, зачем нужно столько памяти.

Привлечь внимание к продукту

Интуитивно механизм внутреннего внимания включает в себя изучение того, как каждый токен связан с любым другим токеном во входном контексте. Следовательно, для контекста длины lвнимание вычисляется изO(l²) пар токенов. Мы хотим использовать более длинные последовательности в качестве входных данных для задач, которым требуются более длинные контексты, например, обобщение документов, ответы на вопросы из документов. По мере увеличения длины последовательности это невозможно, например, для ввода 100 000 токенов нам нужно вычислить 10 миллиардов оценок внимания. Предполагая fp32, это 40 ГБ для одной головки блока внимания с несколькими головками.

Как вычисляется самовнимание? Каждый ввод x проецируется в вектор запроса, ключа и значения, q, k, v соответственно. Каждый из этих векторов, q, k и v, построен с помощью линейной проекции входных данных x. например, q = W_q*x. На практике для каждого пакета размером b векторы запросов и ключей для всех входных токенов упаковываются в тензоры Q и K размера (b, l, d). l — входная длина, q и k — d-мерные векторы. Каждый элемент в скалярном произведении QK^T измеряет, насколько близки запрос и ключ друг к другу и, следовательно, сколько внимания нужно уделить k. > при обработке QK^T имеет вид [b, l, l], следовательно, квадратичен по входной длине l. Большое внимание уделяется повышению сложности модуля внутреннего внимания, поскольку он позволяет обрабатывать более длинные последовательности.

Прямая связь

Выходы многоголового собственного внимания поступают в модуль прямой связи. Слой прямой связи включает в себя два умножения матриц и промежуточную активацию ReLU. Матрицы W_1 и W_2 имеют размер d_ff * d. d_ffобычно 4d. d увеличивается с новыми моделями, например, GPT3 использует d = 12288, поэтому требования к памяти уровня прямой связи нетривиальны.

Вывод:

Архитектура трансформатора имеет высокие требования к памяти как в модулях самообслуживания или (перекрестного внимания), так и в сети с прямой связью. Поэтому последние модели требуют обучения на нескольких графических процессорах. Квадратичные требования к памяти затрудняют обработку длинных последовательностей. Исследования по созданию модели «Эффективный трансформатор» продолжаются с особым упором на внимание к себе.

Ссылки

  1. ZeRO-Infinity: преодоление стены памяти графического процессора для глубокого обучения в экстремальных масштабах https://arxiv.org/pdf/2104.07857.pdf
  2. Эффективные трансформаторы https://dl.acm.org/doi/pdf/10.1145/3530811