Telegram-канал gonzo_ml - gonzo-обзоры ML статей: Unsorted

gonzo-обзоры ML статей

29 Dec 2024 04:34

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

До некоторой степени LCM напоминает Лекуновскую JEPA (https://openreview.net/pdf?id=BZ5a1r-kVsf), которая тоже предсказывает репрезентации следующих наблюдений в пространстве эмбеддингов. Но JEPA фокусировалась на выучивании такого пространства в режиме self-supervised, а LCM фокусируется на точном предсказании в уже существующем пространстве эмбеддингов (но срастить эти два подхода, наверное, тоже просится).

Итак, работаем в пространстве эмбеддингов, для обучения LCM нужно подготовить текстовый датасет и сконвертировать его через SONAR в эмбеддинги, один на каждое предложение. На практике это не так просто, точная сегментация не всегда легка из-за ошибок в датасете или специфического форматирования. Кроме того длинные предложения могут быть слишком сложны для кодирования/декодирования через SONAR, качество пострадает. В итоге для разбиения на предложения выбрали Segment any Text (SaT, https://github.com/segment-any-text/wtpsplit) с дополнительным ограничением на длину сегмента, всё длиннее 250 символов (это число мы ещё встретим) разбивается, этот метод называют SaT Capped.

LCM должна conditionally на основе контекста генерить непрерывный эмбеддинг. Это отличается от работы LLM, где надо выдать распределение вероятности по дискретным токенам словаря. Прямолинейный подход к снаряду был бы обучить трансформер генерить эмбеддинг с objective минимизации MSE лосса. Это будет называться Base-LCM. Это не так просто, потому что у заданного контекста может быть много подходящих, но семантически разных продолжений, это видно на примере генерации картинок диффузионками, из одного промпта получаются довольно разные изображения. И в целом в той области как раз много наработок по выучиванию conditional probability distributions для непрерывных данных, поэтому другой логичный вариант на попробовать — это диффузионная модель, Diffusion-based LCM. Наконец ещё один вариант — квантизация и возврат к задаче генерации дискретных элементов, Quantized LCM.

Пройдёмся подробнее по вариантам LCM.

Base-LCM выступает бейзлайном, это стандартный декодер трансформера, переводящий последовательность предшествующих концептов (эмбеддингов предложений) в последовательность будущих. Трансформер окружён со стороны входа и выхода двумя простыми сетями, PreNet и PostNet, занимающимися нормализацией/денормализацией и проекцией SONAR эмбеддингов в размерность модели и из неё обратно. Обучается на semi-supervised задаче предсказания следующего концепта, минимизируя MSE loss относительно ground truth. К документам обучающей выборки добавляют суффикс “End of text”, что даёт возможность научиться генерить документы переменной длины. В инференсе один из стоп-критериев проверяет близость сгенерённого эмбеддинга эмбеддингу этого суффикса и останавливает генерацию, если близость выше заданного порога; другой стоп-критерий смотрит на косинусную близость между текущим и предыдущим эмбеддингом и останавливает, если она выше порога (оба порога установлены в 0.9).

Diffusion-based LCM тоже авторегрессионно генерит концепты, один за раз, делая заданное количество denoising steps для каждого генерируемого концепта. Использует classifier-free diffusion guidance. Есть версия модели с одним стволом (One-Tower) и двумя (Two-Tower). В первом случае это один ствол трансформера, делающий всё. Во втором отдельный ствол (contextualizer) занимается кодированием предшествующего контекста, а второй (denoiser) генерит новый эмбеддинг концепта и через cross-attention смотрит на контекст из первого ствола.

Quantized LCM использует Residual Vector Quantization и далее работает аналогично обычным LLM, предсказывающим дискретные юниты. Здесь можно использовать температуру, top-p/top-k. Архитектуру стараются собрать максимально похожую на Diffusion-based LCM, чтобы было легче сравнивать.