Руководство по созданию вашего первого GFlowNet в TensorFlow 2
Генеративные потоковые сети (GFlowNets) — это метод машинного обучения для создания составных объектов с частотой, пропорциональной связанному с ними вознаграждению.
В этой статье мы раскроем значение всех этих слов, объясним, чем полезны GFlowNet, поговорим о том, как их обучать, а затем разберем реализацию TensorFlow 2.
Развивайте свою интуицию
GFlowNets были представлены на NeurIPS в 2021 году Эммануэлем Бенджио и соавторами. GFlowNets — это метод глубокого обучения для создания объектов с частотой, пропорциональной ожидаемому вознаграждению за эти объекты в среде.
Мотивирующим примером в их первой статье является открытие новых химических структур (или я также буду называть их «молекулами»). Химические структуры или молекулы являются «композиционными» в том смысле, что они состоят из отдельных строительных блоков (атомов). Химическая структура может быть проверена на интересующее качество, такое как антибиотическая активность. Некоторые молекулы являются сильными антибиотиками — при тестировании против бактерий в лаборатории большинство бактерий погибает, и получается измерение, о котором мы можем думать как о большой «награде». Большинство молекул не убивают бактерии и, таким образом, возвращают небольшое «вознаграждение» в результате лабораторного теста.
GFlowNet — это вероятностная модель, которая строит объекты. Как правило, GFlowNet используется следующим образом:
- С начальной точки он выводит распределение вероятностей по возможным строительным блокам.
- Он выбирает один стандартный блок случайным образом в соответствии с предсказанным распределением. Обратите внимание, что этот шаг является стохастическим. Какие-то действия будут более вероятными, какие-то менее.
- Добавив этот первый строительный блок, он затем принимает построенный объект в этой точке в качестве входных данных, генерирует новое распределение вероятностей по возможным строительным блокам и добавляет следующий.
- Так происходит до тех пор, пока не будет выбрано завершающее действие.
- Действие завершения завершает сгенерированный объект, который затем можно запросить в среде, чтобы получить соответствующее значение вознаграждения.
Некоторые ключевые вещи, которые нужно понять:
- Построение рассматривается как марковский процесс. То есть при выборе следующего строительного блока единственным вводом является текущее состояние строящегося объекта.
- Мы использовали образ «строительных блоков», но «действия», вероятно, более подходящий термин. При построении молекулы вы не просто выбираете атом углерода, вы также выбираете, куда его поместить. И представление о каждом шаге как о выборе действия упрощает представление о других применениях GFlowNet. В приведенном ниже примере кода мы строим путь, перемещаясь по сетке. Каждое действие — это шаг в соседнюю ячейку.
- Распределение вероятностей по возможным действиям изучается (мы скоро обсудим, как это делается). Это прогнозируемое распределение называется «политикой».
- Наиболее интуитивно понятным элементом модели является «форвардная политика». Это конструктивная часть, которая выводит распределение вероятности следующих действий. Модель также включает «обратную политику», описывающую распределение по действиям, которые могут привести к текущему состоянию.
- Полная последовательность действий от начала до конца называется «траекторией».
Теперь вы знаете, что GFlowNet — это вероятностные модели, которые создают объекты. Ключевой особенностью GFlowNets является то, что они создают объекты с частотой, пропорциональной вознаграждению.
Если вы повторите этот стохастический конструктивный процесс много раз, вы получите больше объектов в регионах с высоким вознаграждением и меньше объектов в регионах с низким вознаграждением. Используя наш пример с молекулами антибиотиков, модель будет генерировать больше молекул, похожих на сильные антибиотики, но также будет генерировать несколько молекул, похожих на более слабые. Выходные вероятности будут пропорциональны наградам.
Эммануэль Бенжио написал пост в блоге и сопровождающий его учебник, который стоит изучить.
Приложения GFlowNet
Теперь вы понимаете, что делает GFlowNet. Почему вы хотите один? Для чего они используются?
Они используются для разумной выборки из действительно очень больших сред.
Например, придерживаясь примера с химической структурой, предположим, что у вас есть некоторые обучающие данные для начала (группа химических структур и связанная с ними антибиотическая активность). Сначала вы обучаете модель (отдельно от вашей GFlowNet), которая может предсказывать активность антибиотиков с учетом химической структуры. Назовем ее суррогатной функцией вознаграждения, потому что на данный момент это лучшее представление взаимосвязи между химической структурой и «вознаграждением» антибиотика. Теперь вы хотите использовать свою суррогатную модель, чтобы выбрать новые, непроверенные химические структуры, чтобы опробовать их в лаборатории.
Итак, каков ваш следующий шаг? Одним из способов продолжения было бы исчерпывающее создание всех возможных химических структур, используя данные обучения в качестве отправной точки. Сложность в том, что пространство растет в геометрической прогрессии. Даже начиная с 10 или 20 возможных «строительных блоков», у вас также есть несколько возможных мест размещения на каждой молекуле, что приводит к сотням возможных действий. Если вы исчерпывающе перечислите каждую комбинацию таких действий, всего за несколько шагов, у вас будут миллиарды возможностей, и вы скоро перегрузите свои вычислительные ресурсы, даже не начав исследовать возможное химическое пространство.
Другим способом продолжения будет случайный выбор действий, которые позволят вам исследовать более глубоко, но не исчерпывающе. Проблема с этим подходом заключается в равномерной случайности. Поскольку большинство молекул не являются антибиотиками, вполне вероятно, что подавляющее большинство структур, отобранных для анализа, не принесут пользы.
Вот где GFlowNets вступают в игру. Сначала вы обучаете GFlowNet отражать функцию вознаграждения, изученную суррогатной моделью. Затем вы создаете множество структур с помощью GFlowNet и устанавливаете для них приоритеты с помощью суррогатной модели. Таким образом, вы берете пробы из возможного химического пространства, но таким образом, который обогащает многообещающие кандидаты, а также позволяет исследовать менее перспективные структуры.
Другими словами, сети GFlowNet используются для интеллектуальной выборки из действительно очень больших сред.
Как обучаются GFlowNet
Теперь вы знаете, что такое GFlowNet и как они используются. Как научить модель демонстрировать желаемое поведение? Как обучить модель, которая будет генерировать объекты, пропорциональные вознаграждению?
Для более подробного обсуждения этой темы я отсылаю вас к этой статье, в которой обсуждаются различные целевые функции, которые можно использовать для обучения GFlowNet. Для целей этой статьи мы будем придерживаться баланса траектории.
Теперь мы переходим к части «Flow» GFlowNets. Авторы предлагают нам думать обо всех возможных траекториях в окружающей среде как о сети труб, по которым течет вода. Каждая «труба» — это действие из одного состояния в другое. Вода течет из источника (начальной точки для всех траекторий), через все действия и вытекает в конечных состояниях, где известны «награды». Цель состоит в том, чтобы распределить воду по трубам так, чтобы она вытекала в объемах, пропорциональных наградам.
Функция потери баланса траектории является одним из способов достижения этой цели. Для каждой траектории прямой поток вероятностей (из исходного состояния в конечное) должен быть пропорционален потоку вознаграждения из конечного состояния в исходное. То есть произведение вероятностей прямой политики на траектории должно быть пропорционально произведению вероятностей обратной политики и вознаграждения на той же траектории. Записанная как функция потерь, это выглядит так:
Где:
- Z — это одно изученное значение, которое представляет собой общий поток от источника через всю сеть.
- P_F — это политика прямой передачи, а P_B — политика обратной связи. Обратите внимание, что для каждой траектории вы берете произведение вероятностей каждого действия вдоль траектории.
- R — вознаграждение, связанное с определенной траекторией. То есть это награда за объект, сгенерированный траекторией.
- Θ представляет параметры модели, которые будут изменены, чтобы минимизировать эту функцию.
- n — длина траектории, для которой рассчитывается значение потерь.
Интуитивно внутренняя часть представляет собой отношение прямого потока (от источника к завершению) к обратному потоку (от вознаграждения к источнику). В идеале отношение будет стремиться к единице. Взятие логарифма этого отношения означает, что если прямой поток больше, логарифм будет больше нуля. Если обратный поток больше, журнал будет меньше нуля. Если взять все это в квадрат, то минимум будет равен нулю (то есть отношение внутри логарифма будет равно единице).
Еще одним преимуществом этой схемы является то, что мы избегаем проблем с недополнением при умножении многих вероятностей вместе. Умножения превращаются в суммы под логарифмом.
Эта функция потери баланса траектории вычисляется для каждой выборки в обучающей выборке. Ключевым моментом является то, что вознаграждение и конечное состояние являются фиксированными значениями в ваших обучающих данных, но траектории изучаются. В каждую эпоху обучения прямая и обратная политика будут меняться, и модель будет основываться на последовательностях действий, которые уравновешивают вероятности, чтобы они должным образом отражали вознаграждение.
Создайте GFlowNet в TensorFlow 2
Теперь, когда вы понимаете идеи, лежащие в основе GFlowNet, почему они полезны и как их обучать, вы готовы создать их. Весь свой код авторы предоставляют на github, и вам обязательно нужно его изучить. В то время как лаборатория Bengio реализовала свои GFlowNet с использованием PyTorch, мы будем использовать TensorFlow 2.
Весь код для этого руководства можно найти по адресу: https://github.com/mbi2gs/gflownet_tf2/
Окружающая среда
Во-первых, среда, описанная в оригинальной публикации GFlowNet, представляет собой n-мерный куб со сторонами длиной H и функцией вознаграждения, которая является самой высокой в углах. В этом примере мы будем использовать двумерную сетку 8x8, потому что ее легко построить и увидеть, что происходит. Файл env.py определяет класс RewardEnvironment. Сердцем среды вознаграждения является функция, которая вычисляет вознаграждение на основе входных координат:
def reward_fn(self, coord): """Calculate reward function according to section 4.1 in https://arxiv.org/abs/2106.04399. :param x_coord: (array-like) Coordinates :return: (float) Reward value at coordinates """ assert len(coord) >= 1 assert len(coord) == self.dim r1_term = 1 r2_term = 1 reward = self.r0 for i in range(len(coord)): r1_term *= int(0.25<np.abs(coord[i]/(self.length-1)-0.5)) r2_term *= int( 0.3 < np.abs(coord[i]/(self.length-1)-0.5)<= 0.4) reward += self.r1*r1_term + self.r2*r2_term return reward
Итоговая среда выглядит следующим образом: более низкая награда отмечена фиолетовым цветом, а максимальная награда — желтым.
Выборка вперед и назад
Класс GFNAgent определен в gfn.py. Выборка из форвардной политики влечет за собой:
- Начиная с исходной точки
(0,0). - Создание распределения вероятностей возможных действий (вверх, вправо или завершение) с использованием модели.
- Выбор действия по вероятностям.
- Применение действия для прекращения или перемещения в новую позицию.
- Продолжайте до завершения действия.
Понимание процесса выборки — лучший способ понять, что на самом деле делает GFlowNet. В питоне вот функция прямой выборки:
def sample_trajectories(self, batch_size=3, explore=False):
"""Sample `batch_size` trajectories using the current policy.
:param batch_size: (int) Number of trajectories to sample
:param explore: (bool) If True, mix policy with uniform
distribution to encourage exploration.
:return: (tuple of nd.array) (trajectories, one_hot_actions,
rewards)
"""
# Start at the origin
still_sampling = [True]*batch_size
positions = np.zeros((batch_size, self.dim), dtype='int32')
trajectories = [positions.copy()]
one_hot_actions = []
batch_rewards = np.zeros((batch_size,))
for step in range(self.max_trajectory_len-1):
# Convert positions to one-hot encoding
one_hot_position = tf.one_hot(
positions, self.env_len, axis=-1)
# Use forward policy to get log probabilities over actions
model_fwrd_logits = self.model.predict(one_hot_position)[0]
model_fwrd_probs = tf.math.exp(model_fwrd_logits)
if explore:
# Mix with uniform distribution to encourage exploration
unif = self.unif.sample(
sample_shape=model_fwrd_probs.shape[0])
model_fwrd_probs = self.gamma*unif + \
(1-self.gamma)*model_fwrd_probs
# Don’t select impossible actions (like moving out of
# the environment)
normalized_actions = self.mask_and_norm_forward_actions(
positions, model_fwrd_probs)
# Select actions randomly, proportionally to input
# probabilities
actions = tfd.Categorical(probs=normalized_actions).sample()
actions_one_hot = tf.one_hot(
actions, self.action_space).numpy()
# Update positions based on selected actions
for i, act_i in enumerate(actions):
if act_i == (self.action_space — 1) and still_sampling[i]:
still_sampling[i] = False
batch_rewards[i] = self.env.get_reward(positions[i,:])
elif not still_sampling[i]:
positions[i,:] = -1
actions_one_hot[i,:] = 0
else:
positions[i, act_i] += 1
trajectories.append(positions.copy())
one_hot_actions.append(actions_one_hot)
return np.stack(trajectories, axis=1), \
np.stack(one_hot_actions, axis=1), \
batch_rewards
Некоторые детали, которые следует выделить:
- Координаты положения преобразуются в горячие векторы перед подачей в модель.
- Существует вариант «исследовать», который смешивает единообразное распределение с типовой политикой, чтобы поощрить решения, не соответствующие политике. Гамма — это параметр, контролирующий степень разведки (путем управления соотношением смешивания полиса и униформы).
- Некоторые действия невозможны, но модель сначала не знает об этом (например, действия, выводящие модель из среды). Мы выполняем шаг маскировки, чтобы свести вероятность невозможных действий к нулю. В то же время позже вы увидите, что мы также обучаем модель распознавать эти невозможные шаги.
- Приведенная выше функция отбирает несколько траекторий одновременно. Для каждой завершенной траектории добавляется символ заполнения (
-1), в то время как оставшаяся часть завершает выборку.
Выборка обратных траекторий очень похожа (см. функцию back_sample_trajectory()), но использует обратную политику, а не прямую. Отсталая политика может показаться ненужной, но она является ключом к созданию работающей модели. В ходе обучения каждый раз, когда в модель передается пара (позиция, вознаграждение), она производит выборку из обратной политики, передавая вознаграждение обратно в источник, а затем сопоставляет прямую политику с этой. траектория с обратной выборкой, действие за действием. Мы увидим, как это работает в коде, в следующем разделе, но просто знайте, что обратная политика необходима!
Тренировочный цикл
Процесс обучения начинается с выборки траектории и определения вознаграждения путем опроса окружающей среды. Можно итеративно сэмплировать, затем обучать, затем сэмплировать, затем тренировать. С каждой итерацией модель все лучше аппроксимирует окружающую среду, и если есть какая-то структура вознаграждения, которую можно изучить с помощью GFlowNet, она ускорит процесс исследования за счет интеллектуального отбора проб из неизвестных областей. Также можно начать с обучающего набора данных, собранного независимо от политики GFlowNet (это называется «автономным» обучением). В этом уроке мы используем офлайн-обучение.
Процесс обучения начинается с перемешивания и группирования обучающих данных. В этой демонстрации мы используем функцию генератора train_gen() для создания перемешанных пакетов из 10 пар (позиция, вознаграждение), которые перебирают весь набор данных каждую эпоху.
Функция train() выполняет итерацию через заданное количество эпох. В каждой эпохе он вычисляет значение потерь для каждого обучающего примера, а затем обновляет градиент для весов модели. Если средние потери в течение эпохи ниже, чем в любую предыдущую, веса модели сохраняются. Таким образом, лучший набор весов может быть загружен в конце.
def train(self, verbose=True):
"""Run a training loop of `length self.epochs`.
At the end of each epoch, save weights if loss is better
than any previous epoch.
At the end of training, read in the best weights.
:param verbose: (bool) Print additional messages while training
:return: (None) Updated model parameters
"""
if verbose:
print('Start training...')
# Keep track of loss during training
train_loss_results = []
best_epoch_loss = 10**10
model_weights_path = './checkpoints/gfn_checkpoint'
for epoch in range(self.epochs):
epoch_loss_avg = tf.keras.metrics.Mean()
sampler = self.train_gen()
# Iterate through shuffled batches of deduplicated data
for batch in sampler:
loss_values, gradients = self.grad(batch)
self.optimizer.apply_gradients(
zip(gradients, self.model.trainable_variables + [self.z0])
)
losses = []
for sample in loss_values:
losses.append(sample.numpy())
epoch_loss_avg(np.mean(losses))
# If current loss is better than any previous, save weights
if epoch_loss_avg.result() < best_epoch_loss:
self.model.save_weights(model_weights_path)
best_epoch_loss = epoch_loss_avg.result()
train_loss_results.append(epoch_loss_avg.result())
if verbose and epoch % 9 == 0:
print(f'Epoch: {epoch} Loss: {epoch_loss_avg.result()}')
# Load best weights
self.model.load_weights(model_weights_path)
Градиенты рассчитываются с помощью функции град(). Эта функция использует объект tf.GradientTape() для расчета градиентов для каждого параметра модели с учетом потерь. Обратите внимание, что параметр z0 добавлен к списку обновляемых параметров.
Функция потерь
В основе тренировочного процесса лежит функция потери баланса траектории. Напомним, что есть и другие функции потерь, которые можно использовать, но баланс траектории имеет тенденцию превосходить предыдущие методы (см. сноску 3).
def trajectory_balance_loss(self, batch): """Calculate Trajectory Balance Loss function as described in https://arxiv.org/abs/2201.13259. I added an additional piece to the loss function to penalize actions that would extend outside the environment. :param batch: (tuple of ndarrays) Output from self.train_gen() (positions, rewards) :return: (list) Loss function as tensor for each value in batch """ positions, rewards = batch losses = [] for i, position in enumerate(positions): reward = rewards[i] # Sample a trajectory for the given position using # backward policy trajectory, back_actions = self.back_sample_trajectory( position) # Generate policy predictions for each position in trajectory tf_traj = tf.convert_to_tensor( trajectory[:,…], dtype=’float32') forward_policy, back_policy = self.model(tf_traj) # Use “back_actions” to select corresponding forward # probabilities forward_probs = tf.reduce_sum( tf.multiply(forward_policy, back_actions), axis=1) # Get backward probabilities for the sampled trajectory # (ignore origin) backward_probs = tf.reduce_sum( tf.multiply(back_policy[1:,:], back_actions[:-1,:self.dim]), axis=1) # Add a constant backward probability for transitioning # from the termination state backward_probs = tf.concat([backward_probs, [0]], axis=0) # take log of product of probabilities (i.e. sum of log # probabilities) sum_forward = tf.reduce_sum(forward_probs) sum_backward = tf.reduce_sum(backward_probs) # Calculate trajectory balance loss function and add to # batch loss numerator = self.z0 + sum_forward denominator = tf.math.log(reward) + sum_backward tb_loss = tf.math.pow(numerator — denominator, 2)
Убыток применяется к каждой паре (позиция, вознаграждение) отдельно. Обратите внимание, что первым шагом для заданной позиции обучения является использование обратной политики для выборки траектории обратно к исходной точке. Затем рассчитываются соответствующие форвардные вероятности для каждой позиции на выбранной траектории. Отношение z0 умноженных на прямые вероятности к вознаграждению, умноженному на обратные вероятности, формирует окончательное значение потерь, но помните, что мы применяем логарифм, чтобы умножения превращались в сложения.
Обратите внимание, что обратная политика в этой демонстрации не включает вероятность завершения. Вместо этого первое обратное действие всегда считается завершающим, поэтому мы просто включаем постоянную вероятность на первом обратном шаге. Оттуда и обратно до начала мы отбираем только вероятности перехода. Обратите внимание, что вероятность перехода из исходной точки не определена и не включена в обратную траекторию.
Наконец, мы добавляем некоторые дополнительные условия потери, которые наказывают за невозможные действия, такие как выход из среды. Мы делаем это, определяя позиции на краю среды, суммируя вероятности, ведущие от этого края за пределы среды, и добавляем эти ошибочные суммы к потерям. Мы делаем это как для прямой, так и для обратной политики.
# Penalize any probabilities that extend beyond the
# environment
fwrd_edges = tf.cast(
np.argmax(trajectory, axis=2) == (self.env_len-1),
dtype=’float32')
back_edges = tf.cast(
np.argmax(trajectory, axis=2) == 0,
dtype=’float32')
fedge_probs = tf.math.multiply(
tf.math.exp(forward_policy[:,:self.dim]),
fwrd_edges)
bedge_probs = tf.math.multiply(
tf.math.exp(back_policy[:,:self.dim]),
back_edges)[1:,:] # Ignore backward policy for the origin
fedge_loss = tf.reduce_sum(fedge_probs)
bedge_loss = tf.reduce_sum(bedge_probs)
combined_loss = tf.math.add(
tb_loss, tf.math.add(fedge_loss, bedge_loss))
losses.append(combined_loss)
return losses
Демонстрация
Блокнот gflownet_demo.ipynb объединяет все части в рабочий процесс обучения в автономном режиме.
Мы начинаем с инициализации агента GFlowNet (который также инициализирует среду вознаграждения).
from gfn import GFNAgent agent = GFNAgent()
Вызов agent.model.summary() показывает нам структуру модели, которая по умолчанию состоит из двух плотных слоев по 32 единицы в каждом, плотной прямой политики и плотной обратной политики. Есть 1765 обучаемых параметров плюс еще один для z0.
Необученная политика сразу после инициализации выглядит довольно однообразно. Вероятность вертикального или поперечного перехода (стрелки) или завершения (красные восьмиугольники) практически одинакова в каждой точке.
agent.plot_policy_2d()
Мы можем визуализировать результат этой необученной политики, выбрав множество траекторий и оценив вероятность завершения в каждой точке среды.
l1_error_before = agent.compare_env_to_model_policy()
Обратите внимание, что большинство траекторий заканчиваются вблизи начала координат. Этого следует ожидать, учитывая высокую вероятность завершения на каждой позиции. Цель обучения — сделать так, чтобы этот сюжет больше походил на ландшафт с наградами (режимы в каждом углу). Помните, правильно обученный GFlowNet генерирует объекты с вероятностью, пропорциональной вознаграждению.
Для целей этой демонстрации мы глубоко выбираем из среды, пока большинство позиций и связанное с ними вознаграждение не окажутся в наборе данных. Обратите внимание, что обучающие данные не дублируются, поэтому представлены только уникальные пары (позиция, вознаграждение). В этом случае после выборки 5000 траекторий мы получили все четыре моды (максимумы в каждом углу) и почти все 64 (8x8) позиции. Выборочные данные показаны на рисунке ниже, где размер x пропорционален вознаграждению в этот момент.
agent.sample(5000) agent.plot_sampled_data_2d()
А теперь обучаем модель!
agent.train()
Обучение продолжается 200 эпох, и перед остановкой потери уменьшаются примерно на два порядка.
Start training… Epoch: 0 Loss: 5.823361873626709 Epoch: 9 Loss: 1.9640847444534302 Epoch: 18 Loss: 1.6309585571289062 Epoch: 27 Loss: 1.0532681941986084 … … … … Epoch: 171 Loss: 0.03023693338036537 Epoch: 180 Loss: 0.050495393574237823 Epoch: 189 Loss: 0.046723444014787674 Epoch: 198 Loss: 0.06047442555427551
После обучения мы можем снова визуализировать политику, которая выглядит намного лучше.
agent.plot_policy_2d()
И, наконец, мы визуализируем вероятность попадания в каждую позицию окружения, выбирая 2000 траекторий, используя обученную форвардную политику.
l1_error_after = agent.compare_env_to_model_policy()
Это больше похоже на то, что они вознаграждают ландшафт! Конечно, это не идеально, и вы можете потратить много времени на настройку гиперпараметров, чтобы усовершенствовать модель. Но это, безусловно, обучение разумной политике.
Мы также можем сравнить изученное распределение вероятностей (оцененное путем рисования множества траекторий в соответствии с изученной политикой) и целевое (нормализация среды вознаграждения), используя ошибку L1 (абсолютное значение различий в каждой позиции).
До обучения ошибка L1 составляла 1,5, а после обучения уменьшилась до 0,2.
Заключение
Теперь вы знаете, что такое GFlowNet, чем они полезны и как их обучать. Вы прошли демонстрацию и можете изучить реализацию в TensorFlow 2. Вы готовы создать свою собственную и применить ее в своей работе. Удачи в изучении отличных политик!
