Рисуем с помощью закона Кулона. Как сделать генеративную модель на основе электростатики

Привет, Хабр! Меня зовут Александр Колесов, я исследователь группы «Основы генеративного ИИ» AIRI. У себя в команде мы активно исследуем то, как устроена работа генеративных моделей, ищем новые методы, экспериментируем. Недавно мы обратили внимание на то, что те пути, которые проходят представления данных в диффузионных моделях, очень похожи на пучки силовых линий электрического поля. 

Это не только красивая метафора — мы предложили метод Electrostatic Field Matching (EFM), который позволять извлечь из такой аналогии пользу. Статью с подробным описанием мы недавно свозили на ICLR 2026, там все подробности, теоремы и эксперименты. Здесь же хотелось кратко пересказать основную идею и показать её реализацию на простых примерах.

Генеративные модели на основе физики и их интерпретируемость.

Про диффузионные модели для генерации данных слышали, пожалуй, все, даже за пределами машинного обучения. Если кратко: в них упорядоченные данные строятся за счёт пошагового расшумления хаоса — как будто растворённые в стакане воды чернила снова собираются в каплю, отсюда и название. 

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

Интерпретируемость стала одной из важнейших причин, по которой диффузионные модели стали так распространены на практике, в особенности для генерации изображений, либо превращения одних изображений в другие (генерация типа данные→данные). В тех же состязательных моделях GAN, которые также используют для генерации, не идёт никакой речи об интерпретируемости модели. Там генератор и дискриминатор играют в «кошки‑мышки», и понять, почему модель выдала конкретный результат, почти невозможно.

Тогда возникает логичный вопрос: раз диффузионные модели, основанные на термодинамике, дают такие хорошие преимущества, включающие интерпретируемость, так может нам стоит поискать ещё какие‑то процессы в физике, полезные для генеративного ИИ? 

На какие процессы стоит обратить внимание? 

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

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

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

В идеале нам нужно что‑то, что вообще не нужно интегрировать, а просто вычислять по готовым формулам. Несложно показать, что оба вышеупомянутых требования выполняются для электростатики. Закон Кулона позволяет строить силовые линии электрического поля, порождаемого точечными зарядами, в довольно простом виде, а за сохранение потока отвечает теорема Гаусса.

Краткий ликбез: закон Кулона и суперпозиция

Давайте ненадолго вернемся в 8 класс. 

Закон Кулона в трехмерном пространстве говорит: поле точечного заряда убывает как . Забегая вперёд, отмечу, что в нашем методе мы работаем в многомерном пространстве размерности , и нам нужно обобщение этой формулы. Хорошая новость в том, что попыткам сформулировать электростатику в многомерных пространствах уже более ста лет, и обобщённый закон Кулона давно известен: формула для поля точечного заряда становится чуть сложнее, но суть та же — убывание :

где  — площадь поверхности гиперсферы единичного радиуса. В трёхмерном пространстве () гиперсфера — это обычная сфера, и формула принимает привычный школьный вид. 

Принцип суперпозиции — главная суперсила электростатических моделей. Если рассмотрим произвольную точку пространства, то поле в ней будет равно сумме полей создаваемое всеми частицами. 

Силовые линии. Итак, мы применили принцип суперпозиции и закон Кулона и вычислили вектор поля в некоторой точке. Мы можем на чуть‑чуть сместиться в пространстве вдоль этого вектора, выбрать новую точку и повторить вычисления. Таким образом можно построить кривые, касательные к которым совпадают с направлениями векторов напряженности поля, а густота кривых определяет их модуль — это и есть силовые линии.

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

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

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

Суть метода и идея конденсатора

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

Откуда столько измерений? Это особенность представления данных в вычислительных алгоритмах, включая диффузионные модели. Поясню на примере.

Допустим, у вас есть датасет цветных картинок 32×32 пикселей по 3 цветовых канала на каждый. В таком случае цвет каждого пикселя выступает в качестве независимой величины, и их всех можно объединить в огромный вектор 3072-мерного пространства. 

Таким образом, весь датасет — это множество точек этого пространства, которое мы описываем соответствующим распределением. Мы идём дальше и предлагаем смотреть на это пространство, как на гиперплоскую обкладку конденсатора, где каждой точке ставится в соответствие точечный заряд — положительный для одного распределения (например, шума) и отрицательный для другого (например, картинки, из которой этот шум был получен). 

Что можно сказать, про то, что происходит внутри такого конденсатора?

Во‑первых, каждый заряд с левой пластины и каждый с правой в любой точке пространства создает свое поле. По принципу суперпозиции, поле в любой точке внутри гиперконденсатора — это сумма полей от всех зарядов.

Во‑вторых, все линии поля, которые исходят из положительных зарядов и проникают внутрь гиперконденсатора, упираются в отрицательные заряды — они не порождаются и не исчезают посередине. На более формальном языке речь идёт о сохранении потока. Мы можем убедиться в этом, взяв интеграл по полю через замкнутую поверхность, левый и правый край которой проходят параллельно обкладкам. Помните, сохранение потока было одним из требований, которые мы накладывали на процесс?

В‑третьих, линии электрического поля начинаются слева на положительной частице, а заканчиваются справа на отрицательной. Это означает, что перенос одного распределения в другое осуществляется строго вдоль этих силовых линий (в статье это гарантирует отдельная теорема). В каждой точке этих линий касательный вектор к ним — это в точности вектор поля.

Отсюда идея: если я в каждой точке знаю поле, то знаю касательный вектор к кривой, а значит знаю направлению, вдоль которого надо двигаться, чтобы от левой пластины (одного вероятностного распределения) прийти в другую пластину (другое распределение). Таким образом, задача об обучении генеративной модели может быть сведена к выучиванию поля в гиперконденсаторе.

Нейросеть вместо прямого вычисления 

А зачем обучать сеть, если поле всегда можно посчитать? Действительно, много ли нужно, чтобы просто взять формулу для закона Кулона и подставить в неё нужные значения?

Во‑первых, нейронная сеть в некотором смысле обобщает поле. При использовании других методов машинного обучения (полиномов или деревьев) при движении по полю мы всегда будем получать тренировочные данные на правой пластине. Поскольку мы в конечном итоге хотим генерировать новое, то нам необходимо всё обобщать, а, значит, нам нужно учить нейронки.

Во‑вторых, кажущаяся простота разбивается об объёмы вычислений. На практике нужно будет вычислять поле, порождаемое миллионами точек с обкладок. Вместо этого мы предлагаем обучить нейронную сеть на небольшом объеме данных и быстро аппроксимировать направление электрического поля.

А теперь ответим на вопрос, как учить, и что нам даёт обученное поле?

Что подаём на вход сети? Координаты в пространстве 

где — точка между распределениями на гиперплоскости, а — высота между пластинами (от до ).

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

Как получаем правильный ответ (target)? Подставляем заряды и в формулу суперпозиции 

и считаем поле через Монте — Карло сэмплирование.

В итоге сеть выучивает поле один раз, и на инференсе мы просто «запускаем» силовую линию из любой точки из и шаг за шагом строим её, вычисляя касательный вектор. Там, где она достигнет пластины , и будет наш ответ. 

Эксперименты

Метод EFM работает и для перевода из шума в данные (как в диффузии), и для перевода данные→данные (как CycleGAN, но без циклов). Мы убедились в этом с помощью серии простых экспериментов. 

Ниже представлен результат работы EFM с хорошо известными игрушечными распределениями, для которых  — так их можно наглядно визуализировать. Для перехода Gaussian → Swiss Roll видно, как линии поля плавно перетекают из колокола в спираль. 

На цветных MNIST метод хорошо переводит «3» в «2», не путая цвет и форму. Причём траектории получаются гладкими и детерминированными — никаких случайных семплов шума на каждом шагу.

С традиционной генерацией на MNIST метод EFM также справляется хорошо:

Что в итоге? 

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

Никто, однако, не говорит, что электростатика — это единственная физическая аналогия для переноса распределений. На самом деле, требований к процессу не так уж и много, и они описаны выше: сохранение потока и аналитическая простота. Сюда можно добавить, разве что, требование убывание потенциала на бесконечности — с причины подробно разобраны в статье. Если у вас есть идеи, какие другие физические процессы могли бы быть здесь полезны — велкам!

И ещё кое‑что. Кто‑то из вас может справедливо заметить, что у нас в статье генерация безусловная, но на практике почти все модели рисуют изображения по промпту. На самом деле, метод EFM хорошо работает и здесь — и про это у нас есть отдельная работа, но об этом уже как‑нибудь в другой раз. 

Спасибо, что прочитали!

Авторы статьи: Александр Колесов (Сколтех, AIRI), Манухов Степан (Сколтех, МГУ), Владимир Палюлин (Сколтех), Александр Коротин (Сколтех, AIRI).