Ivan

DenseAttention: No-Compromise Exact All NxN Interactions Algorithm with O(N) Space and Time Complexity

2024-11-03T13:22:42.387Z

Возможны ли нейросети без нелинейностей? Возможно ли сделать сеть только из матричных умножений - наиболее эффективных по вычислениям и с возможностью параллелизма? И самое главное - не потерять при этом точность работы трансформера. В этой статье показывается, что это возможно. Определив и выкинув наиболее слабые части архитектуры, автор заменяет их матричными умножениями, а где-то вводит новые преобразования для повышения эффективности модели. В результате получается DenseAttention - структура с повышенной точностью и эффективностью вычислений.

Source: Architecture's author - Andrew Argatkiny, DenseAttention paper, DenseAttention Github, VK Lab Meeting

Motivation

Основным минусом нелинейностей является их неэффективность в вычислениях. Например, метрика Model FLOPS Utilization (MFU), которая является отношением наблюдаемой пропускной способности к теоретической максимальной пропускной способности, если бы модель работала с пиковым значением FLOPS без накладных расходов на память или связь, довольно низка в современных архитектурах:

MosaicBERT - 40%

PaLM - 46%

FlashAttention 2 - 72%

Так происходит, потому что GPUs не производят никаких вычислений, пока считывается/записывается память. В статье Data Movement Is All You Need показано, что матричные умножения в модели BERT-large составляют 99.8% всех вычислений (FLOPS), но они занимают лишь 61% времени вычисления. А 31% времени тратится на вычисление оставшихся операций (которые составляют 0.02% FLOPS). То есть трансформер, из-за ограниченности памяти, вычислительно неэффективен.

Углубляясь в проблему, приведу метрику Arithmetic Intensity (ArIn) - отношение общего количества операций FLOPS к общему количеству перемещений данных (байт):

Arithmetic Intensity = FLOPS / Bytes

Для эффективности алгоритма необходимо (но не достаточно), чтобы его значение ArIn было выше, чем ArIn ускорителя. Иначе часть времени ускоритель будет простаивать, что и было показано в статье выше.

Чему же равны метрики ArIn современных ускорителей и вычислительных операций, применяющихся в трансформере? У NVIDIA A100 этот показатель равен 156 FLOPS/B, тогда как в трансформере мы имеет следующие значения:

ReLU activation - 0.25 FLOPS/B
Element-wise - 1/3 FLOPS/B
Layer normalization & Softmax < 10 FLOPS/B

То есть мы видим разницу ArIn минимум в 3-4 порядка. Это оказывает колоссальное влияние на время работы трансформера. Добавлю, что эти нематричные операции выполняются не на тензорных ядрах, а на обычных, что также снижает их эффективность.

Однако и с матричными вычислениями не все в порядке. Основная операция Attention - softmax(Q *K^T)*V - имеет 32 FLOPS/B.

Можно ли заменить эти операции и даже избавиться от них и создать новую, более эффективную архитектуру?

Designing DenseAttention

Source

Для начала автор удаляет некоторые составляющие:

Dropouts - на этапе pre-train они не нужны, но их можно добавить в этап fine-tune
Masking - можно убрать с энкодера (и с декодера тоже можно)
Scale - его можно перенести
Softmax
W_keys, W_values, W_output
Между слоями Attention и FNN убираем LayerNorm и skip-connections

Самой большой проблемой является softmax - без него выходы Attention неограничены, возрастают в полиномиальной степени и в итоге сходятся либо к нулю, либо к бесконечности.

Давайте попробуем разобраться почему это происходит. Рассмотрим упрощенную матрицу attention

Y = X * W * X^T * X

Стандартное отклонение каждого элемента матрицы Y ограничено снизу, но не ограничено сверху. Даже если матрицы X и W независимо распределенные, проблема остается - мы не знаем точное распределение X. А это распределение получается с толстыми и тяжелыми хвостами -> на каждом новом слое оно уходит в бесконечность, а значит понять распределение Y невозможно.В этом случае LayerNorm не помогает, потому что опирается на L2-норму.

Давайте попробуем поменять норму. Возьмем бесконечную норму - модуль максимального значения этой матрицы:

||X|| = max(|X_ij|)

Для этой нормы мы можем вывести такие условия, при которых выход attention будет ограничен. Введем для исследования матрицу Z, которая будет произведением трех матриц X:

Z = X * X^T * X

Тогда, если бесконечная норма матрицы Z ограничена, то и выход attention будет ограничен.

В статье приводится детальное доказательство этого факта, основанного на ограничении дисперсии произведения матрицы X и W.

Вводя новый scale factor, равный 1/N^(1/3), норма матрицы будет ограничена сверху размерностью эмбеддинга. Тем самым мы полностью можем избавиться от softmax без потери качества работы алгоритма.

Тогда введем новую операцию - MaxNormActivation:

Такая норма не центрирована, в ней нет bias и нет никаких весов.

Введя такой трюк, мы получаем большую эффективность - без softmax мы получаем ассоциативность матричных умножений:

(Q * K^T) V = Q (K^T * V)

То есть теперь мы можем варьировать нашу вычислительную сложность в зависимости от размера датасета и эмбеддинга. Но в любом случае наш алгоритм будет работать намного быстрее, чем раньше.

Source

Также в DenseAttention автор уменьшает количество голов в архитектуре - рассматривается два варианта: либо одна большая, либо 4 маленьких. Таким образом получается выигрыш по вычислениям и точности модели. Так, используя только одну голову с размерностью d=1024 для модели BERTLarge операции умножения матриц уже составляют 205 FLOPS/B, против 32.

Продолжая тематику удаления вычислений из трансформера, автор удаляет матрицу keys - W_keys. В стандартном механизме attention каждый раз происходит перемножение двух низкоранговых матриц - queries и keys. Они низкоранговые, потому что имеют размерность d эмбеддинга / d головы. В DenseAttention ранг выходной матрицы гораздо выше, а значит операция перемножения ее на низкоранговую избыточна. Поэтому мы можем удалить матрицу keys для экономии ресурсов. Матрицы W_values, W_outputудаляются по той же причине.

Последним удалением в архитектуре является LayerNormи skip-connections. Эмпирически показано, что их удаление не ведет к уменьшению метрик модели.

Теперь в нашей архитектуре нелинейности остались лишь в конце блока attention и FFN!

Финальная архитектура DenseAttention имеет вид:

Source

В формулах для общего случая нескольких голов это выглядит так:

Благодаря этому, обновленная архитектура attention имеет вычислительную сложность 11Nd^2 в случае O(N) и 9Nd^2 + 2dN^2 в случае O(N^2), что вычислительно превосходит стандартную архитектуру, особенно на длинных последовательностях.

Cosine RelPE

Помимо повышения эффективности вычислений, автор вводит новую функцию positional encoding. Дело в том, что современные модели чаще всего используют Rotary Positional Embeddings (RoPE), который применяет преобразования к матрицам Q и K. Однако в работе RoFormer: Enhanced Transformer with Rotary Position Embedding авторы показали, что параметризация, используемая в RoPE приводит к долгосрочному снижению нормы выхода attention. Более того, преобразования RoPE неэффективны в вычислительном отношении, поскольку их вычисление требует дорогостоящих изменений структуры тензора и нескольких поэлементных операций с низкой ArIn, отдельно для Q и K.

Автор вводит новое преобразование g1, которое вычислительно более эффективно, ведь оно допускает только одно element-wise умножение, вместо двух:

Тогда новое преобразование Cosine RelPE:

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

LocalAttention

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

Автор также вводит LocalAttention, однако делает это не ради повышения эффективности вычисления, а ради увеличения качества модели. LocalAttention в данном случае состоит из трех уровней: LocalAttention, ShiftedLocalAttention и global DenseAttention.

LocalAttention классически разбивает последовательность на окна равного размера w, но в таком случае половина информации теряется. Поэтому вводится ShiftedLocalAttention, который смещен на w/2 относительно первого, что позволяет всем токенам иметь симметричное соседство после двух последовательных слоев. Последний слой global DenseAttention охватывает весь контекст последовательности. Все 3 слоя могут соединяться вместе, как обычные слои трансформера.

Experiments

Как обычно, по экспериментам пройдусь быстро.

Long Range Arena - это сложный набор из 6 классификационных тестов, предназначенных для изучения возможностей эффективных моделей с длительным контекстом на больших последовательностях длиной от 1к до 16к.

BERT pre-train с размерностью модели d = 1024 на тех же датасетах - Wikipedia и BookCorpus. Отмечу, что BERT-large был увеличен с 24 до 32 слоев, чтобы сохранять одинаковое количество параметров.

SpeedTest алгоритмов DenseAttention, standard BERT и FlashAttention-2:

На этом все! Спасибо, что дочитали до конца :)

Linear Quantization

2024-10-28T10:40:21.500Z

Разобравшись в основах параллелизма моделей, перейдем к не менее актуальной теме - квантизации. Квантизация моделей машинного обучения стала одним из ключевых направлений оптимизации нейронных сетей в последние годы. Этот метод позволяет значительно уменьшить вычислительные затраты и объем памяти, необходимые для работы моделей, сохраняя при этом высокую точность предсказаний.

Source: КПД, Quantization in Depth

Brief Literary Review

Вообще исследователи уже давно стараются оптимизировать модельки, уменьшать их вес и размер, при сохранении качества работы. О прунинге случайного леса я слышал еще в самом начале своей карьеры, а с появлением LLM эта тема и вовсе получила финансовый толчок от разных компаний, ведь инференсить квантизованную модельку дешевле, удобнее и быстрее, чем тратить миллионы у.е. на деплой Llama70B.

Фундаментальной работой в этой области стало исследование GPT3.int8(): 8-bit Matrix Multiplication for Transformers at Scale в которой авторы показали возможность эффективной 8-битной квантизации крупных языковых моделей. Далее ребята из Яндекса сказали, что это не предел и выкатили Extreme Compression of Large Language Models via Additive Quantization с возможностью сжатия до 2-3 битов, при чем их метод аддитивной квантизации является Парето-оптимальным с точки зрения соотношения точности и размера модели при сжатии до менее чем 3 битов на параметр. Этот же метод потом раскатили и на диффузионных моделях Accurate Compression of Text-to-Image Diffusion Models via Vector Quantization, что дало результаты лучше, чем Q-Diffusion: Quantizing Diffusion Models и Post-training Quantization on Diffusion Models. Ну и вскоре китайцы в работе VPTQ: Extreme Low-bit Vector Post-Training Quantization for Large Language Models показали, что векторная квантизация лучше всех и мы можем сжимать модели до 2 бит, при этом значительно снижать перплексию квантованных моделей и улучшать точность на задачах вопросов и ответов. А тут EDEN: Communication-Efficient and Robust Distributed Mean Estimation for Federated Learning ребята показали, что квантовать можно не только веса и активации, но и градиенты.

Из этого короткого литобзора видна довольно высокая актуальность методов квантизации, что еще раз подтолкнуло меня на разбор некоторых методов более подробно. Однако я не хочу, чтобы кто-то подумал, что я профессионал квантизации и разберу сразу все в своих постах. Если вы хотите углубиться в эту тему или развиваться в ней, то рекомендую канал КПД - там вы найдете море постов про квантизацию моделей, а сам автор является тем самым профессионалом в этой сфере.

Моей целью будет являться разбор базовых понятий, которые помогли бы в дальнейшем понимать сложные вещи, которые я буду разбирать чуть позже. Мне показалось, что сделать это будет проще всего по курсу о квантизации от HuggingFace, так что если вы его не смотрели или у вас не хватало времени на него, то у вас появилась отличная возможность прочитать summary этого курса :)

Linear Quantization

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

Параметры s (scale) и z (zero point) имеют тот же тип данных, как у original value и quantized value соответственно. Из формулы выше закономерно вытекает формула квантизация вектора.

Как видно, линейная квантизация не является каким-то rocket science, однако возникает закономерный вопрос - а как определить те самые s (scale) и z (zero point)? Для этого существуют отдельные формулы, из которых, с помощью границ области определения оригинальных значений и квантизованных, определяются s и z:

В случае, если zero point выходит из области значений q, то он приравнивается ближайшему крайнему значению q. Например, если z < q_min, тогда z = q_min.

Linear Quantization бывает двух типов: asymmetric и symmetric. Asymmetric мы разбирали только что, когда r_min не равен по модулю r_max (аналогично с q_min и q_max). В symmetric нам нет необходимости использовать zero point, ведь если q_min = q_max, то это просто середина отрезка, которая сопадает с серединой отрезка [-r_max; r_max] и равняется нулю. Из-за этого формулы немного изменятся, но суть останется той же:

Granularity

Также стоит добавить, что нет необходимости расчитывать в Linear Quantization параметры zero point и scale для всего тензора - это может негативно повлиять на точность. Лучше рассчитывать отдельные zero point и scale для каждой группы (например для каждой строки или даже n элементов тензора) - это и называется Granularity.

Очевидно, что чем меньше количество элементов в группе, тем точнее будет полученная квантизация, но тем больше памяти необходимо для хранения переменных.

Activations Linear Quantization

Помимо весов, мы можем квантовать активации. В курсе приводится следующее мнение:

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

2) Когда мы квантуем веса и активации, то часто нам необходимы типы данных в 8 битной памяти, поэтому используются такие как INT8 или INT4. Я не очень понял, что они имели в виду, когда писали, что деквантизация не поддерживается на некоторых устройствах, поэтому если вы поняли этот момент, то напишите в комментариях.

Weights Packing

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

А можем ли мы беспрепятственно квантовать до 2 бит? Кажется, мы можем столкнуться с ограничениями системы счисления, ведь довольно странно отображать числа в двух битах. Однако мы можем схитрить и таки перевести 8-битный тензор в 2-битный.

Допустим у нас будет тензор, состоящий из 4 чисел в формате UINT8:

[1, 0, 3, 2]

Их двоичная запись будет выглядеть вот так:

Разумеется, нам совсем не хочется хранить нулевые биты, которые не несут никакой информации:

Поэтому мы уберем их и соединим информативные биты в единое чило формата 8 бит:

Или же в десятичной записи:

[177]

Итак, наша 2-битная квантизация заработала. Она позволяет отображать реальный объем памяти, который занимают 2-битные квантизованные веса. Однако у этого метода есть и недостаток: если мы хотим прогнать инференс модели, нам нужно распаковать параметры, потому что pytorch не поддерживает инференс на 2-битных весах.

В коде это можно иплементировать с помощью операции OR, постепенно сдвигая биты и записывая их в заранее созданный тензор нужного размера:

for i in range(num_values):
    for j in range(num_steps):
        packed_tensor[i] |= uint8tensor[unpacked_idx] << (bits * j)
        unpacked_idx += 1
return packed_tensor

А unpacking:

mask = 2 ** bits - 1

for i in range(uint8tensor.shape[0]):
    for j in range(num_steps):
        unpacked_tensor[unpacked_idx] |= uint8tensor[i] >> (bits * j)
        unpacked_idx += 1

unpacked_tensor &= mask

Beyond Linear Quantization

LLMs становятся все больше и больше, поэтому линейная квантизация, конечно же, уже не обеспечивает должного эффекта экономии памяти и вычислений. Помимо нее существуют и другие методы, такие как LLM.int8, GPTQ, SmoothQuant, AWQ и так далее. Их статьи я буду разбирать в следующих постах.

На этом у меня все :)
Спасибо, что дочитали до конца!

Large Parallelism Post: Part V. FSDP: Fully Sharded Data Parallel

2024-08-19T13:54:47.501Z

FSDP - параллелизм с полным шардингом данных. Реализован за счет разбиения операции AllReduce на две - ReduceScatter и AllGather, а также за счет перегруппировки этих операций. Имея шард модели, данные весов собираются с других GPU за счет AllGather, далее происходит Forward pass, после чего снова собираются веса через AllGather, и только потом проиходит Backward pass. В конце градиенты обновляются с помощью ReduceScatter. Имеет 3 типа шардинга - DDP, Hybrid Sharding и Full Sharding.

Source: HuggingFace, Habr, Meta

Обратимся к разобранному ранее методу DDP - модель копировалась на каждую машину с GPU, далее проходил расчет Forward и Backward проходов, а полученные значения градиентов усреднялись с помощью AllReduce.

Разработчики задаись вопросо - почему бы не изменить пайплайн усреднения и передачи данных? Ведь это можно сделать с помощью разбиения операции AllReduce на две: ReduceScatter и AllGather.

На фазе ReduceScatter градиенты суммируются в виде одинаковых блоков по рангам на каждом GPU на основании индексов их рангов. На фазе AllGather шард-порция агрегированных градиентов, имеющаяся на каждом GPU, делается доступной всем GPU.

Далее операции ReduceScatter и AllGather перегруппировываются таким образом, чтобы каждому DDP-воркеру нужно было бы хранить лишь единственный шард параметров и состояний оптимизатора.

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

При применении FSDP на GPU имеется лишь шард модели. Сведения о весах, для выполнения прямого прохода по модели, собирают с других GPU посредством шага AllGather. Потом, до выполнения обратного прохода, сбор данных о весах выполняется снова. После выполнения обратного прохода локальные градиенты усредняются и распространяются между всеми GPU посредством шага ReduceScatter. Это позволяет каждому GPU обновить свой локальный шард весов

Более подробный пайплайн полного прохода FSDP на трех нодах [0-2] изображен на рисунке ниже.

Source: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

Разработчики также отмечают, что существуют разные стратегии шардинга параметров. Для их классификации они вводят коэффициент F - количество уровней, на которые распределяется модель. Например, если F равен 1, то вся модель будет загружена на каждую моду и мы получим классический DDP - этот случай разобран в прошлом посте. Если F задать количество GPU (обозначим это количество за W), то тогда на каждой ноде будет 1/W часть модели. Также есть гибридный шардинг - когда F принимает значения между 1 и W. Кратко отмечу параметры Full Sharding (F=W) и Hybrid Sharding (1<F<W) для случая с 16 GPU:

Full Sharding - вся модель шардится по всем GPU. На рисунке видно, что все веса разделяются между нодами, при этом они состовляют единую группу шардинга, так как являются частями одной модели. Этот способ обеспечивает наименьший объем занимаемой памяти, но требует наибольших затрат на коммуникацию параметров (в 1,5 раза больше, чем в DDP)

Source: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

Hybrid Sharding - это нечто среднее между DDP и Full Sharding. В этом случае у нас есть частичная репликация модели, а также шардинг ее параметров. Розовым выделены части шардинга модели - то есть модель копируется два раза, а веса каждой копии (репликация) распределяются между GPU (в данном случае 8). Чтобы проще это понять, представьте, что вы взяли 16 GPU и задали параметр шардинга 8 - это значит, что у вас будет 16/8 = 2 группы шардинга на 8 GPU - то есть на 16 GPU будет лежать 2 модели. А веса этих двух одинаковых моделей распределяться по группам репликации - для каждой группы по 2 GPU. Так как каждая модель разбивается на 8 частей - вот и получается 16 GPU. Такой способ используется для моделей средних размеров - они слишком малы для полного шардинга из-за низкой скорости коммуникации и слишком большие, чтобы их тренировать классическим DDP.

Source: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

Для финала приведу показатели TFLPOS на каждое GPU от размера моделей. Видно, что Full Sharding немного выигрывает в каждом варианте - однако не стоит забывать, что у него самый низкий показатель коммуникации параметров.

Source: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

На этом у меня все!
Спасибо, что дочитали до конца :)

Large Parallelism Post: Part IV. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

2024-08-12T20:03:39.166Z

В данной статье разработан метод параллельного обучения моделей с большим количеством параметров (от 100млрд до 1трл). ZeRO позволяет параллельно хранить и вычислять параметры модели, градиенты и параметры оптимизатора, сохраняя при этом низкий объем коммуникаций и высокую гранулярность вычислений.

Source: Arxive, ZeRO&DeepSpeed

Многие методы параллелизма неоптималены, когда речь идет про тяжелые модели и огромное количество данных. При этом хочется, чтобы паралеллизм сохранялся и в модели, и в данных. Ребята из Microsoft придумали Zero Redundancy Optimizer (ZeRO) - метод тренировки очень больших моделей (>100B параметров).

Для начала опишу проблему тренировки больших моделей - модель с 1 триллионом параметров и оптимизатором Adam в точности FP16 требует 16 Тб данных для хранения ее параметров, градиетов и значений оптимизатора. Такой объем данных необходимо разделять не только по памяти, но и вычислительно.

ZeRO имеет два подхода - ZeRO-DP (Data Parallel - сам механизм параллелизма) и ZeRO-R (Residual - оптимизация работы с памятью, чтобы ZeRO-DP работал корректно).

Начнем с ZeRO-DP - он имеет три вида оптимизации, которые соответствуют трем видам разделения параметров памяти:

1) Параметры модели. Это синий цвет и они имеют точность FP16

2) Градиенты точностью FP16 (оранжевый цвет), которые будут использованы для обновления весов на Backward проходе

3) Состояния оптимизатора. Выделены зеленым цветом - в него входят точные значения градиентов, дисперсия, моменты в точности FP32. Если вы используете SGD, то он не будет занимать много памяти, но вот если Adam, то памяти нужно будет тратить гораздо больше. Эти данные используются только после вычисления оранжевых градиентов.

FP16, FP32 - это точность, где 16 и 32 это сколько бит отводится под хранение. Подробнее про стандарт IEEE 754 можно посмотреть вот здесь.

fi - количество параметров модели, K - константа оптимизатора (у Adam K=12), N - на сколько GPU параллелим Source

ZeRO-R разработан для оптимизации работы с остаточной памятью во время работы ZeRO-DP. Вот что он делает:

1) Сохраняя промежуточные активации на Forward pass, чтобы использовать их на Backward pass, можно оптимизировать вычисления (Training Deep Nets with Sublinear Memory Cost), но это не работает на больших моделях. ZeRO-R разделяет и удаляет отработавшие реплики активаций.

2) Определяет соответствующий размер временного буффера для нахождения баланса памяти и вычислений (я не смог найти исходный код, поэтому не могу подробно рассказать как он это делает)

3) Предотвращает фрагментацию памяти. ZeRO-R управляет памятью, основываясь на различном времени жизни тензоров - скорее всего в зависимости от создания тензора он двигает его в ячейках памяти, чтобы не нарваться на OOM ошибку.

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

Давайте разберемся как оба подхода работают вместе.

Допустим у нас есть 4 карточки и тогда разделим входые данные на 4 части. Сама модель тоже делится на 4 части M [0-3], каждая из которых хранится на отдельном GPU (помним, что модель очень большая и хранить ее целиком на одной карте мы не можем). На каждой GPU создаем временные буфферы, в которых будем хранить промежуточные активации - они понадобятся на Backward pass. Далее с помощью broadcast распределяем параметры модели (голубые) с GPU [0] на каждое GPU [1-3] - параметров не так много и эта операция довольно дешевая.

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

Мы уже сталкивались с Activation Checkpointing в прошлой статье, однако я снова хочу остановиться на этом месте и объяснить более подробно, почему мы сохраняем лишь часть актиаций и что мы будем делать с ними потом. Я долго искал кодовую имплементацию ZeRO-R, но нашел лишь ZeRO-DP, где реализован случай полного параллелизма (ZeRO3 - случай os+g+p). Однако в документации AWS Neuron подробно и понятно описано это решение. Все дело в памяти и скорости вычислений - мы хотим сохранять расчитанные активации на Forward pass, чтобы не рассчитывать их снова на Backward pass и экономить время для вычисления градиентов. Однако при работе с большими моделями, у нас не хватит памяти сохранять сразу все активации, поэтому мы сохраняем лишь активации последнего слоя. Когда начинается Backward pass, мы досчитываем необходимые активации между сохраненными слоями, применяем их для вычисления градиентов, обновляем веса и удаляем все рассчитанные активации из памяти. Так мы одновременно экономим на вычислениях и памяти - что-то вроде trade-off между памятью и вычислениями. Хотя я наткнулся на дискуссию в GitHub, где говорят, что использовать такой метод доподсчета активаций не очень выгодно - видимо все зависит от размера модели.

После расчета части параметров модели M0, мы удаляем эти параметры из памяти GPU [1-3] (вот тут начинает работать ZeRO-R), потому что сохранив промежуточные активации, нам нет надобности хранить эту часть данных - нам придется так делать с каждой частью модели, а вся она точно не поместиться.

Так мы делаем для каждой части модели M [0-3] (то есть по очереди бродкастим параметры каждой части модели на остальные карты, считаем forward pass, сохраняем часть активаций и удаляем параметры). Когда процесс заканчивается на части M3, то на каждом GPU вычисляется значение Loss. Далее начинается Backward pass.

На каждом GPU дорасчитываются активации Forward (а часть уже сохранена) и на каждой карточке рассчитываются градиенты. Далее эти градиенты пересылаются на GPU [3], со всех остальных GPU и аккумулируются. Здесь происходит группировка градиентов - в оригинале авторы применяют на каждом процессе используют Reduce вместо AllReduce для экономии памяти.

После расчета градиентов M3 с каждой карты и сохранения их на GPU [3], на других GPU данные промежуточных активаций, градиентов и параметров модели удаляются для освобождения памяти.

Аналогично пройдемся по остальным частям модели M [0-2] (во время Backward pass будем также делать broadcast параметров моделей на каждую карту и дорасчитывать параметры активаций) и в конце получим на каждом GPU параметры градиентов. Далее параметры градиентов запускаем в оптимизатор для обновления весов. Оптимизатор определит новые параметры модели в точности FP32, которые далее переведем в точность FP16. На этом шаге цикл завершается и все прошлые шаги повторяются заново.

Добавлю, что в 2023 году вышла статья ZeRO++: Extremely Efficient Collective Communication for Giant Model Training, в которой сделано 3 улучшения - в совокупности они повышают эффективность работы алгоритма в 4 раза:

Квантование параметров с FP16 до INT8
Иерархическое разбиение, которое позволяет избавиться от повторного вычисления данных
Квантование градиентов, которое позволяет применять all-to-all обмен данными (вместо AllReduce)

Пишите в комментариях, если хотите разбор статьи :)

Results

По результатам пробегуcь как всегда быстро, потому что кто бы публиковал статью с плохими результатами?)))

Видим, что ZeRO-DP отлично превосходит 10 и 15 Pflops на моделях с большим количеством параметров, когда другие методы проседают.

Таблица с пояснением каждой конфигурации ZeRO - понадобиться для следующих графиков. С ZeRO-DP, думаю, все понятно - эти параметры нам встречались и ранее. А вот параметры ZeRO-R я прокомментирую:

CB - Constant Size Buffers - при больших моделях используется постоянный размер буфера (если модели маленькие, то буфер уменьшается)
MD - Memory Defragmentation - понятно, что фрагментация памяти так или иначе возникает при работе ZeRO, однако при работе с очень большими моделями он может выполнять дефрагментацию во время работы, предварительно выделяя смежные участки памяти для контрольных точек активации и градиентов и копируя их в предварительно выделенную память по мере их создания
Pa - Partitioned Activation Checkpointing - стандартный Activation Checkpointing, но в случае очень больших моделей и очень ограниченной памяти, данные могут быть перегружены на CPU

На рисунке 6 показаны размеры моделей при использовании различных оптимизаций ZeRO для фиксированного batch-size = 16. На рисунке 7 показано максимальное количество памяти, кэшируемой PyTorch во время каждой итерации обучения для модели с параметрами 40B и 100B. Интересно, что 100В модель при конфигурации ZeRO 5 показывает примерно ту же кэшируемость, что и 40В на конфигурации 2 и 3. На рисунке 8 показано количество операций в секунду в зависимости от конфигурации ZeRO.

На этом у меня все!

Спасибо, что дочитали до конца :)

Large Parallelism Post: Part III. Mixed Precision Training

2024-07-22T11:54:39.593Z

В статье предлагается способ снижения потребления памяти тренироки сетей в 2 раза за счет перевода величин из формата FP32 в FP16. Авторы статьи исследуют распределения полученных величин по экспоненте и вводят Loss Scaling для предовращения зануления значений активаций и увеличения точности. Последним дополнением является Arithmetic Precision - разграничение проводимых операций в FP32 и FP16 формате.

Source: Arxive

Современные системы обучения глубокому обучению используют single-precision format FP32. Авторы предлагают тренировать сети используя IEEE half-precision format FP16 - то есть тратить в 2 раза меньше битов, а значит и в 2 раза меньше памяти, сохраняя при этом точность.

Примеры числовых форматов. Source

FP32 Master Copy of Weights

Сам пайплайн mixed precision довольно прост и изображен на картинке ниже. Во время обучения веса, активации и градиенты хранятся в точности FP16, но для сохранения уровня точности создается Master Copy весов в FP32 формате, которая в последствии обновляется.

Хотя Master Copy весов FP32 далеко не всегда выгодно хранить, для этого есть причины. Первая: во многих сетях значения обновлений (веса умноженные на learning rate) могут быть меньше 2^(-24) - то есть они становятся просто нулем в оптимизаторе FP16. На гистограме ниже видно, что примерно 5% обновлений зануляются в half-precision format, что может негативно повлиять на точность модели.

Вторая: веса могут занулиться в случае большой разницы между их значениями и обновлениями. Даже если значение весов можно представить в формате FP16, но значение обновления будет довольно большим, то при сложении этих параметров результат выйдет за границы формата FP16 и станет нулем, который невозможно будет восстановить. Именно Master Copy весов помогает избежать этого.

Хоть Master Copy весов FP32 увеличивает на 50% потребление памяти по сравнению с ипользованием только FP16, влияние на общее потребление памяти все равно остается куда меньшим, чем тренировка сети только в FP32 формате. Сохранение активаций каждого слоя при Backprop является самой затратной по памяти операцией, поэтому перевод этого процесса в FP16 снижает потребление примерно в 2 раза.

Loss Scaling

Ниже представлена диаграмма доли значений актиаций от экспоненты (2 в соответствующей степени). Легко заметить, что большинство значений выходят за пределы диапазона формата FP16, а значит они зануляются.

Чтобы решить эту проблему, авторы предлагают увеличить экспоненты в 3 раза (масштабировать в 8 раз) - этого будет достаточно для соответствия точности FP32 формата.

Для эффективного сдвига значений в FP16, достаточно масштабировать Loss, вычисленный за Forward pass, перед началом вычисления Backward pass. Chain rule позволяет сохранять масштабирование, а значит не потребуется проводить дополнительные вычисления. Величина масштабирования выбирается эмпирически.

Arithmetic Precision

Для поддержания точности сетей, авторы обнаружили, что в некоторых сетях необходимо, чтобы векторное произведение накапливалось в формате FP32, а только потом переводилось в FP16 перед записью в память. Скорее всего это вызвано ограничениями точности арифместических операций в FP16. Тензорные ядра в архитектуре GPU Nvidia позволяют накапливать dot-product либо в FP16, либо в FP32 формате - Nvidia tesla v100 gpu architecture.

Большие суммы по элементам вектора (такие как в batch-normalization или softmax) также должны выполняться в FP32 формате. То есть чтение и запись в память производится в FP16, а арифметика в FP32 для сохранения точности. Однако, это не замедляет скорость вычислений, поскольку такие операции ограничены пропускной способностью памяти.

Results

Эксперименты проводились для двух подходов:

Baseline (FP32) - активации, веса и градиенты хранятся в single-precision формате. Все вычисления также проводятся в FP32.
Mixed Precision (MP) - хранение в памяти и некоторые вычисления осущетсвляются в FP16. Веса, активации и градиенты хранятся в FP16, используется Master Copy весов в FP32. Для некоторых сетей используется Loss Scaling. Используются операции Tensor Core с накоплением в FP32 для сверток, полносвязных слоев и матричных умножений в рекуррентных слоях.

По таблице сравнения точности видно, что MP нередко немного превосходит Baseline вычисления.

Обратите внимание, как выбор величины Loss Scaling влияет на тренировку модели.

Также эксперименты проводились с моделями CNNs Detection, Speech Recognition, Machine Translation, Language Modeling, GANs - все они показали результаты на уровне Baseline.

Спасибо, что дочитали до конца :)

Large Parallelism Post: Part II. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

2024-07-01T21:28:29.718Z

Во второй части мы углубимся в Tensor Parallelism на основе статьи Megatron-LM. В ней представлен способ параллельных вычислений внутри блоков MLP и Attention. Благодаря разделению весовых матриц по столбцам и строкам, становится возможным распараллелить блоки MLP и Attention между GPU с минимальными коммуникациями между нодами. Также разберем пайплайн TP+DP.

Source: Arxive, Tensor Parallelism

В прошлом посте Large Parallelism Post: Part I мы разбирали Model Parallelism, который включает в себя два типа параллелизма:

Вертикальный - несколько слоев модели на каждом GPU
Горизонтальный - мы размещаем часть всей модели на каждом GPU

Горизонтальный параллелизм модели и называется Tensor Parallelism (TP) и именно с ним мы сегодня будем разбираться.

Главное отличие горизонтального параллелизма от вертикального в том, что он исключает простаивание GPU благодаря разделению всех слоев модели между нодами. На схеме ниже показаны их основные различия: в вертикальном мы на каждом GPU храним отдельный кусок слоя; в горизонтальном мы делим каждый слой модели между всеми GPU, после снова собитраем в один слой и снова делим между нодами

Source

Для начала обозначим, что именно мы будем параллелить - это классический слой трансформера, который изображен на схеме ниже. Глобально, он состоит из двух блоков - MLP и Attention. Внутри каждого блока есть слой Dropout, который в совокупности также хотелось бы параллелить.

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

В случае параллелизма по столбцам, мы распределяем матрицу весов на два столбца, умножаем на входные данные (важно, что входная матрица X одинакова для A1 и A2 в этом случае), получаем два выходных столбца Y1 и Y2 и конкатенируем их в выходную матрицу Y.

Для параллелизма по строкам, мы уже способны разделить входную матрицу X, но разделяем ее по столбцам, а вот матрицу весов A делим по строкам, производим умножение соответствующих частей, получаем выход Y1 и Y2, складываем их и получаем финальную матрицу Y.

Обратимся теперь к статье - как мы обсуждали ранее, в ней представлен параллелизм двух блоков: Multi-Layer Perceptron (MLP) и Attention.

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

Первые мысли, которые могут возникнуть, когда мы хотим что-то параллелить - давайте поделим матрицы входа и весов (row parallelism):

Однако, это не сработает, потому что GeLU является нелинейной функцией и у нас не получится получить финальную матрицу с помощью сложения:

Однако, мы можем решить эту проблему, если сначала распараллелим в этом случае только матрицу весов (column parallelism). Тогда выходом из такой операции будет два столбца Y, которые мы сконкатенируем:

Обращу внимание, что матрица входных данных остается одинаковой для каждой матрицы A1, A2 и не параллелится.

После получения выходных матриц Y1, Y2, они подаются в следующий слой, где матрицы весов B1, B2 разделены по строкам (row parallelism), после чего формируются выходные матрицы Z1, Z2 и они подаются в оператор g, а дальше формируется финальная матрица Z, после прохода Dropout слоя.

Подробнее расскажу, что за операторы f и g на схеме выше - это AllReduce оператор на Backward и Forward pass соответственно. g работает, когда после Forward pass необходимо собрать все данные для формирования выходной матрицы, а f после Backward pass формирует матрицу градиентов.

Теперь разберем параллелизм Attention слоя. На самом деле, там все работает аналогично MLP, только матриц чуть больше. Сначала применяется column parallelism для матриц Query, Key, Value, далее их результат проходит через Sofmax и Dropout, после чего матрицы Y1, Y2 умножаются на веса B1, B2 в слое Dropout, который работает с помощью row parallelism, а потом все подается в оператор g и формируется финальная матрица Z.

Вот общий пайплайн tensor parallelism - по сути в каждом разобранном блоке у нас есть верхняя и нижняя часть, которые мы можем отправить на отдельные GPU. Тогда у нас остается только 4 операции AllReduce, которые мы будем делать для сбора финальных матриц оператором g и сбора матриц градиентов оператором f.

Source

Еще одна схема пайплайна tensor parallelism трансформера. Я ее прикрепил для понимания, почему всего будет 4 операции AllReduce, а также для комментария, что в пайплайне присутствуют Skip Connections.

Матрицу входных эмбеддингов авторы параллелят по столбцам E = [E1, E2]. Для выходных эмбеддингов Y1, Y2 авторы для них высчитывают Cross-Entropy Loss и после, с помощью AllGather, получают финальные вероятности. Это сделано для сокращения времени коммуникации между GPU, поскольку сводит размер финального пространства к b x s (b - batch-size, s - sequence length).

Также добавлю, что авторы используют Activation Checkpointing. Он заключается в сохранении только конечных активаций каждого слоя на Forward pass, а на Backward pass все активации между сохраненными слоями будут рассчитываться повторно и использоваться для получения градиентов. Это экономит затраты памяти и вычислительные ресурсы.

До этого мы обсуждали, что делим трансформер между двумя карточками, но, конечно же, его можно делить между многими GPU. Я думаю этот факт более чем очевиден, поэтому перейдем к части совмещения TP и DP. Да, так можно делать и вот схема пайплайна:

Если кратко, то вы можете создать ноду с 8 GPU (то есть разделить модель на 8 частей), потом перекопировать эту ноду 64 раза, а в каждой ноде подавать на определенное GPU определенную часть данных (да да, это тот самый самый простой DP). То есть, например на GPU [1, 9, 17 ... 505] в каждой ноде вы подаете данные Group1, на GPU [2, 10, 18 ... 506] данные Group2 и так далее. Градиенты будут аккамулироваться как в DP, а каждая нода теперь способна вместить в себя большую модель (хватило бы только GPU).

Results

Судя по графикам, пайплайн TP+DP довольно неплохо справляется с большим количеством GPU - есть небольшая просадка на несколько процентов по сравнению только TP, но это скорее из-за малой задержки NCCL операций между карточками.

Добавлю еще вот такую табличку, которая показывает сколько дней тратится на одну эпоху, обучая GPT-2 в разных конфигурациях на 512 GPUs (одна эпоха это 68507 итераций).

На этом у меня все!

Спасибо, что дочитали до конца)

Large Parallelism Post: Part I

2024-06-29T15:38:13.812Z

В этой части разобраны самые основные методы параллельного обучения сетей - Data Parallel, Distributed Data Parallel, Model Parallelism и Pipeline Parallelism.

Source: Model Parallelism HF, GPipe

Data Parallel (DP)

Классический параллелизм - параллелизм данных, реализованный в PyTorch и применяющийся одной строчкой:

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

В этом случае модель сначала копируется с "0" GPU на остальные, а данные делятся между карточками. Forward и Backward проходят на каждом GPU отдельно, а результаты полученных градиентов после каждого прохода суммируются в модель на "0" GPU. С нового цикла Forward и Backward все начинается заново.

Source

Важные уточнения:

При каждом Forward проходе модель копируется с "0" GPU на остальные - то есть если существуют процесс модели на "1" карте, то он будет потерян.
Тензоры распределяются между GPU, но типы tuple, list и dict будут скопированы (именно скопированы, без deepcopy). Остальные типы данных будут использоваться сразу всеми карточкмаи, поэтому будут повреждены. Так что лучше все сразу переводить в тензора.
Очень важно подчеркнуть, что DataParallel использует многопоточность GPU, что снижает их эффективность.
DataParallel может параллелить GPU только на одной машине, что ограничивает размер модели и объем данных.

Прокомментирую пункт 3 подробнее. Многопоточность предполагает общую память между потоками GPU, а значит существует риск непредвиденного обновления ячеек памяти. Это может возникнуть из-за неправильного освобождения и ее выделения - такое встречалось в работе с NCCL (вот ссылка на тред).

Также необходимо добавить, что потоки Python зависят от GIL, который может блокироваться C bindings (это структурное связывание двух языков программирования C и Python - подробнее тут). Вычисления на GPU обычно используют C bindings, а значит могут иногда блокировать GIL, что остановит работу основного кода.

Distributed Data Parallel (DDP)

Для решения описанных проблем, сам PyTorch настоятельно рекомендует использовать метод DistributedDataParallel (DDP). Он вызывается также легко, как и DataParallel:

from torch.nn.parallel import DistributedDataParallel as DDP

ddp_model = DDP(model, device_ids=[rank])

Подход остался тем же самым - одна модель с "0" GPU копируется на остальные, а данные распределяются между процессами. Кстати делается это с помощью ссылки на состояние модели (state_dict), которое с помощью broadcast распространяется на все GPU. Это нужно для того, чтобы все копии модели запускались из того же самого состояния, что и оригинальная модель. При этом важно уточнить, что теперь вы не ограничены одной GPU для модели - вы работаете с машинами, где может находится сколько угодно GPU. Это означает, что для запуска больших моделей вы ограничены только бюджетом кластера.

Как я и сказал - подход остается действительно таким же, но все дело в мелочах. Например, DDP работает на мультипроцессинге, что и позволяет запускать его на нескольких машинах (и неважно сколько GPU на каждой машине). К тому же, так как каждый GPU имеет свой выделенный процесс, то это позволяет избежать перерасхода производительности, вызванного блокировками GIL. Пример работы DDP показан на рисунке ниже.

Source

Продолжая разговор о мелочах, необходимо добавить, что каждый процесс имеет свой метод reduce, который нужен для запуска AllReduce усреднения градиентов при Forward и Backward проходах. Для повышения эффективности работы с градиентами, Reducer разбивает градиенты по buckets (размер можно настроить с помощью bucket_cap_mb), чтобы применять операции не к каждому градиенту, а сразу к их группе. Важно отметить, что параметры модели распределяются по бакетам примерно в порядке, обратном порядку Model.parameters() данной модели - это логично, ведь получение градиентов начинается с последних слоев к первым.

Source

Optimizer Step происходит для модели на "0" процессе, а затем ее состояние снова рассылается по всем процессам и цикл начинается заново.

Model Parallelism (MP)

Очень хорошо, когда ваша модель помещается на одном GPU или даже на ноде кластера, но при работе с большими моделями такого ожидать не приходится - что делать если мы не обладаем достаточными мощностями для работы с большими моделями? Ответ очевиден - давайте параллелить модель.

Существует два варианта "резки" слоев модели - вертикальный и горизонтальный. Как обычно бывает в научной среде, в терминологиях путаются и иногда горизонтальный параллелизм называют вертикальным и наоборот (вот в этом видео Aleksa именно так и делает). Я буду придерживаться терминологии HuggingFace. Для начала расскажу про вертикальный параллелизм модели, а к горизонтальному мы вернемся в статье про Tensor Parallelism.

Представьте, что вы можете нарезать слои модели и каждую такую часть отправить на отдельный GPU.

Source

Тогда вы будете пропускать данные сначала через первые слои модели на GPU[0], выход отправите на GPU[1], а после на GPU[2], где рассчитаете Loss, и начнете передавать градиенты с последних слоев к началу GPU[2->0].

В статье GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, к которой мы обратимся чуть позднее, изображена схема такого пайплайна, на которой сразу видны очевидные минусы данного подхода.

Пока проходит Forward или Backward pass на каждом GPU, остальные простаивают. Более того, общие данные придется копировать между каждым GPU, что займет время и ресурсы.

Чтобы решить данную проблему, разработчики Google решили придумать Pipeline Parallelism.

Pipeline Parallelism (PP)

Снова обратимся к работе GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. Pipeline Parallelism очень похож на MP - в нем также реализовано разделение модели на слои, которые хранятся на каждом GPU. Но в нем есть небольшое отличие в работе с входящими данными - каждый mini-batch разбивается на несколько micro-batches, то есть на еще более мелкие пакеты. Это позволяет обрабатывать каждый micro-batch параллельно и зон, где GPU простаивает, становится гораздо меньше. Оставшаяся зона простойки GPU называется Bubble.

Условно, вместо обработки одного большого mini-batch F0 из прошлой схемы, GPU нужно обработать четыре последовательные части F0. После расчета каждой части, результаты можно передавать на следующий GPU. Аналогичная ситуация происходит на Backward pass. Конечно, некоторый простой GPU сохраняется и в этом варианте (Bubble на рисунке), но он значительно меньше, чем при MP, что ускоряет обучение.

Спасибо, что дочитали до конца!
В следующей части углубимся в Tensor Parallelism на основе статьи Megatron-LM.

xLSTM: Extended Long Short-Term Memory

2024-06-11T17:35:16.977Z

В статье разработано улучшение сети LSTM с помощью двух архитектур - sLSTM и mLSTM. Каждая архитектура, направлена на решение проблем оригинальной сети. Добавлены параллельные вычисления, способность корректировать запоминание информации, а также матричное представление данных внутри сети. Эксперименты доказывают, что xLSTM сравнима с GPT моделями.

Source: Arxive

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

Source

Вначале получаем взвешенную сумму входного вектора и вектора скрытого состояния (коэффициенты в этой сумме и есть весовые матрицы).
Forget Gate - результат применения сигмоиды к скрытому и входному векторам. Результат умножается на вектор контекста, решая какую информацию нужно забыть с учетом полученных состояний.
New Cell Content - получается с помощью гиперболического тангенса. Расчитывается новый контекст и одновременно решается какая информация в новом векторе релевантна с помощью умножения на сигмоиду результата взвешенной суммы входного и скртого состояния (Input Gate). Новый контекст прибавляется к прошлому - именно эта операция и отвечает за способность LSTM запоминать длинный контекст.
Вычисляется новое скрытое состояние с помощью гиперболиечского тангенса финального вектора контекста и очередным умножением на сигмоиду суммы входного и скрытого векторов (Output Gate).

У такой архитектуры есть пара проблем:

Невозможно производить вычисления параллельно
Невозможно корректировать решения сети запоминать/забывать информацию
Скалярность памяти вносит некоторые ограничения на ее эффективность

В новой статье ученые решили эти недостатки.

Во-первых уточню, что xLSTM состоит из двух архитектур - sLSTM и mLSTM. Начнем с sLSTM.

sLSTM

Картинку позаимствовал с разбора Data Secrets

Давайте внимательно посмотрим что изменилось в схеме sLSTM.

Сразу замечаем новую красную ячейку памяти n (normalization) над ячейкой контекста.
Первые две сигмоиды заменяют экспоненциальные функции.
Деление вместо гиперболического тангенса при расчете нового скрытого состояния.

Теперь перейдем к формулам.

Вот что пишут сами авторы о нововедениях:

Чтобы наделить LSTM способностью пересматривать решения о хранении, мы вводим экспоненциальные gates (красные) вместе с нормализацией и стабилизацией. В частности, input gates и forget gates могут иметь экспоненциальные функции активации.

Что произошло? Раньше мы не могли получать большие значения из-за ограниченности сигмоиды. Теперь, расчитывая экспоненту, у LSTM есть возможность регулировать релевантность информации в input gate и forget gate.

Например, если входной вектор является очень важным для сети, значение input gate будет высоким, а значит и умножение на вектор контекста даст большой результат. Одновременно с этим сеть понимает, что прошлые векторы были неважными, а значит мы получим малые значения в forget gate (кстати именно поэтому авторы на картинке указывают, что для forget gate можно применять как сигмоиду, так и экспоненту - неважно как сеть занулит прошлые значения).

Теперь перейдем к новой ячейке нормализации. Расчет данного значения можно увидеть в формуле 9, но зачем он нужен? По формуле мы видим, что состояние нормализации постоянно увеличивается за счет накопления значений input gate. Да, есть возможность забывать прошлую информацию с помощью forget gate, однако в этой ячейке все равно аккумулирована вся релевантная информация текста. А теперь обращаемся к формуле 10 - частному вектора контекста и вектора нормализации. По сути здесь сеть вычисляет насколько релевантная информация содержится в векторе контекста по отношению ко всему тексту документа. А также это позволяет архитектуре передать значение этой релевантности в следующий слой.

Перейдем к последней части sLSTM - стабилизации. Мы понимаем, что экспонента может выдавать очень большие значения, которые приводят к переполнению памяти. Поэтому авторы вводят алгоритм стабилизации этих значений - состояние m (которое берут из статьи Online normalizer calculation for softmax). Оно выбирает максимум из двух значений: суммы прошлого состояния стабилизации с логарифмом forget gate и логарифма input gate (формула 15). Далее считаются новые значения input gate и forget gate по формулам 16 и 17.

На инференсе, конечно же, это является трюком для предотвращения взрыва градиентов и переполнения памяти. Однако на этапе тренировки данная часть вычислений очень важна и вот почему.

Когда левая часть максимума с логарифмом forget gate и прошлым стабилизатором превосходит логарифм от input gate, то финальное значение forget gate обнуляется - то есть при высоком значении текущего forget gate и прошлой памяти, сеть умножает на ноль вектор контекста и нормализации.
Напротив, когда логарифм input gate превосходит сумму логарифма forget gate и прошлого стабилизатора, то обнуляется финальное значение input gate, а значит зануляется текущий контекст и информация не добавляется в вектор нормализации.

Мы учим сеть находить баланс между добавлением и забыванием информации. Если инфомация важна, то нельзя сразу присвоить ей большое значение, иначе оно заглушит прошлые знания. Если мы хотим занулить текущую информацию, то, конечно, можно ей выдать высокие значения forget gate, однако на следующем шаге это действие может перевесить нечто важное в input gate, даже если forget gate будет мал, и тогда снова произойдет забывание.

Последнее, что сделали авторы - вместо работы с одной цепочкой блоков sLSTM авторы делают несколько голов с помощью матриц, подражая multi-head attention (они назвали это New Memory Mixing, хотя вообще то об этом известно уже очень давно. Ниже я приведу код для более подробного объяснения). Матрицы Wz, Wi, Wf, Wo, Rz, Ri, Rf, Ro являются теперь блочно-диагональными, где каждый диагональный блок задает отдельную голову. В этом случае, скаляры становятся, очевидно, векторами. Сам Memory Mixing может происходить только внутри каждой головы, а не между голов.

mLSTM

В сети mLSTM авторы увеличивают объем памяти с помощью агрейда скаляра с до матрицы C. Они используют терминологию трансформеров и вводят вектора q, k и v для хранения и извелечения памяти. Извелечение необходимой информации из памяти основано на правиле обновления ковариации, которое позволяет сохранять пары векторов (v, k):

Разберем на простом примере как это работает. Допустим нам нужно сделать 2 итерации по сохранению векторов (v, k) и мы получим :

А теперь нам нужно достать из сохраненной памяти вектор v0:

Потому что

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

Это правило встречалось уже ранее в статье Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention, разбор которой я делал в канале.

Итак, разберемся подробнее. Как я писал ранее - авторы используют вектора q, k и v наподобие трансформеров, вычисление которых вводится в формулах 22, 23 и 24. Здесь нет ничего нового - входной вектор умножается на матрицу весов и к результату добавляется bias. В случае k видим классическое сохранение размерности путем деления на корень из d - это реализовано в ванильном self-attention.

Важно отметить, что в input, forget, output gate теперь не используется скрытое состояние (формулы 25, 26, 27) - это очень важное обновление, поскольку теперь появляется возможность распараллелить процесс обучения.

Добавлю, что для input и forget gate используется тот же процесс стабилизации, что и в sLSTM.

Теперь снова вернемся к правилу ковариации и к формулам 19, 20 и 21.

В 19 формуле мы видим формирование матрицы C путем накопления пары векторов v и k. Да, мы снова прибегаем к forget, input gate для контроля забывания и релевантности накапливаемой информации.

Аналогично мы поступаем в формуле 20 с одним отличием - там происходит накопление только векторов k.

Наконец, в 21 авторы приводят расчет скрытого состояния архитектуры mLSTM - извлекают и нормализуют релевантную информацию из памяти. В правой части в числителе происходит извлечение необходимой информации из матрицы C с помощью вектора q. Да, в разобраном выше примере ковариации я использовал вектор k, однако в таком случае мы извлекали только один конкретный вектор v. В реальности нам необходимо регулировать извлечение релевантных векторов v, в зависимости от текущего контекста, поэтому извлечение происходит с помощью вектора q. Похожее извлечение мы видим и в знаменателе, только там мы работаем только с накопленными векторами k, и берем максимум между получившимся значением модуля произведения векторов и единицы - это необходимо чтобы избежать деления на малые значения, близкие к нулю. В конце мы умножаем результат на output gate, как это делали в sLSTM.

Выходом mLSTM является скрытое состояние h (это будет хорошо видно по коду далее).

xLSTM

Поздравляю, что вы дочитали до этого места, потому что представляю, как сложно за один раз уложить в голове все вышенаписанное (мне потребовалось 2,5 дня полной работы со статьей).

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

Еще раз спасибо DataSecrets за их разбор и комментарии к изображениям!

Для sLSTM все начинается с LayerNorm (LN), через который проходит входной вектор. Далее опционально применяется 1D свертка с окном 4 и нелинейная функция Swish перед подачей данных в input и forget gate. Потом для input, forget, z и output gate данные подаются через блочно-диагональные линейные слои с четырьмя диагональными блоками (или головами).

Этот момент мне кажется довольно непростым для понимания, поэтому я решил порыться в коде имплементации. Если захотите подробно разобраться в этом шаге, то вам нужен класс LinearHeadwiseExpand, но если описывать кратко, то в этом классе входные данные проецируются в более высокую размерность, разделяясь на несколько независимых линейных преобразований (тех самых голов), преобразуя входной тензор x в форму (..., num_heads, in_features // num_heads)

self.weight = nn.Parameter(
            torch.empty(num_heads, out_features_per_head, in_features // num_heads),
            requires_grad=config.trainable_weight,
        )

x = x.view(*shape[:-1], self.config.num_heads, -1)
x = torch.einsum("...hd,hod->...ho", x, self.weight)
x = x.reshape(*shape[:-1], -1)

Каждая голова имеет свой собственный набор весов (по сути это и есть то самое New Memory Mixing). Их результаты объединяются в один выходной тензор.

Далее происходит сама работа sLSTM, где sLSTMCell_vanilla возвращает ячейки памяти:

return torch.stack((ynew, cnew, nnew, mnew), dim=0),
       torch.stack((igate, fgate, zraw, ogate), dim=0)

Которые в дальнейшем (метод forward в классе sLSTMLayer) проходят через Dropout, поступают в GroupNorm и передаются сначала в up-projection для увеличения размерности (снова происходит параллельное разделение с использованием функции GeLU), а после в down-projection для возвращения данных к первоначальному размеру. Здесь происходит что то вроде отсева качественных данных, если можно это назвать так грубо. Не забудем про skip-connections, которые добавляются к результату, чтобы побороть затухание градиента.

Теперь обратимся к сети с mLSTM. При разборе, я опирался на код ее слоя mLSTMLayer. Ее pipeline выглядит примерно также, только входные данные сначала проходят через LayerNorm и up-projection, одновременно разделяясь на 2 потока - один проходит через mLSTM, другой через активацию Swish (или SiLU). Данные, поданные в mLSTM, снова разделяются (авторы очень любят паралеллить, как вы заметили) и для векторов q и k они предварительно проходят через слой 1D свертки с окном 4 аналогично sLSTM (от сюда же данные добавляются через LearnableSkip в пост обработку - этот LS является обучаемым). Точно также с помощью LinearHeadwiseExpand данные переводятся в блочно-диагональный вид (Block Size = 4) для каждого вектора и подаются в mLSTMCell, который возвращает h, c, n и m значения (m - это переменная из стабилизации):

h, (c_state_new, n_state_new, m_state_new)

Вот так это выглядит в коде:

h_tilde_state, mlstm_state = self.mlstm_cell.step(q=q, k=k, v=v, mlstm_state=mlstm_state)

Я не очень понял множество линий передачи матриц q, k и v, которые авторы нарисовали в блоке mLSTM на картинке выше. В коде не нашел ничего подобного - все сводится к скрытому состоянию h, к которому добавляется значение из LearnableSkip. Результат умножается на output gate от результата skip-connection через Swish (или SiLU) функцию, проходит через понижение размерности down-projection, слой dropout и возвращается вместе с вычисленными состояниями mlstm_state, conv_state:

h_tilde_state_skip = h_tilde_state + (self.learnable_skip * x_mlstm_conv_act)

# output / z branch
h_state = h_tilde_state_skip * self.ogate_act_fn(z)

# down-projection
y = self.dropout(self.proj_down(h_state))
return y, {"mlstm_state": mlstm_state, "conv_state": conv_state}

Нужно отметить, что данные действительно проходят через GroupNorm, только реализован он в mLSTMCell в виде:

h_state_norm = self.outnorm(h_state)  # (B, NH, S, DH)

Как же работает xLSTM? Все просто - эти слои соединяются друг с другом, формируя единую сеть с названием xLSTM. Рекомендую посмотреть пример в ноутбуке. Количество блоков той или иной архитектуры регулируется пропорцией, то есть в xLSTM[7:1] будет 7 блоков mLSTM и 1 блок sLSTM (или 42 mLSTM и 6 sLSTM).

Результаты

Их много, поэтому как обычно приведу несколько, которые меня привлекли больше всего (также не хочу увеличивать размер статьи). Если хотите узнать больше, то рекомендую прочитать их тут. Канал gonzo уважаю)

Итак, при сравнении предсказания следующей лексемы при обучении на 15B из SlimPajama, xLSTM показывает лучшие результаты (правда Llama бралась не 70B, а 1.3B). Аналогичная ситуация при трейне на 300B.

Экстраполяция последовательностей в языковом моделировании. Это сравнение больших моделей xLSTM, RWKV-4, Llama и Mamba размером 1,3B при предсказании следующей лексемы на валидационном наборе SlimPajama после обучения на 300B лексем от туда же. Модели обучались на длине контекста 2048, а затем тестировались на длинах контекста до 16384. Слева: оценка сложности лексем при различных длинах контекста. В отличие от других методов, модели xLSTM остаются на низком уровне сложности для более длинных контекстов. Справа: Качество предсказания при экстраполяции на большие размеры контекста в терминах валидационной perplexity (PPL). xLSTM дает лучшие значения PPL.

Итог

В целом видно, что RNN все также сильны и да, к ним в последнее время проявляется интерес. Сложно сказать из-за чего это происходит, но я полагаю, что люди пока не изобрели что-то лучше трансформеров и SSM, поэтому обращаются к прошлому и улучшают его.

Я не думаю, что сейчас все бигтехи массово перейдут на эту архитектуру. И никто не будет переобучать gpt-4 и gpt-4o на xLSTM (хотя xLSTM уже приспособили для задач CV - вот статья Vision-LSTM: xLSTM as Generic Vision Backbone).

Мое мнение - крутой апгрейд RNN, который, вероятно, локально применят в NLP/CV отделах RnD, в стартапах и в науке. Если он хорошо себя зарекомендует, есть вероятность, что увидим развитие этой технологии, а также новые решения в проде :)

На этом у меня все!

Спасибо, что дочитали до конца! Я знаю, что это было трудно, но вы справились)

RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval

2024-05-31T22:49:32.677Z

В статье разработана новая архитектура RAG на основе итерационного процесса кластеризации фрагментов текстов БД и их суммаризации. Архитектура позволяет сети отвечать на тематические запросы, требующие суммаризованного контекста всего документа (пример - краткий пересказ книги).

Source: Arxive

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

Если вкратце перефразировать проблему - чем больше фрагментов текста вы включаете в запрос, тем меньше вам нужен RAG. Ведь по сути вы можете вместе с запросом отправлять в LLM целые книги (контекстные окна сегодняшних моделей позволяют это делать). Например top-K фрагментов текста не смогут дать полного ответа на запрос "Через что прошли дети капитана Гранта, чтобы найти своего отца?".

Также они не смогут полно ответить на сравнительные вопросы типа "Как?": "Как ты думаешь, идти мне на Data Fest или идти на свидание?" из-за большого количества релевантных фрагментов текста. В целом для любой системы RAG найти несколько конкретных частей текста в большом документе является сложной задачей.

Что же делает RAPTOR?

Вместо разделения документов на маленькие фрагменты и сохранения их в векторную БД для последующего извлечения, RAPTOR сначала их кластеризует, а после суммаризует каждый кластер с помощью LLM. Он повторяет этот процесс итерационно, пока не остается один, финальный фрагмент текста, в котором содержится вся информация документа.

Каждая нода содержит следующую информацию:

Индекс ноды
Индексы дочерних нод
Summary текст дочерних нод
Эмбеддинг summary

Именно такая архитектура "от общего к частному" отлично работает на любых запросах - краткое содержание книги, сравнение двух книг или даже конкретные факты из обоих документов с сравнением. Все это извлекается в общих чертах с готовой суммаризированной информацией, а если необходимы факты, то можно опуститься на слой ниже и извлечь более детальное summary.

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

Алгоритм кластеризации основан на Gaussian Mixture Models (GMMs) - такой подход обеспечивает необходимую гибкость и вероятностную структуру. GMM - это вероятностная модель, которая предполагает, что все точки данных генерируются из смеси конечного числа гауссовских распределений с неизвестными параметрами.

В качестве финального распределения вероятности (к каким кластерам принадлежит вектор текста) авторы используют взвешенные Гауссовы распределения.

Первая формула - условная вероятность принадлежности вектора x (эмбеддинга текста некоторой размерности) к некоторому кластеру k. Остальные параметры - параметры гауссого распределения N. Финальное распределение - взвешенная сумма всех распределений. Параметры π определяют принадлежность каждого распределения к кластеру.

Однако высокая размерность эмбеддингов представляет собой проблему для традиционных GMM, поскольку метрики расстояний могут плохо себя вести при измерении сходства в высокоразмерных пространствах. Поэтому авторы используют Uniform Manifold Approximation and Projection (UMAP) для уменьшения размерности. Варьируя параметр количества соседей в кластерах, они понижают размерность эмбеддингов в 2 этапа:

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

Такой подход позволяет захватить как общие темы, так и конкретные детали.

Оптимальное количество кластеров определяется с помощью Bayesian Information Criterion (BIC) - критерий, который сильно штрафует модель за увеличение количества параметров и вознаградает за уменьшение.

Здесь N - количество фрагментов текста, k - параметры модели а L - максимизированное значение функции правдоподобия модели (likelihood function). Благодаря этому критерию, у нас не возникнет ситуации равенства количества фрагментов текста и кластеров.

Для суммаризации текста авторы использовали gpt-3.5-turbo, однако около 4% summary содержали незначительные галлюцинации. Они не распространялись
на родительские узлы и не оказали заметного влияния на решение задач, связанных с ответами на вопросы.

RAPTOR предоставляет 2 способа запроса - traversal и collapsed tree.

Traversal tree - сначала выбирается top-k наиболее релевантных корневых узлов на основе их косинусного сходства с запросом. На следующем уровне рассматриваются дочерние узлы этих выбранных узлов, и из этого пула снова выбираются top-k узлов на основе их косинусного сходства с вектором запроса. Этот процесс повторяется до тех пор, пока мы не достигнем узлов листа.

Collapsed tree - более простой способ поиска релевантной информации за счет одновременного рассмотрения всех узлов дерева. Вместо того чтобы переходить от слоя к слою, этот метод сглаживает многослойное дерево в один слой, выводя все узлы на один уровень для сравнения.

Результаты

RAPTOR превосходит базовые показатели каждого из соответствующих методов поиска (SBERT, BM25, DPR) на наборе данных NarrativeQA, используя в качестве языковой модели UnifiedQA-3B.

Также он превосходит поисковики на датасетах QuALITY и QASPER.

Сравнение оценок F-1 на наборе данных QASPER с использованием трех различных языковых моделей (GPT-3, GPT-4, UnifiedQA 3B) и различных методов поиска. «Title + Abstract» отражает производительность, когда для контекста используются только название и аннотация статей.

Спасибо, что дочитали до конца!

Пишите свои мысли в комментариях)

Advanced RAG Pipelines

2024-05-31T12:20:02.413Z

Начиная писать этот материал я решил не вдаваться в подробности базовой архитектуры Retrieval-Augmented Generation (RAG), поскольку про нее и так много известно (но на крайний случай оставляю ссылки на краткое объяснение и полный гайд). Также хочу поделиться классным репозиторием от LangChain - в нескольких ноутбуках from scratch реализован RAG в нескольких вариантах для разных БД. Шпаргалка по работе RAG именно от туда:

Source

Base RAG

Все знают, что LLM помогают в любой работе - писать код, статьи, проводить собесы и так далее. Однако у них есть несколько проблем - устаревшая информация и галюцинации. Например, вводя запрос "Какой сейчас курс доллара?", LLM, очевидно, не может воспользоваться той информацией, которую получала при обучении, потому что курс все время меняется. Спрашивая ее о конкретных фактах "Какое влияние оказали первые реформы Столыпина на экономику Российской Империи?", LLM может сгенерировать неточную или неверную информацию.

Retrieval-Augmented Generation помогает решить эти проблемы - он ищет по запросам пользователя релевантные документы в базе данных (БД) и прикрепляет их к промпту в LLM. Сами документы разбиты в БД на небольшие фрагменты текста (chunks), которые перекрывают друг друга - это необходимо для подачи в сеть последовательности фрагментов и сохранения контекстной взаимосвязи между ними.

Source

Фрагменты текста хранятся в векторной БД (например QDrant, Chroma, FAISS), а их релевантность с запросом пользователя можно оценивать с помощью CosSim, Cross-Encoder, Bi-Encoder. Так LLM получает точные данные из БД вместе с запросом, что и помогает ей давать релевантные ответы.

Source

Теперь давайте сфокусируемся на более сложных pipeline-ах RAG и попробуем их детально разобрать.

Document Hierarchies

Начнем с иерархии документов. Думаю по названию уже понятен основной смысл - мы можем создать что-то вроде оглавления к RAG чтобы структурировать и сегментировать наши данные. Закономерный вопрос - почему векторная БД не справиться с этой задачей и необходимо придумывать новую архитектуру хранения данных?

Попробуем ответить вопросом на вопросом - а что если в нашей системе миллиарды документов (например - бухгалтерская база экономики страны)?

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

2) Более того, с увеличением объема данных количество шума при каждом извлечении увеличивается, а значит система гораздо чаще находит неверные данные, которые просто оказались ближе друг к другу.

Иерархическая структура документов пытается решить эти проблемы - она извлекает и сегментирует семантические представления текста. С помощью многоуровневого распределения система итеративно уточняет наше пространство поиска, чтобы использовать только тот набор данных, который, как мы знаем, имеет семантически релевантное содержание для нашего первоначального запроса. В этом случае документ (родительская нода) выступает суммарным эмбеддингом дочерних фрагментов текста.

Пример трехуровневой системы:

Кластера или группы документов
Отдельные документы по соответствующим темам
Отдельные части релевантных документов

Source

Еще одну реализацию иерархической структуры документов можно найти у LlamaIndex - HierarchicalNodeParser. В этом варианте иерархия строится на убывании размера чанков текста - от большего к меньшему.

Например:

Размер чанка 2048
Размер чанка 512
Размер чанка 128

Knowledge Graphs

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

Узлы таких графов представляют отдельные сущности, такие как люди, места, объекты или концепции. А ребра представляют взаимоотношения между этими узлами, указывая на то, как они связаны друг с другом.

Source

Но давайте разберемся как это работает.

На самом деле LLM можно использовать еще на базовом уровне, генерируя запрос на языке запросов к графовой БД (например Cypher, GraphQL или Gremlin). Примерный промпт такой задачи Text2Cypher может выглядеть вот так:

You are a NebulaGraph Cypher expert. Based on the provided graph Schema
and the question, please write the query statement.
The schema is as follows:
---
{schema}
---
The question is:
---
{question}
---
Now, write down the query statement:

LLM отлично справляется с выделением ключевых сущностей в текстах, поэтому без проблем перестроит запрос на язык графовой БД и извлечет необходимые данные.

Source

Но можно пойти дальше и построить GraphRAG. Есть 2 возможности - работать с обычным текстом или с его эмбеддингами. В первом случае ничего не меняется - у нас есть узлы и отношения между сущностями в виде ребер в графе и все это в текстовом формате. Во втором случае все интереснее - узлы представлены эмбеддингами, как и связи между ними. А это значит можно проходить по графу для поиска наиболее релевантных на запрос сущностей (привет CosSim).

В этом случае LLM выполняет 3 действия:

Если работаем с эмбеддингами, то извлекаем топ-N семантически похожих узла из KG
С помощью LLM делаем запрос в KG для релевантных узлов, связанных с извлеченными сущностями
Формируем подграф контекста и формируем промпт для LLM вместе с запросом

Source

Фреймворки для построения графовых БД - ontotext, NebulaGraph и Neo4j

Hypothetical Document Embeddings

Представьте, что вы задаете довольно общий, но простой вопрос LLM (например "Какие самые продаваемые бургеры?"). Проблема заключается в том, что найти контекст из БД документов для детализации этого вопроса иногда бывает сложно из-за его общности. К тому же по контексту этого запроса в БД слишком много релевантных фрагментов и становится струдной задачей определить какие именно являются релевантными.

Но можно пойти по другому. Если до этого момента мы пытались загрузить в сеть наш запрос вместе с релевантной информацией, взятой из документов БД, то исследователи в работе Precise Zero-Shot Dense Retrieval without Relevance Labels предположили обратное - что если просить сеть сгенерировать первичный гипотетический ответ?

Source

Чтобы решить проблему "холодного старта", исследователи предложили брать первичный ответ из LLM на запрос пользователя, учитывая, что он может быть неточным. Но нам это и нужно - зная, что ответ содержит в себе лишь гипотетическую информацию (и возможно он не совсем верен), мы можем сравнить этот эмбеддинг с эмбеддингами правильных ответов из нашей БД и, благодаря похожему контексту, быстро найти точный ответ на наш запрос.

Улучшением такого подхода является генерация сразу нескольких гипотетических ответов и последующее их усреднение для нахождения более точного эмбеддинга из БД.

Необходимо уточнить, что это не работает, когда модели подается запрос, не являющийся wide-knowledge ("Объясни физический смысл лагранжиана квантовой хромодинамики") - в данном случае высока вероятность генерации галлюцинаций.

Contextual Compressors & Filters

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

Source

Суть проста - у нас есть базовый retriever, который выдает набор релевантных документов, а compressor сжимает их в короткие ответы и оставляет только необходимую информацию (только то, что полезно для ответа на вопрос).

Фильтром может быть что угодно - например стороняя fine-tuned LLM на эмейлы пользователей (картинка ниже).

Source

Вы даже можете построить собственный pipeline фильтра - например добавлять эмбеддинг важной информации (или формата ответа - JSON) в конец и только потом отправлять получившийся запрос в LLM.

Source

Multi-Query Retrieval

Метод мульти-запросов тоже довольно прост в своей реализации. Между Retriever-ом и пользовательским запросом находится еще одна LLM, которая генерирует множетсво вариантов первичного запроса, рассматривая его с разных точек зрения.

Source

Например на запрос "Расскажи о правлении Петра I" сеть может сгенерировать дополнительные вопросы "Экономическая стабильность во времена правления Петра I", "Войны и итоги во времена правления Петра I", "Градостроительство во времена Петра I" и так далее.

Все эти запросы подаются Retriver-у для нахождения фрагментов документов, соответствующих запросам, после чего отбираются релевантные и подаются в LLM вместе с основным запросом.

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

RAG-Fusion

Один из самых новых pipeline-ов RAG, который вышел обычной публикацией на TowardsDataScience, предлагает устранить вечный разрыв между тем, что пользователи явно задают в запросе и тем, что они собираются спрашивать. Ведь иногда пользователь хочет узнать много фактов о какой-то теме и просто вводит ее в виде запроса. С этой проблемой также сталкиваются все поисковые системы.

Первое - мы генерируем multi-query на основании запроса пользователя. По каждому запросу (включая оригинальный пользовательский) находятся релевантные фрагменты текста из векторной БД. Далее все результаты ранжируются и объединяются с помощью Reciprocal Rank Fusion (RRF).

Source

RRF (или взаимное слияние рангов) - это метод, который основан на объединении нескольких результатов поиска для получения единого унифицированного рейтинга. Один запрос не может охватить все аспекты запросов пользователей, и он может быть слишком узким для предоставления полных результатов. Вот почему при создании нескольких запросов необходимо учитывать все различные элементы и предоставлять тщательно подобранный ответ, с учетом ранга релевантности каждого фрагмента текста.

Ниже представлен пример. Четыре retrieval системы выдали отранжированные списки докумнтов. Далее для каждого документа в каждом случае вычисляется его обратный ранг. И после для каждого документа вычисляется финальное значение ранга - оно состоит из обратных сумм, где в знаменателе стоят два слагаемых - k и r(d). k - постоянный параметр, который обычно равняется 60. r(d) - вычисленный обратный ранг каждого документа в каждом случае. Далее суммы ранжируются по убыванию и мы получаем финальную выдачу документов по их релевантности. Необязательно подавать все фрагменты в контекст - можно брать top-K наиболее релевантных.

Source

Недостатки RAG-Fusion:

Существует риск выдачи слишком подробных ответов - в таком случае можно использовать Contextual Compressors & Filters
Нагрузка на контекстное окно модели

Multimodal RAG

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

Классическим примером где нужен RAG такого типа может быть поломка какой-то бытовой техники. Вы присылаете ее фотографию в LLM (например диодов на пылесосе) и формируете запрос с проблемой и просьбой ее решить. По фотографии и описанию, LLM получит от RAG контекст таких фотографий, который поможет определить модель техники, а также похожие ее проблемы. Фрагменты документов инструкций помогут понять как решить эту проблему - возможно ли починить дома или стоит обратиться в сервисный центр.

Работает Multimodal RAG аналогично векторному поиску по эмбеддингам текста, только теперь добавляется пространство эмбеддингов фотографий. Вот пример с конференции Google поиска похожих фотографий в векторном пространстве.

Source

Необходимо добавить, что эмбеддинги текстового описания картинки и самой картинки находятся рядом в векторном пространстве. Использование взаимосвязных эмбеддингов текста и картинки помогает LLM создавать более детальные генерации и уменьшает вероятность галлюцинаций, если БД картинок была тщательно отобрана и в ней не присутствуют асбтракции.

Source

Классический pipeline мультимодального RAG представлен на рисунке ниже. Все как обычно - берем запрос, векторным поиском находим релевантные фрагменты данных (текст и картинки), формируем на основе этого контекста ответ модели.

Source

Вот примеры генераций LLM на запрос "Imagining a day at the beach". Видим более детальное и насыщенное описание с использованием Multimodal RAG:

Generated without MM-RAG: I imagine a day at the beach would be very relaxing. I would sit in the sun and listen to the waves crash along the shore. Maybe I would go for a swim or build a sandcastle. It would be nice to get away from normal life for a while and enjoy the peaceful atmosphere.

Generated with MM-RAG: I imagine a day at the beach filled with golden sandy shores and the rhythmic crash of bright blue waves lapping gently at the coastline. My toes would sink into smooth sand as I breathed in the fresh and briny sea air under a bright sky, dotted with puffy clouds. I’d love to go snorkeling and glimpse the colorful fish swimming below or just nap on a towel, lulled into a rest by the glittering water and crying seagulls swooping overhead. Maybe later I could crack open a coconut or build an elaborate sandcastle with bridges and moats before taking a long walk at sunset, watching the glowing orange disk sink below the horizon.

На этом у меня все)

Спасибо, что дочитали до конца!

Пишите свои мысли в комментарии)