Симметрично-осведомленные GFlowNets

Основная информация

• ID статьи: 2506.02685
• Название: Symmetry-Aware GFlowNets
• Авторы: Hohyun Kim, Seunggeun Lee, Min-hwan Oh (Сеульский национальный университет)
• Классификация: stat.ML cs.LG
• Конференция: ICML 2025 (42-я Международная конференция по машинному обучению)
• Ссылка на статью: https://arxiv.org/abs/2506.02685

Аннотация

Генеративные потоковые сети (GFlowNets) предоставляют мощный фреймворк для выборки графов пропорционально вознаграждению. Однако существующие методы страдают от систематических смещений из-за неточности расчета вероятностей переходов состояний. Эти смещения коренятся в присущей графам симметрии и влияют как на схемы генерации на основе атомов, так и на основе фрагментов. Для решения этой проблемы в статье представлены симметрично-осведомленные GFlowNets (SA-GFN), которые интегрируют коррекцию симметрии в процесс обучения посредством масштабирования вознаграждения. Путем прямого встраивания коррекции смещения в структуру вознаграждения SA-GFN устраняет необходимость в явных расчетах переходов состояний. Экспериментальные результаты демонстрируют, что SA-GFN достигает несмещенной выборки, одновременно повышая разнообразие и постоянно генерируя графы с высоким вознаграждением, которые тесно соответствуют целевому распределению.

Исследовательский контекст и мотивация

Основная проблема

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

Важность проблемы

1. Смещение при генерации молекул: В открытии лекарств более 50% молекул обладают множественными симметриями, 18% содержат 4 или более симметрии. Игнорирование симметрии приводит к неправильному моделированию и снижению точности генерации молекулярных структур.
2. Систематическое смещение: Смещение систематично, при генерации узлов оно смещает в сторону графов с меньшей симметрией, при генерации фрагментов — в сторону компонентов с симметрией.
3. Вычислительная сложность: Точный расчет вероятностей переходов требует дорогостоящих тестов графовой изоморфности.

Ограничения существующих методов

• Ma et al. (2024) предложили использование позиционного кодирования для приблизительного обнаружения эквивалентных действий, но требует применения при каждом переходе, имеет большие вычислительные затраты и является только приблизительным решением.
• Традиционные целевые функции GFlowNet (TB, DB и т.д.) не могут избежать проблемы эквивалентных действий, так как основаны на формализации переходов состояний.

Основные вклады

1. Теоретический вклад: Предоставляет строгую формализацию авторегрессивной генерации графов в рамках GFlowNet, явно решая проблему эквивалентных действий
2. Простое и эффективное решение: Предлагает метод масштабирования вознаграждения на основе размера группы автоморфизмов, требующий минимальных изменений существующих алгоритмов обучения
3. Несмещенный оценитель: Выводит несмещенный оценитель правдоподобия модели
4. Экспериментальная верификация: Экспериментально проверяет теоретические результаты, демонстрируя эффективность метода в генерации разнообразных высокооцениваемых образцов

Подробное описание метода

Определение задачи

Дано функция вознаграждения R(x), целью GFlowNets является обучение политики pA таким образом, чтобы вероятность выборки терминального состояния была пропорциональна его вознаграждению: p̄A(x) = R(x)/Z, где Z — нормализующая константа.

Основной теоретический фреймворк

1. Графовая изоморфность и отношения эквивалентности

• Графовая изоморфность: Два графа G и G' изоморфны (G ≅ G'), если существует перестановка π такая, что π(E) = E'
• Группа автоморфизмов: Группа автоморфизмов графа G — Aut(G) — это множество всех перестановок, сохраняющих структуру графа
• Орбита: Орбита узла u — Orb(G,u) = {v ∈ V : ∃π ∈ Aut(G), π(u) = v}

2. Формализация эквивалентных действий

Определение 4.1 (Эквивалентность переходов): Графовые переходы (G₁,G'₁) и (G₂,G'₂) являются переходно-эквивалентными, если G₁ ≅ G₂ и G'₁ ≅ G'₂.

Определение 4.2 (Орбитальная эквивалентность): Графовые действия (G₁,t₁,u₁) и (G₂,t₂,u₂) являются орбитально-эквивалентными, если типы действий совпадают и существует перестановка π такая, что π(G₁) = G₂ и π(u₁) = u₂.

Теорема 4.3: Орбитально-эквивалентные действия приводят к переходно-эквивалентным переходам.

3. Ключевые теоретические результаты

Лемма 4.5: Для действия AddEdge выполняется $\frac{|\text{Orb}(G,u,v)|}{|\text{Orb}(G',u,v)|} = \frac{|\text{Aut}(G)|}{|\text{Aut}(G')|}$

Теорема 4.6 (Коррекция автоморфизмов): При использовании перестановочно-эквивариантных функций $\frac{p_{\bar{A}}(a|s)}{q_{\bar{A}}(a|s')} = \frac{|\text{Aut}(G)|}{|\text{Aut}(G')|} \cdot \frac{p_E(G'|G)}{q_E(G|G')}$

Метод симметрично-осведомленной коррекции

1. Масштабирование вознаграждения при генерации узлов

Следствие 5.1 (Коррекция TB): Потеря траекторного баланса должна быть $L_{TB}(\tau) = \left(\log \frac{Z\prod_{t=0}^{n-1} p_E(G_{t+1}|G_t)}{|\text{Aut}(G_n)|R(G_n)\prod_{t=0}^{n-1} q_E(G_t|G_{t+1})}\right)^2$

Решение: Масштабировать вознаграждение как $\tilde{R}(G) = |\text{Aut}(G)|R(G)$

2. Коррекция при генерации фрагментов

Теорема 5.3 (Коррекция фрагментов): Для терминального состояния G, сгенерированного соединением k фрагментов {C₁,...,Cₖ}: $\tilde{R}(G) = \frac{|\text{Aut}(G)|R(G)}{\prod_{i=1}^k |\text{Aut}(C_i)|}$

3. Несмещенный оценитель правдоподобия модели

$\bar{p}_A(x) = \mathbb{E}_{\tau \sim q_E(\tau|G_n)}\left[\frac{p_E(\tau)}{|\text{Aut}(G_n)|q_E(\tau|G_n)}\right]$

Технические инновации

1. Теоретическая элегантность: Упрощает сложную коррекцию на уровне переходов до масштабирования вознаграждения терминального состояния
2. Вычислительная эффективность: Избегает дорогостоящих тестов графовой изоморфности на каждом шаге, требуя расчета размера группы автоморфизмов только один раз
3. Универсальность: Применима к множеству целевых функций GFlowNet (TB, DB, FM и т.д.)
4. Точность: Предоставляет точное решение, а не приблизительное

Экспериментальная установка

Наборы данных

1. Иллюстративные примеры: Начальное состояние с 6 отключенными узлами, 112 терминальных состояний
2. Синтетические графы: Гетерогенные графы с максимум 7 узлами, 72,296 терминальных состояний
3. Генерация молекул:
   • На уровне атомов: Задача предсказания зазора HOMO-LUMO
   • На уровне фрагментов: Задача предсказания энергии связывания мишени sEH

Метрики оценки

• Ошибка L1: Ошибка L1 между целевой вероятностью и вероятностью терминального состояния модели
• Разнообразие: Среднее расстояние Танимото между молекулами
• Метрики Top K: Разнообразие и вознаграждение K молекул с наивысшим вознаграждением
• Уникальность: Доля уникальных молекул в сгенерированных образцах

Методы сравнения

1. Vanilla GFlowNets: Без учета графовой симметрии
2. Transition Correction: Идентификация переходно-эквивалентных действий посредством множественных тестов изоморфности
3. PE (Ma et al., 2024): Использование позиционного кодирования для приблизительной идентификации орбитально-эквивалентных действий
4. Reward Scaling (данная работа): Коррекция путем модификации сигнала вознаграждения
5. Flow Scaling (данная работа): Умножение на коэффициент симметрии при каждом переходе

Результаты экспериментов

Основные результаты

1. Иллюстративные эксперименты

• Вероятности терминального состояния Vanilla модели группируются по |x|, демонстрируя явное смещение
• Reward Scaling достигает такой же производительности, как Transition Correction
• Оцененная нормализующая константа Z: Reward Scaling — 112 (истинное значение), Vanilla — 26,706

2. Эксперименты на синтетических графах

• Целевая функция TB: Reward Scaling значительно снижает ошибку L1, производительность сравнима с Transition Correction
• Целевая функция DB: Reward Scaling сходится медленнее, но в итоге достигает той же точности
• Метод PE как приблизительное решение показывает производительность между Vanilla и точными методами

3. Эксперименты по генерации молекул

Результаты генерации на уровне атомов:

• Разнообразие: 0.929→0.959 (Vanilla→Reward Scaling)
• Уникальность: 0.93→1.0

Результаты генерации на уровне фрагментов:

• Вознаграждение Top K: 0.941→0.952 (Vanilla→Reward Scaling Exact)
• Использование фрагмента циклогексана: 5220→1042 (значительное снижение чрезмерного использования симметричных фрагментов)

Абляционные эксперименты

• Приблизительная коррекция vs точная коррекция: Приблизительный метод уже значительно улучшает производительность
• Различные целевые функции: Как TB, так и DB могут эффективно корректироваться посредством масштабирования вознаграждения

Анализ вычислительной эффективности

• Время расчета автоморфизмов: 0.010 мс на наборе данных QM9, 0.022 мс на ZINC250k
• По сравнению с прямым распространением нейронной сети при выборке траектории вычислительные затраты пренебрежимы
• Сравнение времени обучения: Reward Scaling примерно на 15% быстрее, чем Transition Correction

Связанные работы

Авторегрессивная генерация графов

• Методы матрицы смежности: Сохраняют информацию о порядке узлов, менее подвержены проблеме эквивалентных действий
• Методы последовательности графов: Легко создают эквивалентные действия, проблема становится острой при необходимости расчета вероятностей переходов состояний

GFlowNets

• Существующие целевые функции (потоковое соответствие, детальный баланс, траекторный баланс и т.д.) не могут избежать проблемы эквивалентных действий
• Ma et al. (2024) впервые выявили проблему, но предоставили только приблизительное решение

Выразительная способность графовых нейронных сетей

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

Заключение и обсуждение

Основные выводы

1. Теоретический вклад: Впервые строго анализирует проблему эквивалентных действий в GFlowNets, доказывая, что она приводит к систематическим смещениям
2. Практическое решение: Масштабирование вознаграждения предоставляет простой, точный и эффективный метод коррекции
3. Широкая применимость: Метод применим к генерации на уровне атомов и фрагментов, а также к различным целевым функциям обучения

Ограничения

1. Зависимость от дизайна действий: Теоретические гарантии зависят от конкретного набора предопределенных графовых действий
2. Специфичность задачи: Главным образом проверена на наборах данных, связанных с открытием лекарств
3. Выразительная способность GNN: Ограниченная выразительная способность GNN может привести к дополнительным смещениям

Направления будущих исследований

1. Исследование задач с различными паттернами симметрии и структурами вознаграждения
2. Разработка архитектур GNN с повышенной выразительной способностью
3. Расширение на более сложные сценарии генерации графов

Глубокая оценка

Преимущества

1. Теоретическая строгость: Предоставляет полный математический фреймворк и строгий теоретический анализ
2. Простота метода: Решение чрезвычайно простое, легко реализуется и интегрируется
3. Практическая ценность: Демонстрирует реальный эффект в важных приложениях, таких как генерация молекул
4. Вычислительная эффективность: Избегает дорогостоящих онлайн-тестов графовой изоморфности
5. Сильная универсальность: Применима к множеству целевых функций обучения GFlowNet

Недостатки

1. Ограниченный диапазон экспериментов: Главным образом сосредоточена на задачах генерации молекул, ограниченная верификация на других задачах генерации графов
2. Теоретические предположения: Зависит от предположений о конкретном дизайне действий и архитектуре GNN
3. Приблизительные методы: Приблизительная коррекция при генерации фрагментов не имеет теоретических гарантий
4. Масштабируемость: Для очень больших графов расчет автоморфизмов может стать узким местом

Влияние

1. Академическая ценность: Предоставляет важное дополнение к теории GFlowNets, решая фундаментальную проблему
2. Практическая ценность: Имеет прямой вклад в приложения, такие как открытие лекарств
3. Воспроизводимость: Метод простой, легко воспроизводится и применяется
4. Вдохновляющее значение: Предоставляет идеи для обработки симметрии в других генеративных моделях

Сценарии применения

1. Молекулярный дизайн: Открытие лекарств, дизайн материалов и другие приложения химической информатики
2. Генерация графов: Задачи генерации графов, требующие учета структурной симметрии
3. Комбинаторная оптимизация: Задачи комбинаторной оптимизации с ограничениями симметрии
4. Обучение с подкреплением: Задачи RL с симметричным пространством состояний

Библиография

1. Bengio et al. (2021) — Основы GFlowNet
2. Ma et al. (2024) — Первое выявление проблемы эквивалентных действий
3. Malkin et al. (2022) — Целевая функция траекторного баланса
4. Jain et al. (2023) — Применение многоцелевых GFlowNets

Общая оценка: Это отличная статья, сочетающая теорию и практику, решающая важную, но упущенную фундаментальную проблему в GFlowNets. Метод элегантен и прост, теоретический анализ строг, экспериментальная верификация полна. Имеет значительный вклад как в теоретическое развитие GFlowNets, так и в практические приложения, и, как ожидается, будет иметь продолжительное влияние на соответствующие области.