Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Исследователи из Renmin University of China и Tsinghua University предложили QAT-FM, новую схему сопоставления исходного шума и обучающих данных для Flow Matching. В препринте, не проходившем рецензирования, авторы утверждают, что она устраняет пересечения между траекториями из разных частей данных и требует меньше вычислений, чем полный оптимальный транспорт. Работа важна командам, которые обучают генеративные модели методом Flow Matching и выбирают между независимым сопоставлением и более дорогими схемами.
Что сделали
Flow Matching обучает модель поля скоростей, которое переводит выборку из исходного распределения, обычно гауссовского шума, в распределение данных. Для обучения берут пару из точки шума и объекта данных, проводят между ними линейную траекторию и требуют от модели предсказать её направление.
Результат зависит от схемы сопоставления пар (coupling). При независимом выборе шум и объект данных никак не согласованы, поэтому разные траектории могут пересекаться. В точке пересечения модель получает несколько несовместимых направлений и при обучении по среднеквадратичной ошибке усредняет их. Полный оптимальный транспорт, или OT, подбирает более согласованные пары, но для набора из N объектов размерности d авторы оценивают стоимость его построения как O(N³ + N²d).
В QAT-FM данные рекурсивно делят бинарным деревом. В каждом узле выбирают координату и порог, после чего разделяют объекты на две группы. Для гауссовского распределения строят дерево той же формы, но порог задают через соответствующий квантиль: доли вероятностной массы в дочерних узлах должны совпасть с долями объектов на стороне данных.
В результате каждому листу с объектами соответствует область гауссовского пространства с той же вероятностной массой. Обучающую пару получают, выбирая объект из листа и точку шума из связанной с ним области. При стандартном гауссовском распределении такие области являются прямоугольниками, а координаты можно независимо выбирать из усечённых нормальных распределений.
Авторы оценивают построение дерева как O(Nd log N), а получение одной пары как O(d). Архитектура модели, линейные траектории и функция потерь Flow Matching при этом не меняются. Для условной генерации дерево разрешено делить как по признакам объекта, так и по условию, например классу или текстовому представлению.
Что показали
Авторы доказали для своей конструкции два свойства. Траектории, ведущие к объектам из разных листьев дерева, не пересекаются. Кроме того, ожидаемое расстояние между промежуточными точками двух траекторий оказывается больше, чем при независимом сопоставлении, если разбиения дерева невырождены. Это снижает локальную неоднозначность целевой скорости, но само по себе не гарантирует лучшего качества обученной модели.
Эксперименты авторов согласуются с теоретической мотивацией. В условной генерации ImageNet QAT-FM получил FID 3,56 против 3,82 у независимого сопоставления. FID измеряет различие между признаками настоящих и сгенерированных изображений, поэтому меньшее значение считается лучшим.
В генерации CelebA по текстовому описанию авторы измерили FID 2,68 у QAT-FM и 2,98 у OT-FM. При этом показатель соответствия изображения тексту остался на уровне других условных методов из сравнения. По представленным абляциям результат также мало зависел от выбранного правила разделения дерева и параметров предварительного сокращения размерности.
Ограничения
Гарантия отсутствия пересечений относится только к траекториям из разных листьев. Внутри одного листа остаётся независимое сопоставление, поэтому работа не исключает близкие или пересекающиеся пути для всех обучающих пар.
Метод проверяли на синтетических данных и задачах генерации изображений: без условий, по классу и по тексту. Крупнейший описанный набор содержит более миллиона обучающих изображений, а качество условных моделей оценивали по 50 тысячам генераций. Латентную модель высокого разрешения из-за вычислительных ограничений обучали 64 эпохи.
На крупном латентном эксперименте QAT-FM сравнивали только с независимым сопоставлением. Поэтому из работы нельзя определить, как метод соотносится в этом масштабе с мини-пакетным OT, SD-FM и другими структурированными схемами. Сравнение вычислительной сложности охватывает построение сопоставления и выбор пар, но не показывает, какую долю полного времени обучения модели эти операции занимают в разных инфраструктурах.
Что это значит
Для команды, уже использующей линейный Flow Matching, работа предлагает локальную замену одного компонента: способа формирования обучающих пар. Если независимое сопоставление создаёт неоднозначные траектории, а полный OT не помещается в вычислительный бюджет, QAT-FM стоит добавить в экспериментальный план. Переделывать сеть и основную функцию потерь для этого не требуется, однако понадобятся построение дерева, хранение границ листьев и выборка из усечённого гауссовского распределения.
Планы по архитектуре продукта эта работа пока не меняет. Она не сравнивает QAT-FM с другими классами генеративных моделей и не доказывает сокращение полной стоимости обучения или вывода. Практическое решение следует принимать после проверки на собственных данных: сравнить качество, время подготовки сопоставления, устойчивость дерева к размерности признаков и корректность связи между объектами и условиями.
Если команда только выбирает базовый подход, QAT-FM добавляет аргумент в пользу Flow Matching для задач, где требуется глобально согласовать большой набор шума и данных без полного OT. Если же сопоставление пар не является измеренным узким местом, оснований менять текущий план по одному этому препринту нет.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
