PaperlayerView exact revision on arXiv

Source-linked AI summary

NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems

Filippo Vicentini, Damian Hofmann, Attila Szabó, Dian Wu, Christopher Roth, Clemens Giuliani, Gabriel Pescia, Jannes Nys, Vladimir Vargas-Calderon, Nikita Astrakhantsev, Giuseppe Carleo

arXiv:2112.10526v2quant-phcs.LGcs.MSphysics.comp-ph

TL;DR

Many-body quantum simulations need machine-learning tools that handle quantum objects without forcing researchers into specialized framework engineering or inefficient from-scratch implementations. NetKet 3 addresses this by combining modular quantum primitives with JAX-based, customizable neural-network states and algorithms; it provides broad support for dynamics, symmetries, large models, and extensible workflows, while retaining limitations in regularization, continuous-group encodings, and some platform or feature support.

  • Problem

    Standard machine-learning frameworks are poorly suited to non-standard quantum problems, while from-scratch implementations require expertise and can be sub-optimal.

  • Method

    NetKet 3 is a modular Python toolbox built on JAX that exposes quantum primitives, customizable neural-network ansätze, automatic differentiation, compiled execution, and extensible algorithms.

  • Results

    NetKet 3 adds support for dynamics, continuous and composite systems, million-parameter QGT and stochastic-reconfiguration calculations, arbitrary discrete symmetries, and modular reuse of individual components.

  • Takeaways & Limitations

    The toolbox provides accessible high-level workflows while allowing advanced users to inspect, modify, extend, or selectively reuse its components for complex quantum-mechanical problems.

  • Takeaways & Limitations

    Diagonal-shift regularization can bias real-time evolution, continuous symmetry groups lack a compact computational-basis representation, and several advanced capabilities remain planned.

Abstract

from arXiv · show

We introduce version 3 of NetKet, the machine learning toolbox for many-body quantum physics. NetKet is built around neural-network quantum states and provides efficient algorithms for their evaluation and optimization. This new version is built on top of JAX, a differentiable programming and accelerated linear algebra framework for the Python programming language. The most significant new feature is the possibility to define arbitrary neural network ansätze in pure Python code using the concise notation of machine-learning frameworks, which allows for just-in-time compilation as well as the implicit generation of gradients thanks to automatic differentiation. NetKet 3 also comes with support for GPU and TPU accelerators, advanced support for discrete symmetry groups, chunking to scale up to thousands of degrees of freedom, drivers for quantum dynamics applications, and improved modularity, allowing users to use only parts of the toolbox as a foundation for their own code.

1 Introduction

NetKet 3 addresses the difficulty of applying standard machine-learning frameworks to non-standard quantum problems by combining quantum primitives with JAX in an accessible, extensible Python toolbox. It adds customizable neural-network states, accelerator support, dynamics, symmetries, modularity, and other capabilities for many-body quantum simulations.

  • Motivation: Standard machine-learning frameworks are not designed for many-body quantum problems, while from-scratch implementations can require expertise and deliver sub-optimal performance.
  • Core design: NetKet 3 integrates quantum-mechanical primitives with JAX, enabling pure-Python models, just-in-time compilation, automatic gradients, and CPU/GPU execution.The framework also integrates JAX with MPI to scale across hundreds of computing nodes.
  • Core design: Users can define custom neural-network quantum states in JAX, prototype quickly, avoid manual gradient implementation, and port models to GPUs.
  • New capabilities: NetKet 3 supports unitary and Markovian dissipative dynamics, continuous systems, composite Hilbert spaces, large-scale QGT and stochastic reconfiguration, and arbitrary discrete symmetries.The QGT and stochastic-reconfiguration implementations scale to models with millions of parameters.
  • Extensibility: Multiple dispatch lets users override internal algorithms without editing NetKet’s source, supporting custom objects and algorithms for novel problems.
  • Limitations: JAX custom classes require proficiency with JAX-compatible code, while TPU support was limited to half-precision float16 and required workarounds for precision and gradient underflow.At the time described, JAX primarily supported macOS and Linux, and GPU use required Linux with a recent CUDA version.
  • Usability and scope: NetKet provides high-level defaults for beginners alongside lower-level interfaces, modular quantum objects, workflow drivers, symmetry tools, examples, benchmarks, and experimental APIs.The toolbox includes components for Hilbert spaces, operators, variational states, samplers, QGTs, optimization, dynamics, and symmetries.

2 Quantum-mechanical primitives

NetKet’s quantum-mechanical primitives represent Hilbert spaces, operators, and composite systems while supporting efficient calculations beyond exact diagonalization scales. Its operator formats and row-access interface avoid storing exponentially large matrices explicitly.

  • Workflow: NetKet workflows define a Hilbert space and Hamiltonian or super-operator before combining a neural-network model with a stochastic sampler in a variational state.
  • Hilbert spaces: Hilbert-space classes cover discrete spin, qubit, Fock, and fermionic-orbital systems, continuous particles, constrained subspaces, composite spaces, and doubled spaces for density matrices.Discrete spaces can impose constraints such as fixed total spin or particle populations.
  • Hilbert spaces: Local quantum numbers define general discrete Hilbert spaces, while supported constraints restrict the basis to a subspace such as the zero-magnetization sector.
  • Hilbert spaces: Different Hilbert spaces can be composed with exponentiation and multiplication operators to model coupled systems such as bosonic cavities with spin degrees of freedom.
  • Operators: NetKet represents operators relative to a Hilbert space and provides ready-made Pauli, bosonic, projection, and common Hamiltonian operators.
  • Operators: Standard sparse matrices remain exponentially large, so NetKet offers custom LocalOperator, PauliStrings, FermionOperator2nd, and specialized implementations with use-case-specific scaling.LocalOperator scales linearly with operator terms and degrees of freedom but exponentially with locality K; PauliStrings has negligible memory cost for qubit-like spaces.
  • Operators: Operators support Python arithmetic and conversion to dense or SciPy sparse matrices, but are optimized for efficient row indexing through get_conn(sigma).The row-query interface returns connected basis-state indices and corresponding non-zero matrix elements without requiring full matrix storage.

3 Variational quantum states

NetKet 3 organizes variational quantum states around a common interface for pure and mixed states, combining parameterized models, samplers, and reusable evaluation algorithms. It supports exact and Monte Carlo state types, automatic differentiation, multiple dispatch, and customizable neural-network ansätze.

  • Variational-state interface: A variational state represents a parameterized pure state or density matrix through a common interface that works with high-level optimization and dynamics drivers.The interface exposes parameters, expectation values, gradients, and quantum geometric tensors.
  • Variational-state types: NetKet provides ExactState, MCState, and MCMixedState for exact summation, MCMC sampling, and sampling density-matrix diagonal entries, respectively.The first two represent pure states, while MCMixedState represents mixed states.
  • Algorithm selection: Multiple dispatch selects the most specific expectation-value or gradient algorithm for each variational-state and operator combination while permitting user-defined dispatch rules.JAX just-in-time compilation supports this modular algorithm-selection approach.
  • Models and ansätze: Models map input configurations to quantum log-amplitudes or density matrices, while variational-state classes separately store the parameters.Log-amplitudes expose log-derivatives through automatic differentiation and help represent amplitudes spanning several orders of magnitude.
  • Models and ansätze: Custom neural-network ansätze can be defined as JAX-compatible Python models, primarily through Flax modules, with compatibility for Haiku and init/apply function pairs.NetKet also supplies predefined modules through nk.models and supports custom Flax implementations.
  • Sampling and differentiation: General ansätze use Metropolis–Hastings MCMC sampling, whose transition correction accounts for asymmetric proposal kernels.Direct sampling is available for some models, while complex-parameter gradients use conjugate gradients for real-valued objectives.
  • Quantum geometric tensor: QGT diagonal-shift regularization is effective for ground-state search but can bias real-time evolution away from the physical trajectory.SVD-based alternatives can suppress physically relevant dynamics and increase computational time.

4 Algorithms for variational states

NetKet provides high-level drivers for ground-state optimization, steady-state search, and time-dependent dynamics, while exposing lower-level building blocks for custom algorithms and optimization loops.

  • Built-in drivers: NetKet supplies drivers for VMC ground-state search, time-dependent variational Monte Carlo, and steady-state open-system dynamics.These drivers are constructed from physical models, variational states, and optimization components.
  • Ground-state optimization: VMC estimates Hamiltonian gradients from fresh variational-state samples drawn from the Born distribution.The force vector uses the Hamiltonian local estimator and wave-function log-derivatives.
  • Ground-state optimization: Gradient updates can use standard optax optimizers or stochastic reconfiguration through a quantum-geometric-tensor preconditioner.The preconditioned gradient is passed to the optimizer before updating the variational parameters.
  • Steady states: The steady-state driver minimizes the Frobenius norm of the density-matrix time derivative to find stationary states of an arbitrary super-operator.Its stochastic gradients use the Born distribution of the vectorized density matrix.
  • Time-dependent dynamics: TDVP-based t-VMC propagates pure and mixed variational states, with γ = −i corresponding to approximate real-time Schrödinger evolution.NetKet provides a TDVP driver and JAX-based Runge–Kutta solvers, described as preview functionality.
  • Custom algorithms: Users can implement custom drivers and optimization loops from NetKet’s core NQS building blocks, supporting rapid prototyping beyond the included use cases.The built-in TDVP driver adds features such as error handling, step-size control, higher-order integrators, and just-in-time compilation.

5 Symmetry-aware neural quantum states

NetKet supports symmetry-aware neural quantum states for discrete permutation groups, including nontrivial irreducible representations, while its approach does not cover continuous symmetry groups such as SU(2) or SO(3).

  • Symmetry-aware states: Discrete symmetry restrictions can improve variational convergence and accuracy, and can target excited states that are lowest in selected symmetry sectors.NetKet supports groups isomorphic to permutations of the computational basis.
  • Limitations: The symmetry machinery cannot currently represent continuous groups such as SU(2) or SO(3) because they lack compact computational-basis representations.Efficient encodings for continuous groups remain an open research problem.
  • Group convolutional neural networks: GCNNs generalize convolutional networks to finite, possibly non-commuting groups and can transform under arbitrary irreducible representations.The trivial representation recovers an invariant model.
  • Symmetry groups and representation theory: NetKet can construct lattice space groups, compute irreducible representations and character tables, and generate nontrivial irrep characters for symmetry projection.Space-group irreps can be specified through symmetry-related wave vectors and little-group irreps.
  • GCNN construction: Pre-built GCNNs and reusable layers support custom symmetric ansätze, with symmetry groups specified through permutation-group or lattice interfaces.GCNN outputs are fully symmetric by default, while other irreducible representations can be selected through characters.
  • GCNN construction: GCNN implementations use either group Fourier transforms or fast Fourier transforms on translation-group cosets, with equivalent behavior but different performance and calling sequences.NetKet provides an automatic mode in addition to the two explicit implementations.

6 Quantum systems with continuous degrees of freedom

NetKet provides continuous-space Hilbert spaces, operators, samplers, and VMC workflows for variational simulations of particle systems. Examples cover harmonic confinement and interacting bosons with periodic boundaries, symmetry requirements, and specialized ansätze.

  • Continuous-space Hilbert spaces represent particle positions in individual domains, including R^d or finite boxes with periodic boundaries.
  • Continuous operators support expectation values and gradients but cannot be converted to matrices or used for exact diagonalization.
  • MetropolisSampler supports continuous degrees of freedom, while GaussianRule proposes moves using customizable random shifts.
  • For 10 particles in three dimensions with harmonic confinement, VMC recovers the exact ground-state energy E0 = 15 to very high accuracy.
  • The one-dimensional interacting-boson example combines permutation-invariant DeepSets, Kato’s cusp condition, periodic-distance handling, and VMC energy optimization.

7 Example: Finding ground and excited states of a lattice model

NetKet demonstrates symmetry-aware VMC for ground and excited states of the 10×10 square-lattice J1–J2 model. Spatial symmetry projection, sector selection, chunking, and transfer learning support practical GPU calculations.

  • The 10×10 square-lattice J1–J2 example uses J1 = 1 and J2 = 0.5 inside the spin liquid phase.
  • A GCNN ansatz enforces spatial symmetries, while sampling restrictions select singlet states and specified spin-parity sectors.
  • Chunking splits O(NsamplesL2) network evaluations into smaller calculations to reduce GPU memory consumption.
  • About 30 minutes on an NVIDIA A100 GPU produces a fairly accurate variational energy for the ground-state calculation.
  • Low-lying excited states are obtained by projecting onto different irreducible representations and initializing from parameters optimized for the ground-state sector.

8 Example: Fermions on a lattice

NetKet’s experimental fermion functionality supports finite-orbital lattice systems, fermionic operators, fixed spin populations, and compatible sampling schemes. The examples connect native operators with OpenFermion representations and transformations.

  • Fermion functionality remains under netket.experimental because parts of its API may still change as community feedback is gathered.
  • SpinOrbitalFermions represents fermions occupying orbitals with optional spin and supports fixing up- and down-spin particle numbers.
  • FermionOperator2nd represents second-quantized fermion operators without separating spin and orbital indices.
  • NetKet can construct fermionic operators from OpenFermion objects and convert OpenFermion operators into PauliStrings representations.
  • Sampling must preserve conserved total spin magnetization, so MetropolisExchange moves fermions rather than changing site populations independently.

9 Example: Real-time dynamics

NetKet applies neural quantum states and variational dynamics to unitary quenches and Lindblad open-system evolution. The examples compare variational trajectories with exact solutions and identify a time-step-dependent accuracy boundary for dissipative dynamics.

  • The transverse-field Ising example uses an RBM neural quantum state to simulate a quench and estimate transverse magnetization.
  • Complex-valued RBM weights are required during time evolution to represent the wave function’s complex phase.
  • The unitary quench uses second-order Heun time stepping with dt = 0.01 and compares variational dynamics with an exact solution.
  • The dissipative example evolves a random density matrix under the Lindblad Master equation and compares convergence and magnetization observables with exact dynamics.
  • Dissipative evolution is near-exact where dissipative terms dominate, but develops sizable error when unitary dynamics contributes; smaller time steps can reduce the error.

10 Benchmarks

The benchmarks evaluate NetKet’s VMC, QGT, and MPI scaling on transverse-field Ising models, and compare its QGT-based performance with jvmc. NetKet shows strong CPU and GPU performance, while iterative solvers and MPI improve scalability toward larger models and systems.

  • Benchmark setup: NetKet’s benchmarks measure VMC scaling with NQS complexity, computational resources, and QGT construction combined with repeated gradient-vector multiplications.The experiments use 1D/2D transverse-field Ising models and report timings after just-in-time compilation.
  • Comparison with jVMC: NetKet outperforms jvmc by almost an order of magnitude on a 32-core CPU using SVD-based solvers.The comparison uses complex-valued RBMs on a 64-site one-dimensional transverse-field Ising model.
  • Comparison with jVMC: NetKet outperforms jvmc by about 50% in a full VMC iteration on GPU.Both packages scale poorly from one GPU to two because QGT diagonalization is the bottleneck and cannot be parallelized in this benchmark.
  • MPI scaling: Increasing MPI ranks reduces total computation time by a factor approaching the number of ranks before speedup saturates.Multiple Markov-chain samplers can run in parallel even on a single node, where JAX otherwise cannot use multiple CPU cores except for very large matrices.
  • QGT solvers: Iterative QGT solvers can scale to potentially millions of parameters and larger systems, with the CG solver becoming significantly more efficient on GPU.This approach avoids the diagonalization bottleneck that limits parallelism in the SVD-based benchmark.

11 Discussion and conclusion

NetKet 3 is presented as a modular toolbox for complex quantum-mechanical problems, with arbitrary neural-network ansätze, reusable components, and improved scalability. The discussion also identifies missing capabilities and future development priorities.

  • Discussion and conclusion: NetKet 3 lets users define arbitrary neural-network ansätze for wave functions or density matrices with JAX and use only selected toolbox components.Users can separately employ architectures, samplers, the quantum geometric tensor, or operators without the VariationalState interface.
  • Discussion and conclusion: More efficient geometric-tensor implementations and other algorithmic improvements target systems and models with millions of parameters.The JAX foundation also enables effective GPU use without manual low-level programming for those platforms.
  • Future directions: Future work includes native fermionic systems, more general continuous-system geometries, broader dynamics support, additional drivers, and more general cost functions.The authors describe these as capabilities they would like to integrate in future versions.
  • Symmetry-aware architectures: GCNNs construct symmetry-related wave functions and combine them with a projection operator to form symmetric wave functions.Their layers must remain equivariant under space-group transformations, and the final features are converted into wave functions.
  • Symmetry-aware architectures: Summing many terms in the GCNN ansatz appears to improve variational-optimization stability for sign-problematic Hamiltonians.The passage presents this as an observed effect in addition to support for nontrivial symmetries.

A.2 Fast group convolutions using Fourier transforms

The section develops two Fourier-based approaches to reduce the memory and runtime costs of group convolutions, replacing the straightforward enlarged-kernel implementation. Group Fourier transforms apply to arbitrary finite groups, while FFTs exploit the Abelian translation subgroup of space groups and can scale better on large lattices.

  • Motivation: The direct implementation expands each kernel into a |G| × |G| matrix, requiring O(f_inf_out|G|^2) time and substantial memory.It is easiest to interpret and code, but becomes problematic on modestly sized lattices and networks.
  • Group Fourier transforms: Group Fourier transforms preserve input, output, and kernel memory while evaluating convolutions through representation-space matrix products.The forward transform maps functions on group elements to matrices for group irreducible representations, and the backward transform sums over inequivalent irreps.
  • Group Fourier transforms: For large space groups, the group Fourier-transform runtime scales as O[(f_in + f_out)|G|^2 + f_inf_out|G||P|].Most irreps are defined on stars of |P| wave vectors and have dimension |P|, producing the stated scaling.
  • FFT-based approach: The FFT-based approach transforms kernels and features over translation-group indices, contracts Fourier components pointwise, and transforms the result back.The translation subgroup is Abelian and corresponds to valid lattice vectors, so the translation sum is a standard convolution that can use cuDNN routines or FFTs.
  • Implementation caveat: Complex irreducible representations can temporarily double memory use because Fourier transforms of real inputs, outputs, or kernels become complex.This caveat applies when some irreps cannot be represented by real-valued matrices.
  • FFT-based approach: FFT-based convolutions scale as O[(f_in + f_out)|G| log |T| + f_inf_out|G||P|], improving asymptotically over group Fourier transforms for large lattices.On CPUs the FFT-based approach tends to be faster, while GPU performance depends on efficient parallelization and matrix-multiplication behavior.

B Implementation details of the quantum geometric tensor

The implementation computes quantum-geometric-tensor actions efficiently using Jacobian products rather than always constructing or storing full matrices. NetKet provides both Jacobian-based and on-the-fly approaches, using forward- and backward-mode automatic differentiation.

  • QGT implementation: The quantum geometric tensor can act on a vector without storing the full matrix.This reduces the computational burden compared with constructing the full quantum geometric tensor.
  • Jacobian representation: For batched neural-network states, the Jacobian rows correspond to input samples and columns to variational parameters.The function is evaluated over samples s_k, producing one gradient row per sample.
  • Automatic differentiation: Jacobian–vector products use forward-mode automatic differentiation, whereas vector–Jacobian products use backward-mode automatic differentiation.JAX exposes both primitives and builds higher-level differentiation functions from them.
  • QGTJacobian: QGTJacobian stores the full sample Jacobian and computes QGT–vector products without explicitly forming the full quantum geometric tensor.The Jacobian can be assembled row by row using vector–Jacobian products, with one-sample backpropagation in a batch.
  • QGTOnTheFly: When the full Jacobian cannot fit in memory, QGTOnTheFly computes the required expressions directly with vector–Jacobian and Jacobian–vector products.Its first and last operations are implemented with jax.vjp and jax.jvp, respectively.
Loading 2112.10526v2…