gfnx

Библиотека gfnx — быстрая и масшатбируемая JAX-платформа для обучения и оценки Generative Flow Networks. В отличие от существующих решений вроде torchgfn, где логика среды работает на процессоре и постоянно обменивается данными с видеокартой, gfnx переносит и среды, и обучение целиком на GPU. Вся вычислительная цепочка компилируется целиком с помощью JIT, что резко уменьшает накладные расходы и упрощает крупные серии экспериментов. Библиотека включает набор типовых сред из литературы по GFlowNet: синтетические гиперрешётки, генерацию двоичных последовательностей, задачи проектирования ДНК-последовательностей (TFBind8), генерацию молекул (QM9), синтез антимикробных пептидов (AMP), построение филогенетических деревьев, байесовское восстановление структуры графов и выборки из модели Изинга. Эксперименты показывают серьёзный выигрыш по времени работы при том же качестве выборок. На задачах последовательностей и молекул gfnx ускоряет обучение по сравнению с PyTorch-реализациями в десятки раз. В гиперрешётках и биологических задачах библиотека достигает тех же метрик расхождения с истинным распределением, что и базовые реализации, но при значительно большем числе итераций в секунду. Работа может быть полезна исследователям GFlowNet и обучения с подкреплением, а также тем, кто занимается амортизированным семплированием в комбинаторных и научных задачах.

статья | код