Ученые ускорили генеративные модели, встроив физические законы через GPU-оптимизацию

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

В новой работе на arXiv авторы предлагают метод SNAP-FM (Sparse Nonlinear Accelerated Projection for Physics-Constrained Generative Modeling). Они задействуют структуру, которую физические ограничения естественным образом порождают в подзадачах проекции: блочно-разреженные системы Якоби и ККТ. Эта структура раскрывается с помощью пакета ExaModels.jl, а решение разреженных нелинейных задач выполняется через MadNLP.jl с использованием GPU-разреженной факторизации.

Такой подход позволил существенно ускорить проекцию нелинейных ограничений, сохраняя при этом их точное выполнение. Эксперименты проводились на эталонных задачах PDE с линейными, нелинейными, одномерными и двумерными ограничениями в рамках Physics-Constrained Flow Matching (PCFM).

Результаты показывают, что разреженная нелинейная оптимизация на GPU может стать практической основой для генеративной выборки с ограничениями в научном машинном обучении. В отличие от плотных тензорных операций в стандартных ML-фреймворках, новый метод эффективно использует разреженность, обусловленную локальными связями PDE и пакетной обработкой.

Разработка открывает путь к более быстрым и надёжным генеративным моделям, которые можно применять в аэродинамике, материаловедении, климатическом моделировании и других областях, где физическая согласованность критична.