Объяснение обучения матрёшечных представлений: метод, лежащий в основе эффективных текстовых эмбеддингов OpenAI
Разработка реальных моделей машинного обучения всегда предполагает компромисс между стоимостью и производительностью. Например, чем больше модель и набор данных, которые мы используем во время обучения, тем более способной в итоге будет наша обученная модель. Однако эта повышенная способность сопровождается более длительным временем и более высокими вычислительными затратами, поскольку обучение требует значительно больше времени. Тот же принцип применим и во время инференса, где более крупные модели, как правило, генерируют более крупные представления признаков, требующие больше памяти для хранения.
Поскольку потребности при обучении и инференсе машинного обучения могут значительно различаться от одного сценария использования к другому, важно иметь метод, который позволяет пожертвовать небольшой частью производительности модели в обмен на снижение затрат. Именно здесь вступает в игру такой метод, как Matryoshka Representation Learning (MRL) . Например, модель OpenAI text-embedding-3-small использует MRL, чтобы дать разработчикам возможность сокращать embeddings, сохраняя их основные свойства представления концепций. Позволяя настраивать размерности embedding, MRL помогает найти идеальный баланс между экономической эффективностью и производительностью модели.
В этой статье мы рассмотрим, как работает MRL, его реализацию и то, как он обеспечивает масштабируемые и эффективные модели машинного обучения. Начнем с мотивации, лежащей в основе MRL.
Мотивация, лежащая в основе Matryoshka Representation Learning (MRL)
Компромисс между стоимостью и производительностью всегда учитывается перед разработкой или использованием любых моделей машинного обучения.
Во время обучения параметры модели и объем обучающих данных напрямую влияют на ее итоговую производительность. Более крупные модели и обучающие наборы данных, как правило, приводят к созданию более способных моделей. Однако вычислительные затраты на обучение таких моделей также выше.
Между тем во время инференса более крупные модели создают более крупные представления признаков. Это увеличение представлений признаков означает, что нам требуется больше памяти для их хранения.
Рассмотрим BERT в качестве примера в контексте информационного поиска. BERT — одна из первых моделей глубокого обучения на основе Transformer, которая достигла передовых результатов на нескольких эталонных наборах данных в классификации текста, распознавании именованных сущностей (NER), ответах на вопросы (QnA) и других задачах.
Рисунок: рабочий процесс генерации embedding с использованием базовой модели BERT.
Вкратце, BERT принимает на вход предложение или слово и преобразует его в embedding фиксированного размера. Этот embedding отражает семантический смысл исходного входа, а его размер зависит от конкретного используемого варианта модели.
У BERT есть два основных варианта: базовая модель и большая модель. Базовая модель создает embedding размерностью 768, тогда как большая модель генерирует embedding размерностью 1024.
Допустим, мы хотим хранить 10 миллионов embeddings в векторной базе данных для сценария информационного поиска. Если мы используем базовую модель BERT в формате FP32, нам потребуется примерно 768×10M×4=30.72 ГБ памяти для хранения этих embeddings. Напротив, использование большой модели BERT в том же сценарии потребовало бы около 40.9 ГБ. Кроме того, больший размер embedding может замедлить процесс поиска за счет увеличения вычислительной сложности поиска по сходству. Однако основное преимущество использования более крупных embeddings заключается в их повышенной способности извлекать высокорелевантную информацию по сравнению с более мелкими embeddings.
Одно из идеальных решений этих проблем в информационном поиске — разделение всего процесса на две части: формирование короткого списка и перера ранжирование. Формирование короткого списка означает извлечение первоначального набора документов-кандидатов из огромной коллекции в нашей векторной базе данных, такой как Milvus. Затем перера ранжирование берет этих кандидатов из короткого списка и переупорядочивает их, чтобы максимизировать релевантность итоговых результатов.
Рисунок — рабочий процесс формирования короткого списка и перера ранжирования
Рисунок: рабочий процесс формирования короткого списка и перера ранжирования.
Для формирования короткого списка алгоритм вычисляет сходство между эмбеддингами запроса и большой коллекцией эмбеддингов, хранящихся в векторной базе данных. Поэтому использование меньших размеров эмбеддингов может сделать это вычисление более эффективным и быстрым. И наоборот, перера ранжирование делает акцент на точности, а значит, нам были бы полезны большие размеры эмбеддингов.
Проблема этого метода заключается в том, что после выбора конкретной модели размер эмбеддингов или представлений признаков, которые она создает, фиксирован. Эта нехватка гибкости ограничивает нашу возможность использовать эмбеддинги разных размеров из одной модели.
Рисунок — рабочий процесс модели, генерирующей эмбеддинг фиксированного размера..png
Рисунок: рабочий процесс модели, генерирующей эмбеддинг фиксированного размера.
Подход Matryoshka Representation Learning (MRL) предлагает интересное решение этой проблемы: модели, обученные этим методом, могут создавать эмбеддинги различных размеров. Мы подробно рассмотрим механизм этого подхода в следующем разделе.
Что такое Matryoshka Representation Learning (MRL)?
Matryoshka Representation Learning (MRL) — это метод обучения нейронных сетей созданию многоуровневых представлений в рамках одной модели. Вдохновленный матрешками, где меньшие куклы помещаются внутри больших, MRL позволяет модели выдавать представления различного размера (от грубых до детальных) за один прямой проход. Этот подход позволяет моделям глубокого обучения лучше понимать сложные взаимосвязи и нюансы в данных. MRL особенно эффективен для таких задач, как семантический поиск, информационный поиск, многоязычная обработка и любое приложение, требующее нюансированных представлений данных на разных уровнях абстракции.
Визуализация эмбеддингов Matryoshka с несколькими уровнями детализации
Рисунок: Визуализация эмбеддингов Matryoshka с несколькими уровнями детализации
Популярные модели эмбеддингов, в которых принят подход Matryoshka Representation Learning (MRL), включают OpenAI’s text-embedding-3-large, Nomic’s nomic-embed-text-v1 и Alibaba’s gte-multilingual-base.
Как работает подход MRL
Подход MRL позволяет нам извлекать представления признаков из любой модели машинного обучения в нескольких размерах. Например, вместо использования исходных 1024 измерений эмбеддинга мы можем использовать первые 16, 32, 64, 128 или 256 измерений (или любые измерения, которые выберем). Ключ к этой возможности заключается в том, как обучается модель при реализации подхода MRL.
Рисунок — обучение функции потерь MRL и пример ее использования во время инференса
Рисунок: обучение функции потерь MRL и пример ее использования во время инференса. Источник.
Во время обучения модели с MRL мы оптимизируем не одну функцию потерь, как это обычно бывает при стандартном обучении модели, а несколько. Каждая функция потерь нацелена на оптимизацию представления признаков для конкретной размерности, которую мы заранее определяем. Если мы зададим пять разных размерностей признаков, нам нужно будет оптимизировать пять отдельных функций потерь во время обучения.
Допустим, мы хотим обучить большую модель BERT с использованием подхода MRL. Как упоминалось в предыдущем разделе, эта модель генерирует эмбеддинг размером 1024. Если мы зададим размерности признаков 32, 64, 128, 256 и 1024, мы будем оптимизировать пять разных функций потерь во время обучения, чтобы модель научилась генерировать оптимизированные эмбеддинги в этих размерностях.
Оптимизация по этим различным размерностям довольно проста: MRL разбивает общую функцию потерь на сумму потерь для каждой отдельной размерности. Используя наш пример выше, это можно выразить так:
Процесс оптимизации для модели, обученной с помощью MRL, можно более формально представить следующим образом:
Рисунок: Уравнение оптимизации подхода MRL. Источник.
Как видите, приведенное выше определение функции потерь является довольно общим. Поэтому подход MRL можно применять практически к любой модели, и он полностью не зависит от архитектуры модели. Мы также можем дообучать предварительно обученные модели, такие как BERT, или любые другие модели на основе трансформеров, чтобы они выдавали эмбеддинги различных размеров. Продолжая наш сценарий выше, теперь мы можем генерировать 1024-мерные эмбеддинги из нашей обученной большой модели BERT, а также 32-, 64-, 128- и 256-мерные эмбеддинги.
В целом после обучения модели с MRL начальные размерности признаков, генерируемых обученной моделью, несут более значимую информацию, чем последующие размерности. Первые несколько размерностей содержат высокоуровневые детали, тогда как последующие размерности сосредоточены на более детализированной информации, имитируя структуру матрёшки.
Однако это не означает, что более короткие эмбеддинги всегда являются просто усеченными версиями более длинных. Значения каждого элемента в более коротких эмбеддингах могут отличаться от значений в более длинных эмбеддингах из-за коэффициента масштабирования, применяемого к каждой размерности признаков во время обучения. Однако если мы зададим одинаковый коэффициент масштабирования в каждой размерности, значения каждого элемента между более короткими и более длинными эмбеддингами могут выглядеть идентичными.
Экспериментальные результаты MRL
Подход MRL был оценен на моделях машинного обучения в различных модальностях, включая текст, зрение и зрение-текст. Модели ResNet50 и ViT представляют модели на основе зрения, BERT представляет модели на основе текста, а ALIGN представляет сочетание зрения и текста. Эти модели в основном оценивались для двух распространенных сценариев использования: классификации и поиска. Сначала обсудим классификацию.
При сравнении производительности модели ResNet50, обученной на наборе данных ImageNet-1K, с независимо обученной стандартной ResNet50 модель MRL достигает сопоставимой top-1 точности при различных размерах представления признаков.
Чтобы дополнительно оценить полезность представлений признаков для последующих задач, также была измерена точность 1-ближайшего соседа (1-NN) для каждого размера представления признаков. Настройка для вычисления 1-NN выглядит следующим образом: при наличии 1,3 тыс. образцов изображений в базе данных задача состоит в том, чтобы найти ближайшего соседа для каждого из 50 тыс. запросов. При такой настройке ResNet50, обученная с MRL, оказывается до 2% точнее при каждом размере представления признаков по сравнению со своим аналогом с фиксированными признаками.
Figure- Top-1 accuracy and 1-NN accuracy of ResNet50 on ImageNet-1K
Рисунок: точность Top-1 и точность 1-NN ResNet50 на ImageNet-1K. Источник.
Между тем производительность модели ViT, обученной с MRL на наборе данных JFT-300M, также очень конкурентоспособна для всех размеров представлений. Ее точность 1-NN сопоставима с точностью ViT, обученной с представлениями признаков фиксированного размера. Как показано на рисунке ниже, производительность модели MRL также выше, чем у ее аналога с фиксированным размером при меньших представлениях признаков, отчасти потому, что случайные признаки выбираются из модели фиксированного размера для представления меньших размерностей. Похожая тенденция наблюдается и у модели ALIGN, обученной с использованием подхода MRL; ее производительность соответствует производительности модели ALIGN, обученной с представлениями фиксированного размера.
Figure- 1-NN accuracy of ViT and ALIGN on ImageNet-1K
Рисунок: точность 1-NN ViT и ALIGN на ImageNet-1K. Источник.
Главное преимущество использования подхода MRL заключается в его гибкости при переключении между разными размерами представлений признаков из одной и той же модели. Чтобы в полной мере использовать это преимущество, также были проведены тесты адаптивной классификации с использованием ResNet50 на ImageNet-1K.
Эта настройка включает обучение максимальной вероятности softmax для определения переходов от меньших представлений признаков к большим. Результаты тестов показывают, что модель MRL с 37-мерным представлением признаков столь же точна, как модель ResNet50, обученная с фиксированным 512-мерным представлением признаков, при этом уступая 2048-мерной модели всего на 0,8%.
Figure- Top 1 accuracy of ResNet50 model in adaptive classification (left) and mAP of ResNet50 in image retrieval use case (right)
Рисунок: точность Top 1 модели ResNet50 в адаптивной классификации (слева) и mAP ResNet50 в сценарии поиска изображений (справа). Источник.
Далее обсудим поиск. Как вы, возможно, знаете, цель поиска — найти набор похожих контекстов для заданного запроса. Поэтому было протестировано качество поиска между ResNet50, обученной с MRL, и представлениями признаков фиксированного размера с использованием mean Average Precision (mAP). Модель MRL демонстрирует улучшение mAP до 3% по сравнению со своим аналогом с фиксированным размером при каждом размере представления признаков, как показано на визуализации выше.
Поскольку с MRL мы можем использовать несколько размеров представления признаков, особенно интересно изучить, как можно ускорить процесс поиска, сохраняя точность с MRL. Мы можем протестировать это в сценарии, называемом адаптивным поиском.
Как упоминалось ранее, в процессе поиска есть два основных этапа: предварительный отбор и переранжирование. Настройка адаптивного поиска включала использование 16-мерного представления признаков для предварительного отбора набора из 200 кандидатов, тогда как 2048-мерное представление признаков используется для переранжирования. Этот подход сравнивается с другим методом, который использует 2048-мерное представление признаков как для предварительного отбора, так и для переранжирования.
В ImageNet-1K использование модели MRL теоретически может ускорить процесс поиска до 128 раз. Мы называем это теоретическим ускорением, потому что в реальных приложениях мы редко используем наивный алгоритм ближайшего соседа для задач поиска. Вместо этого для повышения скорости поиска обычно применяются алгоритмы приближенного ближайшего соседа (ANN), такие как FAISS, ANNOY или HNSW. В этом случае конфигурация с использованием модели MRL достигает 14-кратного ускорения по сравнению с процессами поиска, использующими алгоритм HNSW на идентичном аппаратном обеспечении.
Похожий результат наблюдался с датасетом ImageNet-4K, где авторы статьи использовали модель MRL с 64-мерным представлением признаков для предварительного отбора и 2048-мерным представлением для повторного ранжирования. Результаты показывают теоретическое ускорение в 32 раза и реальное ускорение в шесть раз.
Компромисс между mAP@10 и MFLOPs/Query для Adaptive Retrieval (AR) на ImageNet-1K (слева) и ImageNet-4K (справа).
Компромисс между mAP@10 и MFLOPs/Query для Adaptive Retrieval (AR) на ImageNet-1K (слева) и ImageNet-4K (справа). Источник.
Реализация подхода MRL
В этом разделе мы рассмотрим простую реализацию подхода MRL. В частности, мы будем использовать базовую модель MPNet, обученную с подходом MRL на датасете NLI, для генерации эмбеддингов разных размеров. Как вы, возможно, уже знаете, исходная базовая модель MPNet может генерировать только 768-мерный эмбеддинг. Однако эта обученная модель MPNet может генерировать эмбеддинги размерностью 768, 512, 256, 128 и 64, где вес каждого измерения во время обучения установлен равным 1.
Мы можем использовать эту модель с помощью библиотеки SentenceTransformers следующим кодом.
from sentence_transformers import SentenceTransformer
matryoshka_dim_short = 64
matryoshka_dim_long = 768
text = ["The weather is so nice!"]
short_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_short).encode(text)
long_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_long).encode(text)
print(f"Shape: {short_embedding.shape, long_embedding.shape}")
print(short_embedding[0][0:10])
print(long_embedding[0][0:10])
"""
Output:
Shape: ((1, 64), (1, 768))
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
"""
В приведенном выше коде мы используем обученную базовую модель MPNet для генерации эмбеддингов двух разных размеров: 64 и 768. Поскольку вес или коэффициент масштабирования каждого измерения установлен равным 1, два эмбеддинга имеют идентичные элементы, как вы можете видеть в первых 10 элементах двух эмбеддингов выше.
Поскольку значение каждого элемента идентично, сходство между двумя эмбеддингами будет равно 1.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
Если вы хотите обучить собственную модель с подходом MRL, вы можете ознакомиться с официальным репозиторием GitHub этого метода.
Заключение
Подход MRL предлагает решение для баланса между стоимостью и производительностью в машинном обучении. Позволяя любой модели машинного обучения создавать представления признаков различных размеров, MRL обеспечивает гибкость для оптимизации либо по скорости, либо по точности в зависимости от нашего сценария использования и ресурсов. Мы увидели, что такая адаптивность ценна в приложениях, таких как классификация и поиск, где переключение между меньшими и большими представлениями может значительно повысить эффективность без существенного ущерба для производительности.
Экспериментальные результаты показывают, что MRL соответствует точности традиционных моделей фиксированного размера и часто превосходит ее в различных областях, включая текст, компьютерное зрение и мультимодальные задачи. Кроме того, совместимость подхода с существующими архитектурами, такими как ResNet50, ViT и BERT, подчеркивает его применимость к различным модальностям моделей. Благодаря потенциалу для существенного ускорения в реальных сценариях и улучшенному качеству поиска MRL представляет собой многообещающий шаг вперед к более эффективным и универсальным решениям машинного обучения.
Связанные ресурсы
Читать далее
What Is a Vector Lakebase?
A Vector Lakebase is a unified, lake-native data architecture for AI that combines vector-database-grade serving with open lake storage, reusable lake-level indexes, and a shared semantic layer.
What is the K-Nearest Neighbors (KNN) Algorithm in Machine Learning?
KNN is a supervised machine learning technique and algorithm for classification and regression. This post is the ultimate guide to KNN.
Building RAG Pipelines for Real-Time Data with Cloudera and Milvus
explore how Cloudera can be integrated with Milvus to effectively implement some of the key functionalities of RAG pipelines.