Scalable Neural Wave Functions via Information Mixing and Hierarchical Network Architectures

arXiv:2610.30257 · quant-ph, cond-mat.dis-nn · Submitted 2026-10-21 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: Quantum Radio. Generated commentary on the latest quantum physics and condensed matter papers.

Kai: We're looking at the paper: "Scalable Neural Wave Functions via Information Mixing and Hierarchical Network Architectures".

Mira: Here's the gist of the paper. The gist: The MLP mixer architecture, when combined hierarchically with a transformer model, outperforms other state-of-the-art architectures for large-scale quantum systems by efficiently capturing long-range correlations.

Kai: Theoretical studies of systems of interacting quantum particles remain a daunting task owing to the exponential scaling of the Hilbert space with system size The strength and computational efficiency of ANNbased wave function representations depend strongly on the underlying ANN architecture.

Mira: While recurrent neural networks (RNNs) are computationally efficient, the recurrent nature can limit their capacity to capture long-range correlations.

Kai: Transformer models with their underlying self-attention mechanism excel at capturing these long-range correlations, but incur increased computational requirements, restricting applications to small physical systems.

Title and authors: Kai: So we're talking about this paper right now, "Scalable Neural Wave Functions via Information Mixing and Hierarchical Network Architectures." It’s basically looking at how to make neural networks actually useful for simulating really big quantum systems, where the Hilbert space just explodes.

Mira: Right. The core idea here is that we can use a multi-layer perceptron mixer model combined hierarchically with a transformer to efficiently capture those long-range correlations that standard recurrent networks just can't handle well on their own.

Lev: But from an engineering standpoint, it's important to know this architecture is built entirely from feed-forward operations, which makes it computationally efficient compared to the other models we’ve looked at in this field.

Kai: Exactly. The paper explores how you can use patches of input—say, a K by K array of particles—and mix information both locally within those patches and globally across all the patches.

Mira: And what’s really interesting is the way they structure that mixing: first, sharing information globally via trainable non-linear transformations over all input patches, and then sharing it locally within each patch with parameters shared between global and local layers.

Lev: That sounds like a lot of parameter tuning, which means we have to think about how much memory that takes up when you actually run the simulation on real hardware.

Kai: Well, the authors address that by adding a causal mask to the global information mixing layer, which helps achieve autoregressive behavior while still allowing it to capture those long-range dependencies in quantum systems.

Mira: And they tested this setup variationally for two specific problems: the two-dimensional transverse field Ising model and arrays of Rydberg atoms, comparing it against recurrent neural networks and transformer models.

Lev: So what’s the big finding there? What actually worked when you put these architectures to the test?

Kai: It shows that while recurrent neural networks are accurate for smaller systems, this MLP mixer architecture outperforms both of those models on accuracy and efficiency for larger system sizes, specifically beyond eight by eight qubits.

Mira: That’s because they're using a hierarchical approach—a main model that handles large patches and then decomposes it into smaller patches if needed—which helps manage the exponentially growing network output dimension.

Lev: I see what you mean about the decomposition, but does that hierarchy actually translate to a practical speedup on actual hardware? Because for me, we're always worried about the overhead of running these complex networks.

Kai: The paper suggests that this hierarchical construction gives us a competitive representational accuracy while maintaining relatively low computational costs compared to some other patched models.

Mira: And they specifically point out that the MLP mixer model scales linearly with input sequence length, which is better than the quadratic scaling you see in standard transformer models for certain types of tasks.

Lev: Linear versus quadratic scaling is a big deal when you're trying to scale up, especially since we’re dealing with systems where the state space grows so fast.

Kai: And they show that for larger system sizes, like twenty-four by twenty-four qubits, the hierarchical MLP mixer architecture shows an energy error similar in size to what we see with a hierarchical recurrent network.

Mira: So, while it’s competitive in accuracy at those larger scales, the comparison against the RNN still suggests that for smaller systems, the RNN is still more accurate and efficient.

Lev: That makes sense from a resource constraint point of view; if you only have limited resources, you stick with what gives you the best performance for that specific size.

Kai: The overall implication of this paper is that there isn't one single best architecture, but rather the right choice depends heavily on the system we're studying and what resources we have available.

Mira: They are essentially showing that combining a transformer model with an MLP mixer model in a hierarchical way opens up new ways to represent quantum many-body states beyond the limitations of conventional ANN methods.

Lev: It sounds like they’ve given us a framework for building something more robust for these large-scale simulations, even if the computational cost is still substantial.

Kai: So, before we wrap this up, I want to make sure we cover what the authors themselves admit about their method's limitations.

Mira: They do admit that single-level patching alone isn't sufficient when dealing with larger system sizes because you need that hierarchical decomposition to get the best scaling results.

Lev: That’s a fair limitation to point out; they aren't claiming this is a universal solution for every quantum problem.

Kai: So, to summarize "Scalable Neural Wave Functions via Information Mixing and Hierarchical Network Architectures," we see that the MLP mixer model combined with hierarchical structures is a strong candidate for finding ground state representations of large-scale quantum systems.

Mira: It’s about using this architecture to efficiently capture long-range correlations while keeping the computational cost manageable, especially when compared to standard transformer models.

Lev: For real hardware, it means we should be looking at these hierarchical models when we want to simulate systems that are bigger than what a single-layer patched network can handle effectively.

Kai: That’s the picture for this paper. We'll be looking at other papers on arXiv coming in next.

The paper's summary: Kai: So, we're looking at how this paper handles the big picture of simulating quantum systems that are just too massive for traditional methods, you know, where you can’t just brute force it anymore.

Mira: Yeah, they’re not just throwing another transformer at the problem; they're using this specific setup with a multi-layer perceptron mixer and a hierarchical structure to capture those long-range correlations in a much more organized way.

Kai: It means instead of trying to learn every single interaction point in the whole system at once, this AI mixes information both locally—within small patches of particles—and globally across all those patches.

Mira: And they do that by sharing information across all input patches first through some trainable non-linear transformations, and then they refine that mixing locally within each patch using parameters shared between those global and local steps.

Kai: That whole architecture is designed to be computationally light while still being able to handle the complexity, which is a big deal when we're talking about running this on real hardware.

Mira: They test this against recurrent neural networks and transformer models for things like the transverse field Ising model and Rydberg atoms, and they find that this hierarchical approach helps recover the ground state much faster than those other methods.

Kai: That speedup is something Lev would be very interested in, because running variational training on a real quantum computer takes time, so less training iterations means we save cycles.

Mira: The key part is how they handle the system size; single-level patching doesn't work well for truly large systems like twenty-four by twenty-four qubits, which is why they built that hierarchical extension to decompose the large input patches into smaller ones.

Kai: So, what this really changes for us is that we have a more flexible toolkit now; you can use the patched RNN or transformer for smaller systems where they might be faster, but you switch to this hierarchical MLP mixer when the system gets big and you need better long-range correlation capture.

Mira: It’s about finding a way to get high accuracy on these huge, higher-dimensional quantum many-body systems without the network output dimension just exploding into infinity.

Kai: And that’s what makes it interesting because it suggests we can build wave function ansätze that are both accurate and computationally feasible for these large scale problems.

Mira: But then you have to look at the scaling; they show that this MLP mixer model has a linear scaling with the input sequence length, which is better than the quadratic scaling of standard transformers for certain tasks.

Kai: That’s a crucial distinction when we think about putting this on superconducting qubits or trapped ions, because linear scaling is definitely more manageable than quadratic.

Mira: So it’s not just about getting an answer; it’s about finding a way to get that answer efficiently enough to actually be useful for simulating real physics at scale.

Kai: That's what I find really compelling—it moves the discussion from just "can we simulate this?" to "how do we simulate this efficiently?".

The paper's improvements: Tom: So, we're looking at what the authors suggest as improvements for this architecture, moving beyond just having one big hierarchical model that’s already working on large systems.

Kai: They point out that even with the hierarchical transformer and MLP mixer, you still need to be careful about how you set up those input patches; it’s not just one way to chop up the data.

Mira: Yeah, they suggest there are ways to further decompose those large input patches into even smaller ones at multiple scales, which helps manage that exponentially growing network output dimension better than just a single level of hierarchy.

Kai: That means we could potentially build an even deeper hierarchy, breaking down the system representation into more granular pieces if needed.

Mira: The goal there is to keep the computational cost low while still maintaining that accuracy we saw in the larger system simulations, like on twenty-four by twenty-four qubits.

Kai: It sounds like they’re proposing a modular approach where you can swap out different decomposition strategies depending on how complex the specific quantum problem is.

Mira: Right, and this modularity means we might find better efficiency trade-offs; one strategy might be great for spin systems, another for something else entirely.

Kai: It makes sense because not every large quantum system has the same type of correlation structure, so a flexible setup allows us to tailor the AI architecture to fit the physics.

Mira: And they also hint at making the mixing strategy itself more adaptive, perhaps changing how much global versus local information is shared based on what’s happening in those specific patches.

Kai: That kind of adaptive mixing would be powerful because it could potentially find more efficient pathways to the ground state representation without having to manually tune every single parameter.

Mira: So, the suggestion isn't just about adding another layer; it's about making the entire information-mixing process smarter and more adaptable to the specific quantum data you feed it.

Kai: That’s where things get interesting for hardware guys because if the AI itself can help optimize its own structure based on that adaptive mixing, we might see better convergence speeds during training.

Mira: And ultimately, this points toward a future where we don't just use one fixed neural network ansatz for every quantum problem, but a whole library of intelligently constructed architectures.

Kai: It suggests that the next step isn't just scaling up the existing model, but making the model itself more intelligent about how it processes information at different scales.

Conclusion: Kai: So we’re wrapping up on "Scalable Neural Wave Functions via Information Mixing and Hierarchical Network Architectures," which basically shows how to use an MLP mixer and a hierarchical transformer to make neural networks work for massive quantum systems.

Mira: It boils down to using that specific way of mixing information—globally then locally—to capture those long-range correlations without the network just getting way too big.

Kai: The big picture is that this means we can build wave function ansätze for quantum many-body problems that are actually feasible to train and run on current hardware, not just theoretical constructs.

Lev: For me, what’s interesting is the scaling behavior; they show linear scaling with input sequence length for the MLP mixer, which is way better than the quadratic scaling we usually see in transformers.

Mira: Exactly because that linear scaling helps keep things manageable when you start feeding it really large patches of data from a physical system.

Kai: We’ve seen them test this on models like the two-dimensional TFIM and Rydberg atoms, and they get good results, though they admit the choice of architecture still depends on whether you have small or large systems to start with.

Lev: I'd add that for actual error correction hardware, we need to think about the FLOPs per forward pass; if this linear scaling holds up in practice, it could mean much lower overhead for variational training than some of the older methods.

Mira: The implication is that instead of always defaulting to the most complex model, we have a toolset where you can dynamically choose between patched RNNs or this hierarchical structure based on the system size and what you need to capture.

Kai: So, it’s not just one perfect solution; it’s a flexible framework that lets us adapt our AI architecture to the specific physics of the quantum problem we are trying to solve.

Lev: It gives us a solid baseline for how much computational overhead we can expect when building these state representations for real qubits.

Mira: And I think the most important part is that it opens up new avenues for studying complex quantum condensed matter physics because we can finally tackle those higher-dimensional systems more efficiently.

Kai: We’ll be keeping an eye on how this hierarchical decomposition works as we look at those next papers coming in, especially since the authors hinted at even deeper levels of decomposition.

Lev: I’m looking forward to seeing if these scaling results hold up when we try to map them onto a real noisy quantum processor.

Mira: And I'm curious to see what kind of new physical insights emerge from using this architecture on those large-scale models.

Department of Physics, University of Ottawa

quant-ph, cond-mat.dis-nn

Submitted: 2026-10-21

Updated: 2026-10-21

Comments: 9+5 pages, 9 figures

Code: https://github.com/APRIQuOt/mlp-mixer-wave-function

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 73/100

The gist: The gist: The MLP mixer architecture, when combined hierarchically with a transformer model, outperforms other state-of-the-art architectures for large-scale quantum systems by efficiently capturing

Key concepts

MLP Mixer Model: This architecture captures long-range correlations by mixing information globally and locally through feed-forward vector-matrix multiplications. It processes input as patches, allowing it to distinguish between local information within a patch and global information shared across all patches.
Hierarchical Architecture: This involves a main ANN model processing large input patches and feeding its output into a subsequent model that decomposes the large patch into smaller sub-patches. This multi-scale decomposition allows the system to handle very large input sizes efficiently, similar to patched transformer models.
Ground State Representation: The goal is to train these neural networks variationally to find the lowest energy state of a quantum system given its Hamiltonian. The accuracy and speed of finding this state determine how well the ANN can represent the complex quantum wave function.
Patch-based Models: Instead of processing single data points sequentially, these models take input as arrays (patches) of $K imes K$ particles. This approach reduces the effective sequence length compared to single-particle models, leading to faster training times and better efficiency for large systems.

Summary

The gist: The MLP mixer architecture, when combined hierarchically with a transformer model, outperforms other state-of-the-art architectures for large-scale quantum systems by efficiently capturing long-range correlations.

Introduction

Theoretical studies of systems of interacting quantum particles remain a daunting task owing to the exponential scaling of the Hilbert space with system size The strength and computational efficiency of ANNbased wave function representations depend strongly on the underlying ANN architecture. While recurrent neural networks (RNNs) are computationally efficient, the recurrent nature can limit their capacity to capture long-range correlations. Transformer models with their underlying self-attention mechanism excel at capturing these long-range correlations, but incur increased computational requirements, restricting applications to small physical systems. In this work, we address these limitations by exploring the multi-layer perceptron (MLP) mixer model as a wave function ansatz. By mixing information globally and locally through the system via a computationally lightweight mixing strategy, this architecture promises to combine the ability to capture long-range correlations with low computational costs.

Artificial Neural Network Architectures

The focus of this work lies in comparing the performance and efficiency of different artificial neural network (ANN) architectures when used as wave function ans¨atze Specifically, we consider recurrent neural networks (RNNs), transformer models, and multi-layer perceptron (MLP) mixer models. All three models can efficiently process sequential data in the form (σ1,..., σN) with sequence length N. By using the network output at sequence step i as input in the next sequence step i + 1, the ANN architectures can be defined in an autoregressive manner, which allows for sequential data generation. In this work, we define a one-dimensional path through two-dimensional physical systems in a row-major order. The state σi of a quantum particle then defines the ANN input at sequence step i. The corresponding network output defines a probability distribution, pθ (σi+1 σ0,..., σi) = pθ (σi+1 σ≤i), (1) from which the state σi+1 of the next particle is sampled conditioned on the state of previous particles in the sequence.

Multi-Layer Perceptron Mixer Models

The MLP mixer model achieves computational efficiency by reducing internal operations to feed-forward vector-matrix multiplications and scalar non-linear transformations. Combined with reshaping and transposition operations, the architecture can efficiently capture longrange correlations between input elements. Rather than considering a sequence of individual input elements, the MLP mixer model considers a sequence of patches, each consisting of an array of K × K elements. In such sequences, we distinguish between local information covered within each input patch and global information shared across all patches at individual positions. In the MLP mixer setup, information is first shared globally via trainable non-linear transformations applied over all input patches. In a second step, information is similarly shared locally within each patch, where trainable parameters are shared over all global and local transformations. To become a suitable wave function ansatz, we need to adjust the original MLP mixer architecture. An autoregressive behavior can be achieved by adding a causal mask in the global information mixing layer, eliminating information from input elements that appear later in the sequence.

Patched and Hierarchical Architectures

In consistency with the MLP mixer model, which naturally considers patches of K × K input elements, we further consider patched RNN and transformer models. These receive the state of a two-dimensional patch of K × K particles as input and return the probability distribution over all possible states of the next patch conditioned on all previous sequence patches. However, the number of possible states in a patch of particles scales exponentially with the patch size, leading to an exponentially growing network output dimension. To allow for an efficient usage of larger input patches, we consider hierarchical extensions consisting of a main ANN model that receives large input patches with K > 2 and feeds its output into a subsequent ANN model that hierarchically decomposes the large input patch into a sequence of smaller patches of 2×2 particles. This setup is similar to large, patched transformer models and leads to a fixed output dimension.

Physical Systems Comparison

We compare the accuracy and computational efficiency of the different ANN architectures in finding ground state representations for the two-dimensional transverse-field Ising model and a two-dimensional array of Rydberg atoms. To find ground state representations, we train the ANN models variationally, tuning the weight and bias parameters to minimize the energy expectation value ⟨E⟩ for a given Hamiltonian Hˆ as evaluated on Ns samples σn drawn from the generative ANN. In both cases, results from the nine considered ANN architectures are shown in different colors. However, we find that the patched transformer model, the MLP mixer model, and all hierarchical models recover the ground state significantly faster, in less than 5,000 training iterations.

Results

In Fig. 4, we plot the average absolute energy-density error after 15,000 training iterations versus the number of FLOPs used in a single forward-pass in each ANN architecture for the two-dimensional TFIM Hamiltonian. For N = 4 × 4 qubits, the patched transformer architecture and the patched RNN show the best performance, while the MLP mixer requires more FLOPs. For N = 8 × 8 qubits, hierarchical architectures and MLP mixer models appear stronger in the larger system size. The hierarchical MLP mixer model requires a comparable amount of FLOPs to the MLP mixer model and the hierarchical transformer architectures while achieving a competitive representational accuracy.

Discussion and Outlook

Our study thus shows that the best choice of ANN architecture for ground state representations depends strongly on the system under consideration. However, using the hierarchical network architecture consisting of a transformer model and an MLP mixer model as variational wave function ansatz allows for ground state representations of large-scale quantum systems beyond limitations of conventional ANN-based wave function ans¨atze. The computational cost of the RNN scales linearly with the considered sequence length, while the transformer model shows quadratic scaling with the input sequence length. Patch-based models, such as the patched RNN and the patched transformer model, reduce the effective sequence length compared to single-qubit considerations, leading to faster run times. The MLP mixer model shows linear scaling with input sequence length while being naturally patch-based. Our results show furthermore that single-level patching remains insufficient at large system sizes. The hierarchical constructions outperform the patched models by introducing a decomposition at multiple scales.

Appendix A and B

The MLP mixer architecture relies on MLPs, reshaping, transposing, and scalar nonlinearities. The first operation mixes information per location in the input across multiple channels, while the second operation mixes information across all locations in each input channel. In Appendix B we implement the autoregressive setup in the MLP mixer architecture. For this, we apply a causal mask before information is shared globally across inputs. The skip connection that appears after the tokenmixing operation naturally breaks the order of the input sequence as information from the current sequence element can leak in.

Appendix C and D

The forward-pass FLOPs provide a model-intrinsic, hardware-agnostic measure of computational complexity. The results show that for large system sizes, combining the hierarchical transformer architecture with the MLP mixer model thus shows the most promising combination. The hierarchical MLP mixer architecture shows the most promising FLOP scaling behavior. However, at N = 24 × 24, the hierarchical MLP mixer shows an energy error of similar size to the hierarchical RNN. While the consideration of larger system sizes is necessary to draw a certain conclusion, our simulations indicate that the most promising choice in terms of FLOPs and representational accuracy for investigating large-scale quantum systems is the hierarchical transformer architecture with MLP mixer model.

Acknowledgments

We thank Sehmimul Hoque and the APRIQuOt team for fruitful insights and discussions We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), NSERC CREATE/575280-2023 - Training in Materials for Quantum Technologies (MaQTech) This research was enabled in part by computational resources and support provided by the BC DRI Group (bcdri.ca) and the Digital Research Alliance of Canada (alliancecan.ca) Additional computational resources were provided by Google Colaboratory (colab.research.google.com) and Lambda (lambdalabs.com).

Code Availability

We wrote our codes in the Python programming language. We used the PyTorch (v2.

Improvements for AI systems

Adaptive architecture selection based on system size: The system can dynamically choose between models based on required resources, as the best choice of ANN architecture for ground state representations depends strongly on the system under consideration. For small systems like 4x4 qubits, it selects the patched RNN or patched transformer model for efficiency, while for larger systems beyond 8x8 qubits, it switches to the hierarchical transformer with MLP mixer model.
Enhanced representation of long-range correlations: The system can capture complex quantum states by using a hierarchical combination of this model with the transformer model which outperforms other state-of-the-art architectures for large-scale quantum systems. This allows it to investigate novel insights into quantum condensed matter physics by enabling efficient and accurate studies of large-scale, higher-dimensional quantum many-body systems.
Improved computational efficiency for scaling: The AI can achieve low computational costs while maintaining accuracy because the MLP mixer model shows linear scaling with input sequence length while being naturally patch-based, which is superior to the quadratic scaling of transformer models for certain tasks, especially when combined hierarchically.