NetKet et JAX : comment construire un modèle Transformer pour les systèmes de spins frustrés

Трансформерные архитектуры начинают занимать место и в вычислительной квантовой физике: в новом практическом руководстве показано, как на базе NetKet и JAX собрать полноценный пайплайн Neural Quantum States для сложной задачи — фрустрированной спиновой цепочки Гейзенберга J1-J2. Речь не о теории в вакууме, а о воспроизводимой схеме, где модель, сэмплер, оптимизация и проверка точности собраны в один исследовательский контур. Главная идея гайда в том, что архитектура Transformer хорошо подходит для описания многочастичных квантовых состояний, где важны дальнодействующие корреляции между частицами.

Обычные численные методы быстро упираются в размерность пространства состояний, особенно если система фрустрирована, то есть конкурирующие взаимодействия мешают ей прийти к простому упорядоченному минимуму энергии. В таких условиях Neural Quantum States позволяют представить волновую функцию как параметризуемую нейросеть, а затем оптимизировать её через Variational Monte Carlo. NetKet здесь выступает как готовая среда для квантовых вычислений, а JAX — как движок для высокоточной и ускоряемой оптимизации.

В руководстве сначала настраивается базовая физическая часть задачи. Автор задаёт одномерную цепочку длины L с периодическими граничными условиями, где ближайшие соседи взаимодействуют с коэффициентом J1, а следующие за ними — с J2. Именно это сочетание и создаёт фрустрацию, из-за которой задача становится интересной и нетривиальной.

Для описания системы используется граф в NetKet, гильбертово пространство спинов 1/2 с фиксированной полной проекцией и оператор гамильтониана, собранный через GraphOperator. Параллельно включается 64-битная точность JAX, что важно для устойчивых расчётов в задачах такого класса. Дальше начинается машинное обучение в чистом виде.

Волновая функция задаётся собственной моделью TransformerLogPsi на Flax: спиновые конфигурации кодируются как токены, затем получают эмбеддинги и позиционные представления, после чего проходят через несколько блоков self-attention и feed-forward слоёв. В примере используются размерность скрытого пространства 96, четыре головы внимания и шесть слоёв трансформера. На выходе модель возвращает комплексный логарифм амплитуды волновой функции — это критично, потому что квантовое состояние нельзя адекватно описать только вещественным скаляром.

После агрегации информации по всей цепочке через усреднение сеть получает глобальное представление конфигурации и может выражать более сложные корреляции, чем локальные анзацы. Отдельная ценность гайда в том, что он не останавливается на определении модели. Для обучения автор собирает полный VMC-контур: сэмплер MetropolisExchange, вариационное состояние MCState, оптимизатор Adam и Stochastic Reconfiguration как аналог natural gradient descent для квантовых состояний.

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

После обучения пайплайн используют уже как исследовательский инструмент. Автор прогоняет расчёты по нескольким значениям J2 в диапазоне от 0 до 0,7 для цепочки длиной 24 узла, фиксирует итоговые энергии и оценивает пик структурного фактора. Это позволяет не просто подобрать параметры нейросети, а посмотреть, как меняется физическое поведение системы при усилении фрустрации и где могут проявляться переходы между различными фазами магнитного порядка.

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

Для ML-инженеров это хороший пример того, как трансформер можно использовать вне текста, изображений и стандартных табличных данных. Для физиков — понятный шаблон, как перейти от абстрактной идеи Neural Quantum States к воспроизводимому эксперименту с конкретными метриками, бенчмарками и наблюдаемыми величинами. А для тех, кто работает на стыке этих областей, гайд даёт основу, которую можно расширять дальше: переходить к более крупным решёткам, добавлять симметрии, изучать запутанность или строить более сложные временные симуляции.

Что это значит: Transformer-подход постепенно становится рабочим инструментом не только для классических AI-задач, но и для моделирования квантовых систем, где цена ошибки высока, а точные методы быстро заканчиваются. Если NetKet и JAX уже есть в рабочем стеке, этот материал даёт практически готовую отправную точку для экспериментов исследовательского уровня.