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

### Submission summary

 As Contributors: Giuseppe Carleo · Damian Hofmann · Attila Szabó · Filippo Vicentini Preprint link: scipost_202112_00047v1 Code repository: https://github.com/netket/netket Date submitted: 2021-12-20 15:10 Submitted by: Vicentini, Filippo Submitted to: SciPost Physics Codebases Academic field: Physics Specialties: Condensed Matter Physics - Computational Quantum Physics Approach: Computational

### Abstract

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.

###### Current status:
Editor-in-charge assigned

### Submission & Refereeing History

Submission scipost_202112_00047v1 on 20 December 2021

## Reports on this Submission

### Strengths

1) Innovative topic

2) Timely toolbox for quantum many-body physics

3) Clear language and demonstrative example codes

### Report

This manuscript introduces the new version 3 of the NetKet python package. The NetKet package (or library) is used to study mangy-body quantum systems by leveraging recent machine learning techniques.

Machine learning and quantum computational technologies are getting a great deal of attention because of their efficiency in solving important problems in condensed matter physics and quantum chemistry and material sciences. Therefore, this manuscript is very timely.

Moreover, it is well-written and has a broad range of introductory examples helping readers to get started with the NetKet 3 package right away (the GitHub page and the Documentation page for the package are also extremely well-constructed and easy to use).

Given the interesting topic and the clarity of the material, I recommend publication in SciPost Physics Codebases.

As a further comment, the introduction of JAX as acceleration framework and the multiple dispatch through PLUM are very welcomed additions that make this new version of NetKet much faster than before, and therefore of better applicability to large and challenging physical problems.

Minor comments and clarifications are included below.

### Requested changes

1) NetKet 3 is based on JAX, a differentiable programming framework. Even though the high-level primitives as well as most of the examples shown do not require any knowledge of JAX, to really use the full potential of NetKet 3 users need to be fluent in JAX. I think this may not be stressed enough, but I may be wrong. May I suggest to add a clear sentence to Sec. 1.1 or 1.2 discussing how important a good knowledge of JAX is for using NetKet 3? (not just for building NQS, but also for the creating and debugging algorithms)

2) Linear operators introduces in section 2.2 are some of the most important objects in the library. However, I do not find in the text any explicit mention that an operator can only be defined on a fixed Hilbert space (e.g. the first argument in the constructor of nk.operator.spin.sigmax on the example of page 8). May I suggest to explicitly add a sentence on the dependence of operators on Hilbert spaces in section 2.2?

3) In equation (7) the upper limit of the sum over j, M, is not defined. It would help to have a definition of M for this simple case of a one-layer fully-connected ansatz.

4) On page 20, the first code snippet introduces F, while the second code snippet introduces g . May I suggest to add an explanation to clarify their meaning?

5) Section 4 introduces the 3 main algorithms (drivers) of NetKet 3. However, while 4.2 and 4.3 include "full" examples, 4.1 does not have one (and 4.3 actually defers the example to Sec. 8). Since 4.1 and 4.2 are rather similar in their "syntax" and the examples are not too long or complicated, would it be possible to add a full exampled in section 4.1 as well? (the one in Sec 7 is longer and more demonstrative, so I am only asking to get a short overview where all the "pieces" are shown: Hilbert, Operator, Sampler, NQS, VMC)

* typo in code snippet line 2 "latticfe"
* after equation (27) the operator.LocalLiovillian is missing a "u" in the name and should be referred to also as "lind" which issued for "lind.hilbert" in the subsequent lines.

7) Section 5.1. I have a question here. NetKet has support for some discrete (only?) symmetry groups. Can this support be extended and how difficult it is? I think it would be helpful to have a sentence/paragraph related to this.

8) Continuous Hilbert spaces are a nice addition to NetKet 3. The example in section 6.4 for the harmonic oscillator modeled with a Gaussian is interesting to showcase the NetKet computational elements. I wonder if it could be more challenging to add interactions between the the oscillators so that the ground state is not just n _x = n_y = n_z  =0 for all particles.
Some extra comments on this section:
* L=(jnp.inf..) How is an infinite length for the space interpreted, together with the boundary conditions?
* Code snippet line 19: "sab" should be the sampler "sa"
* In Fig. 1 the energy for the variational state starts from a really large value. What is the initialization of the Gaussian NQS?

9) In section 7.2, a full example of the optimization of a GCNN variational state for a lattice model is presented. Here I see for the first time (correct me if I am wrong) the use of "chunk_size" in the code snippet at line 21.
"Chunking" is mentioned in the abstract and then never again. As far as I understand from NetKet's documentation, chunking allows to use large NN models that would not fit in memory. Since chunking is mentioned in the abstract, I would like to see a description of it in a dedicated subsection perhaps. If I have missed it, please forgive me.

10) Question about footnote 18 on page 33. Although this "non-standard MC scheme" can be used to draw (almost) independent samples, a phase of burn in is required. The code snippet on line 20 seems to imply that this burn-in phase is reduced to 0. That would not give rise to the correct sampling at the beginning. Am I missing something? Maybe related to how n_discard_per_chain=0 actually works?

11) The addition of dissipative dynamics to NetKet is very important. In fact, references 26, 26 and 27 have been published on PRL with a Veiwpoint (see https://physics.aps.org/articles/v12/74 ). From the Viewpoint I see that the last paper, the PRB one, is missing from the citations of this manuscript. Can it be added? Or is it not relevant to the techniques implemented in NetKet? Moreover, given the importance and novelty of dissipative dynamics solved with NQS, is it possible to add a figure with results in section 8.2?

12) In the benchmark section 9.2, the next to last paragraphs refers to Fig. 5, but I think it should be Fig. 4. Also, why are there two blue lines in the right panel of fig 5? In the caption, the discussion about samples should refer to the left and center panels and not the left and right panels.

• validity: high
• significance: high
• originality: high
• clarity: high
• formatting: excellent
• grammar: excellent

### Strengths

1)This manuscript describes the Version 3 of the NETKET library to simulate quantum systems using artificial neural networks. The main innovation compared to version 2 is the use of the JAX python framework, which allows easy utilization of just in time compilation, of GPUs, and it provides automatic differentiation.

2) The library allows users to simulate ground states, unitary dynamics, and open system dynamics. Several primitive functionalities are exposed to the user, so that scientists can implement user-defined algorithms to tackle quantum many body problems via machine learning tools.

3) Several examples of the computational tasks are presented with sufficient details to allow new users get started with the utilization of NETKET 3.

### Weaknesses

1) As the authors mention in the conclusions, there is no native support for fermionic systems.

2) As the authors discuss, one of the main strengths of NETKET 3 is that it allows users to implement arbitrary neural network structures. What is missing, perhaps, is some discussion or some guidance on which structure is more efficient or more promising for different physical models and different computational task.

### Report

The version 3 of the NETKET library is described. It is implemented to be flexible, allowing users to address various computational tasks in quantum many-body physics. It provides users useful ready-to-use classes for relevant neural networks states, and also primitive tools to build user-defined algorithms.
The manuscript is well structured. It presents the design philosophy, the user interface, as well as the building blocks that can be used in a flexible manner. Several case studies are presented in a pedagogical manner.
The main innovation compared to NETKET 2 is the use of the JAX library, providing easy utilization of GPUs.
In this respect, NETKET 3 is analogous to the jVMC library (arXiv:2108.03409, Ref.[38]). However, NETKET 3 provides further functionalities not included in jVMC. A performance comparison is presented, showing a computational advantage for NETKET 3.

The use of machine learning tools to simulate equilibrium and dynamical properties of quantum systems is a growing trend in condensed matter physics. I think that the novel developments in NETKET 3 will help computational physicists to develop novel and more efficient algorithms. In this manuscript, the library is well described with various use cases. Therefore, I think that the manuscript is suitable for publication in SciPost Physics Codebases.
In the requested changes, I report a few comments that should be addressed before publication. They are aimed at improving the clarity of some parts of the manuscript.

### Requested changes

1) The authors discuss the implementation of complex-valued wave-functions, mentioning the need for complex weights and biases. In Ref.[38], three main approaches (Single holomorphic network, Single non-holomorphic network, Two real (non-holomorphic) networks) to implement complex valued wave functions are described. I encourage the authors to clearly discuss the different possibilities for this task available in NETKET. A small paragraph or section about this topic would help the readers.

2) In the note 13, the authors briefly mention the possible instabilities in the calculation of the QGT. I encourage the authors to provide a more exhaustive discussion about this problem, with some guidelines to avoid it.

3) In subsection 6.1, the authors discuss a continuous-space problem. Periodic boundary conditions are implemented with a specific flag (pbc=True). It is not clear if pbc=Flase corresponds to hard-wall type boundaries. I suggest to include a specific comment.

4) The use case for continuous-variable problems corresponds to a non interacting system. I suspect that including interactions leads to more severe challenges. Can the authors provide a more challenging testbed, or some comments and guidelines on how to deal with interactions in continuous-variable problems?

5) In figure 5, it is not clear what time unit is used for the color scale of the right.

6) As the authors discuss, one of the main strengths of NETKET 3 is the possibility to implement arbitrary neural network structures. However, the authors do not provide a guideline on which structure is more suitable, depending on the problem being addressed. In the presented use cases, they sometimes employ restricted Boltzmann machines, sometimes GCNNs. Different network architectures have already been explored in the literature, e.g., dense layers, sparse networks, convolutional, graph NN (see, e.g, Phys. Rev. X 8, 011006 (2018), Nat. Comm. 9, 5322 (2018), Phys. Rev. E 101, 063308 (2020), arXiv:2011.12453, 2020), finding different performances for different problems. I encourage the authors to provide some comments to motivate and guide users in the exploration of different architectures.

7) In the conclusions, the authors might provide more comments on how they plan to implement support for fermionic systems. Notice that different approaches have been explored in the literature:
Nature Chemistry 12.10, 891-897 (2020), arXiv:2112.03491, Physical Review Research 2, 10.1103 (2020), Nature Communications 11, 2368 (2020).

• validity: high
• significance: high
• originality: good
• clarity: high
• formatting: excellent
• grammar: excellent