Machine-Learned Hamiltonians for Quantum Transport Simulation of Valence Change Memories

TL;DR

This work tackles the prohibitive cost of constructing ab initio Hamiltonians for large, disordered devices by learning Hamiltonians with an equivariant graph neural network and augmented partitioning, demonstrated on a TiN-HfO$_2$-Ti/TiN VCM cell with about $N \approx 5{,}268$ atoms. The EGNN predicts the Hamiltonian blocks with MAE in the $3.39$–$3.58\ \mathrm{meV}$ range, allowing the use of a quantum transport solver via NEGF to obtain $T(E)$ and $I_d$ with a high level of qualitative agreement to DFT, while achieving a dramatic speedup (≈$2$ s per forward pass vs. hours for DFT). This enables rapid, large-scale transport studies of evolving morphologies and supports exploration of devices beyond current ab initio capabilities. The approach holds promise for efficient simulation of phase-change and other resistive-memory systems where large amorphous regions and vacancy dynamics determine electronic properties.

Abstract

The construction of the Hamiltonian matrix \textbf{H} is an essential, yet computationally expensive step in \textit{ab-initio} device simulations based on density-functional theory (DFT). In homogeneous structures, the fact that a unit cell repeats itself along at least one direction can be leveraged to minimize the number of atoms considered and the calculation time. However, such an approach does not lend itself to amorphous or defective materials for which no periodicity exists. In these cases, (much) larger domains containing thousands of atoms might be needed to accurately describe the physics at play, pushing DFT tools to their limit. Here we address this issue by learning and directly predicting the Hamiltonian matrix of large structures through equivariant graph neural networks and so-called augmented partitioning training. We demonstrate the strength of our approach by modeling valence change memory (VCM) cells, achieving a Mean Absolute Error (MAE) of 3.39 to 3.58 meV, as compared to DFT, when predicting the Hamiltonian matrix entries of systems made of $\sim$5,000 atoms. We then replace the DFT-computed Hamiltonian of these VCMs with the predicted one to compute their energy-resolved transmission function with a quantum transport tool. A qualitatively good agreement between both sets of curves is obtained. Our work provides a path forward to overcome the memory and computational limits of DFT, thus enabling the study of large-scale devices beyond current \textit{ab-initio} capabilities

Figures (4)

• Figure 1: Training (left) and testing (right) TiN- HfO$_2$- Ti/TiN VCM structures with the Hf and O atoms omitted for better visualization. A total of 20 slices with 1 Å thickness was extracted from the Random 1 and 2 devices with arbitrary vacancy distributions to train the ML model. Another unseen 1-Å slice from Random 2 was used for validation. The trained model was then tested on unseen samples with either two fully-formed (Filament 1 and 2) or two broken (Broken 1 and 2) filament configurations.
• Figure 2: Higher level overview of the equivariant graph neural network architecture used in this work. It is a strictly local network consisting of a single message passing layer, with node (edge) embeddings $n_i$ ($e_{ji}$) representing atoms (interactions between atoms). The nodes aggregate messages from each other, weighted by an attention layer, with the output embeddings being used to update the edges. The updated embeddings $n_i'$ and $e_{ji}'$ are then fed into an output head to reconstruct the Hamiltonian sub-blocks.
• Figure 3: General overview of the workflow to train and test our machine learning model. The VCM cells considered (top left) consist of a HfO$_2$ layer (blue: Hf atoms; red: O atoms) with TiN and Ti/TiN electrodes (gray: Ti; dark blue: N). The device dimensions are set 93.4$\times$26.2$\times$26.3 Å$^3$ along the $x$, $y$, and $z$ directions. Slices in the $x$-$z$ plane (yellow) are used to train our EGNN (middle). A so-called augmented partitioning method (bottom left) xia2025 is leveraged to ensure that the atomic connectivity is preserved at the slice boundaries. The trained model is then used to directly infer the Hamiltonian matrix $H_{pred}$ of test structures generated with KMC at different times of a full "I-V" sweep. Finally, $\mathbf{H}_{pred}$ serves as input to quantum transport calculations. The transmission function $T(E)$ and electrical current $I_d$ are the final outcomes.
• Figure 4: Prediction results for the four TiN- HfO$_2$- Ti/TiN VCM test structures from Fig. \ref{['fig:structures']}. Each of them contains two sets of data. Left: Violin plot of the error distribution for the predicted Hamiltonian matrix $\mathbf{H}_{pred}$. The 5th and 95th percentile values of $\epsilon_{n}$ (node error) and $\epsilon_{e}$ (edge error) are indicated using red and green dashed lines, respectively. Note that for the "Broken 1" case, the 5th percentile lies outside of the plotted range. Right: Comparison between the energy-resolved transmission function $T(E)$ as obtained with a predicted Hamiltonian ($\mathbf{H}_{pred}$, solid red curves) with a reference DFT Hamiltonian ($\mathbf{H}_{DFT}$, dashed blue curves) for the Filament 1, Filament 2, Broken 1, and Broken 2 TiN- HfO$_2$- Ti/TiN VCM configurations from Fig. \ref{['fig:structures']}. A bias of 1 V is applied between both metallic electrodes of the devices. The corresponding Fermi window at room temperature is delimited by the black dashed lines. The apparent HfO$_2$ band gap is indicated by the double arrows.