Physics-Informed Machine Learning — одно из перспективных направлений в ML. Здесь ключевую роль играют физически-информированные нейронные сети (Physics-Informed Neural Networks, PINNs). Однако процесс их обучения сопряжён с рядом трудностей. В этой статье кратко обсудим, как можно решить эти проблемы и повысить эффективность обучения PINNs.

Классическая постановка задачи для PINNs

Physics-Informed Neural Networks (PINNs) предназначены для решения дифференциальных уравнений — как обыкновенных, так и в частных производных. Классические численные методы часто сталкиваются с проблемами, когда уравнение сложное, а его размерность велика — вычислительная стоимость становится слишком высокой.

Кроме того, эти методы, как правило, требуют использования фиксированной расчетной сетки. Однако в реальных задачах область определения может быть сложной геометрически и даже меняться со временем. PINNs свободны от этого ограничения: они являются mesh-free (безсеточными) и работают с произвольным набором точек в пространстве и времени.

Помимо самого уравнения, существуют граничные (условия Дирихле, Неймана) и начальные условия (условия Коши). Если эти условия сложны или их много, численное решение усложняется, а иногда требует нетривиальной аппроксимации. PINNs же позволяют естественным образом встроить эти условия прямо в функцию ошибки нейронной сети.

Как это работает и какие преимущества дает?

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

Важнейшее применение PINNs — решение обратных задач. Часто уравнения содержат неизвестные параметры. Если у исследователя есть набор экспериментальных или исторических данных, PINN может откалибровать параметры уравнения так, чтобы его решение наилучшим образом соответствовало этим данным.

Пример из финансов: Для уравнения Блэка-Шоулза в качестве данных могут выступать рыночные цены опционов для различных страйков. Калибровка PINN по этим данным позволяет оценить такие параметры, как волатильность и процентная ставка. Таким образом, модель одновременно учитывает физику процесса (уравнение) и подстраивается под реальные данные.

Кроме того, PINNs эффективны для интерполяции и экстраполяции решений. Подробнее о задачах можно почитать тут.

В чем сложности?

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

Итак, перейдем к обсуждению полезных методик.

1. Оптимизаторы

Классический алгоритм Adam хорошо зарекомендовал себя в задачах машинного обучения, в том числе для обучения физически-информированных нейронных сетей. С его помощью можно эффективно исследовать пространство параметров на начальном этапе. Однако на более поздних стадиях обучения предпочтительнее применять более точные методы оптимизации.

На сегодняшний день стандартом для обучения PINNs стала комбинированная стратегия. Она заключается в следующем: сначала пространство параметров исследуется с помощью адаптивных методов, таких как Adam, или, что ещё лучше, его модификации AdamW. В AdamW регуляризация (weight decay) отделена от обновления градиентов, что повышает стабильность процесса обучения.

На втором этапе используются более точные методы оптимизации, например, Limited-memory BFGS (L-BFGS). Этот квазиньютоновский метод позволяет улучшить сходимость модели, обеспечивая более точную настройку параметров, что особенно критично для решения обратных задач.

Что касается планировщиков (lr scheduler, шедулер), то в большинстве задач неплохо себя показывает cosine annealing schedule (в Pytorch называется CosineAnnealingLR). Он плавно меняет параметр скорости обучения и как понятно из названия, при помощи косинуса. Начинать обучения для Adam можно в районе 1e-3 и доходить, например, до 1e-5. Но здесь конкретные цифры зависят от задачи и лучше поиграться с этими гиперпараметрами.

Здесь также стоит упомянуть общие техники улучшения оптимизации, такие как warm restarts. Например, в PyTorch для этого реализован планировщик CosineAnnealingWarmRestarts.

Суть этой техники заключается в периодическом резком увеличении скорости обучения (learning rate) с последующим её плавным снижением в течение определенного количества шагов. Это позволяет оптимизатору "выпрыгнуть" из неудачного локального минимума, если он в него попал, и с некоторой вероятностью найти более качественный минимум в другой области пространства параметров.

Из недавнего также появилась работа. В ней авторы предлагают помимо комбинации Adam+L-BFGS добавить третий оптимизатор второго порядка под кодовым названием NNCG. Этот метод использует низкоуровневую аппроксимацию матрицы Гессе для построения эффективного предобуславливателя для метода сопряжённых градиентов.

2. Архитектура

Архитектура PINN, конечно, зависит от конкретной задачи. Если вы обучаете PINN для работы с изображениями, логичнее использовать сверточные слои. Если же цель — просто решить дифференциальное уравнение, то архитектура чаще всего представляет собой последовательность обычных линейных (полносвязных) слоев.

Что касается функции активации, то здесь предпочтительнее использовать гладкие функции, такие как гиперболический тангенс или сигмоида. Поскольку предполагается, что решение уравнения является гладкой функцией необходимого порядка, выбор функции активации может сыграть важную роль. А функции по типу ReLU могут наоборот замедлять сходимость сети.

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

Если решение обладает как высокочастотными, так и низкочастотными компонентами, можно использовать Фурье-слои или другие периодические слои, способные эффективно разделять частоты.

Также следует быть осторожнее с применением dropout. Исходя из предположения о гладкости решения, dropout может серьезно ухудшить сходимость. Вместо него лучше применять гладкие методы регуляризации, такие как L2.

При использовании достаточно глубокой сети, даже состоящей только из линейных слоев, не забывайте добавлять skip-connections (остаточные блоки). Это поможет стабилизировать обучение и предотвратит забывание важной информации.

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

Кроме того, в архитектуру можно изначально встраивать необходимые ограничения в виде жестких граничных условий (hard constraints). Например, если известно, что решение в определенной области должно быть равно нулю, не обязательно добавлять это условие в функцию ошибки — вместо этого можно напрямую спроектировать слои или преобразования в архитектуре, чтобы гарантировать выполнение требуемых условий.

3. Функция ошибки

Функция ошибки для PINNs обычно состоит из трех компонент: 1) ошибка данных; 2) ошибка, связанная с граничными условиями и 3) ошибка, связанная с уравнением.

Из-за такой комбинированной функции ошибки, PINNs склонны обучаться плохо. Об этом написано много в научной литературе. И часто ошибка связанная с уравнением на порядок больше чем ошибка данных и ошибка граничных условий. Поэтому следует правильно подбирать коэффициенты перед компонентами функции ошибки. Существуют даже адаптивные алгоритмы взвешивания функции ошибки. Подробнее тут и тут.

4. Область определения

Если поведение решения значительно отличается в различных подмножествах области определения, то тогда может быть целесообразней распараллелить PINN и сделать отдельные ветки нейронной сети для различных областей. При этом чтобы сохранить гладкость решения, можно использовать адаптеры или, например, линейную модуляцию (FiLM).

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

Что касается точек в которых вычисляется PINN, то здесь рекомендую использовать не просто равномерное распределение по всей сетки переменных, а лучше использовать сэмплирование с приоритетом для точек, которые вызывают наибольшие проблемы (residual-based sampling). И аналогично для вычисления ошибки в граничных точках, лучше использовать такое сэмплирование, которое приоритетом будет выбирать именно граничные точки (importance sampling) или напрямую задать их.

Можно также воспользоваться последовательностями с низким расхождением (еще называются квазислучайные). Это последовательности генерируется по заданному детерминированному правилу, при этом они эффективно и равномерно покрывают сетку области определения, обладая некоторыми свойствами случайных последовательностей. Такой подход особенно популярен в финансах (отдельно можно почитать о методе квази-Монте-Карло).

Помимо этого существует адаптивные методы генерации точек на основе производных второго порядка (матрицы Гессе).

В данной области простор велик и можно пробовать разные методы сэмплирования. Наиболее эффективны квази методы, importance sampling и residual-based sampling.

Итог

Надеюсь данный мини-обзор методов улучшения сходимости PINNs будет вам полезен на начальном этапе построения моделей. Как часто бывает, все задачи уникальны и требуют особого подхода, но всегда можно начать с общих принципов. Обо всем этом пишу у себя в тг @kirill_zakharov_blog.

Как улучшить сходимость в физически-информированных нейронных сетях