Внутренности PyTorch: как нейросеть считает градиенты для физических задач

Исследователи представили детальный разбор механизма автоматического дифференцирования (AD) в PyTorch применительно к физически-информированным нейронным сетям (PINN). Работа, опубликованная на arXiv, впервые шаг за шагом с численными значениями демонстрирует, как фреймворк вычисляет градиенты двух уровней: производную от выхода сети по времени и градиенты параметров по функции потерь.

PINN — это особый класс нейросетей, которые решают дифференциальные уравнения, встраивая физические законы прямо в процесс обучения. Для этого требуется дважды дифференцировать: сначала получить производную решения (например, скорость изменения температуры), а затем — градиенты весов сети по составной функции потерь. PyTorch справляется с этим с помощью встроенного движка autograd.

Авторы взяли простейшую задачу: уравнение y' + y = 0, y(0)=1, и полносвязную сеть с архитектурой 1-3-3-1 (один вход, два скрытых слоя по три нейрона, один выход). Они проследили каждый узел вычислительного графа, построенного при прямом проходе, и показали, как обратный проход (reverse-mode) за один проход вычисляет все 22 градиента параметров.

Особое внимание уделено механизму create_graph=True, который позволяет дифференцировать через физический остаток (physics-informed residual). Этот флаг заставляет PyTorch строить граф для градиентов, что необходимо для вычисления вторых производных. Все значения согласованы с аналитическими выкладками из предыдущей работы Tahimi (2026), что подтверждает корректность численных результатов.

Практическая ценность работы в том, что она даёт чёткое понимание внутренней работы автограда PyTorch в нетривиальном сценарии. Для разработчиков PINN это может помочь отлаживать модели, оптимизировать память и избегать ошибок при использовании высоких производных. Кроме того, показана связь между формализмом чувствительности P/Q и векторно-якобиановыми произведениями, используемыми в PyTorch.

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