Meta AI Представляет Adjoint Sampling: Масштабируемое Генеративное Моделирование Без Данных

Проблема Недостатка Данных в Генеративном Моделировании

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

Adjoint Sampling от Meta AI

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

Этот метод эффективен, когда доступна только не нормализованная энергетическая функция. Он генерирует образцы, соответствующие целевому распределению, избегая затратных вычислительных методов, таких как importance sampling или MCMC.

Технические Основы Adjoint Sampling

Основой алгоритма является стохастическое дифференциальное уравнение (SDE), описывающее эволюцию траекторий образцов. Алгоритм обучается управлять дрейфом u(x, t), чтобы конечные состояния траекторий приближались к нужному распределению, например, распределению Больцмана. Ключевое нововведение — функция потерь Reciprocal Adjoint Matching (RAM), которая позволяет выполнять градиентные обновления, используя только начальные и конечные состояния траекторий, избегая обратного распространения через весь путь диффузии и значительно повышая вычислительную эффективность.

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

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

Результаты и Эффективность

Adjoint Sampling демонстрирует лучшие результаты в синтетических и реальных задачах. На синтетических бенчмарках, таких как Double-Well (DW-4) и потенциалы Леннарда-Джонса (LJ-13 и LJ-55), он значительно превосходит методы DDS и PIS, особенно по энергоэффективности. Например, там, где DDS и PIS требуют 1000 вычислений энергии на шаг градиента, Adjoint Sampling достигает аналогичных или лучших результатов всего с тремя вычислениями, оцениваемыми по расстоянию Вассерштейна и эффективному размеру выборки (ESS).

На практике алгоритм оценивался на крупномасштабной генерации молекулярных конформеров с использованием энергетической модели eSEN, обученной на датасете SPICE-MACE-OFF. Вариант с декартовыми координатами и предобучением достиг 96.4% recall и 0.60 Å среднеквадратичного отклонения, превзойдя широко используемый базовый метод RDKit ETKDG по всем метрикам. Также метод хорошо обобщается на датасет GEOM-DRUGS, показывая значительный рост recall при сохранении конкурентной точности.

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

Новый Взгляд на Генеративное Моделирование с Наградами

Adjoint Sampling — важный шаг вперед для генеративного моделирования без данных. Используя скалярные сигналы награды и эффективный on-policy метод обучения, основанный на стохастическом управлении, он обеспечивает масштабируемую тренировку диффузионных сэмплеров с минимальными вычислениями энергии. Поддержка геометрических симметрий и способность к обобщению делают его ключевым инструментом в вычислительной химии и смежных областях.