Quantum TEA Library: Advanced aTTNs

• Quantum TEA Library is an open-source toolkit implementing augmented tree tensor networks (aTTNs) that integrate two-site disentanglers for enhanced entanglement capture.
• The library features advanced optimization routines, including SVD-based updates and mixed-device execution to efficiently manage computational resources.
• It enables precise ground state energy searches in complex lattice models, outperforming standard TTNs especially near quantum critical points.

An augmented tree tensor network (aTTN) is a tensor network ansatz that enhances the expressive power of a standard tree tensor network (TTN) by introducing a layer of two-site unitary disentanglers applied to the physical layer. This construction is designed to capture additional entanglement, particularly relevant for simulations of higher-dimensional quantum lattice models where the entanglement entropy grows with lattice size, even for states exhibiting area law scaling. The open-source Quantum TEA library provides a practical implementation of aTTN algorithms with detailed methodological guidance and performance analysis across prototypical many-body models (Reinić et al., 28 Jul 2025).

1. Definition and Structure of Augmented Tree Tensor Networks

An aTTN state is constructed by applying a layer of two-body unitary disentanglers to a conventional TTN. Formally, the ansatz is

$$|\psi_\mathrm{aTTN}\rangle = D(u)\,|\psi_\mathrm{TTN}\rangle$$

where $D(u) = \prod_{k} u_k$ denotes the product of disentangler unitaries $u_k$, each acting on a pair of physical sites. The TTN, $|\psi_\mathrm{TTN}\rangle$, arranges its tensors in a hierarchical tree structure. The addition of the disentangler layer enables encoding of more complex entanglement structures—beyond what is possible with standard TTNs—by locally rotating basis states at the bottom of the tree.

The function of the disentanglers is to "preprocess" entanglement on the physical sites such that the hierarchical coarse graining of the TTN is able to efficiently compress and propagate correlations in higher-dimensional settings. Each disentangler is variationally optimized to minimize the variational energy subject to the unitarity constraint ($u_k^\dagger u_k = \mathbb{1}$).

2. Algorithmic Implementation in the Quantum TEA Library

The Quantum TEA library includes comprehensive support for constructing, optimizing, and measuring observables in aTTNs. Architectural elements and functionality include:

• DELayer class: Maintains the position and content of the disentangler tensors. Disentanglers can be selectively placed on specific bonds as determined by the underlying lattice geometry.
• Environment and contraction routines: Functions (e.g., ATTN.get_environment(), ATTN.iteratively_contract_along_path()) enable efficient computation of the global environment tensor $\Gamma_k$ required for disentangler optimization at each position.
• Self-consistent and gradient-based update schemes: Disentanglers may be optimized via a MERA-like self-consistent loop exploiting the singular value decomposition (SVD) of the environment tensor or by using Riemannian gradient descent to respect the unitarity constraint.
• Mixed-device execution: Key tensor operations and linear algebra (notably, large tensor contractions and Lanczos iterations for DMRG) can be efficiently offloaded to GPUs, while ancillary tasks run on the CPU, optimizing both memory usage and computational throughput.
• Interactive tutorials: The library provides Jupyter Notebooks and pedagogical documentation, illustrating the process from ansatz construction to optimization and measurement.

3. Ground State Search: Variational Algorithm and Optimization

The ground state search follows a cyclical variational minimization of the energy

$$E = \langle \psi| \hat{H} | \psi \rangle, \quad \text{subject to} \quad \langle\psi|\psi\rangle = 1$$

for $|\psi\rangle = D(u)|\psi_{\mathrm{TTN}}\rangle$. The optimization alternates between the disentangler layer and the TTN tensors. The main steps are:

1. Disentangler Optimization: For each $u_k$, the cost function is expressed locally as

$E = \operatorname{Tr}(u_k\,\Gamma_k)$

where $\Gamma_k$ is the (environment) matrix obtained by contracting the full network with all tensors except $u_k$. The SVD of $\Gamma_k = U \Sigma V^\dagger$ yields the update

$u_k = - V U^\dagger$

which minimizes the local cost. The energy at the minimum is $E = -\sum_i \sigma_i$, where $\sigma_i$ are the singular values. In the gradient-based approach, the unitarity constraint is enforced using Riemannian optimization techniques.

1. Hamiltonian Dressing: The full Hamiltonian is updated via conjugation by the optimal disentangler layer,

$\hat{H}' = D(u) \hat{H} D(u)^\dagger$

1. TTN Optimization: A standard TTN variational sweep (e.g., DMRG update) is performed on $|\psi_{\mathrm{TTN}}\rangle$ with respect to the dressed Hamiltonian $\hat{H}'$.

The process is iterated until convergence. Alternating these updates ensures that both the local entanglement (through $u_k$) and nonlocal correlations (through the TTN hierarchy) are efficiently captured.

4. Observable Measurement and Contraction Strategies

Measurement of observables in aTTNs is adapted from TTN techniques but incorporates additional dressing by the disentanglers:

• For a local operator $\hat{O}$ acting on site $i$, one computes the reduced density matrix $\rho_i$ after "isometrizing" the TTN and contracts with the dressed operator

$\hat{O}' = D(u) \hat{O} D(u)^\dagger$

• For two-body or more general observables, the contraction can result in the observable's support spreading to multiple sites depending on the positioning of the disentanglers relative to the operator.

Notably, observables that are locally diagonal remain local or become two-body after dressing, while those originally nonlocal may be transformed to three- or four-body operators depending on the disentangler pattern.

5. Benchmarking: Performance and Comparative Analysis

The aTTN framework was benchmarked on large two-dimensional systems, up to $32 \times 32$ sites, for models including the square lattice quantum Ising model (Hamiltonian $$\hat{H} = - \sum_{\langle ij \rangle} \sigma_x^{(i)}\sigma_x^{(j)} - h\sum_i \sigma_z^{(i)}$$ at critical $h \approx 3$) and the triangular lattice Heisenberg model.

Key findings:

• For fixed TTN bond dimension $m$, aTTNs consistently yield ground state energies lower than those produced by standard TTNs, with the advantage most pronounced near quantum critical points where entanglement is maximal.
• The memory scaling of TTN/aTTN contractions follows $\mathcal{O}(m^3)$, while matrix product states (MPS) scale as $\mathcal{O}(m^2)$.
• aTTN's overhead, due to the need to contract the disentangler layer into the Hamiltonian before each TTN sweep, increases the prefactor of computational cost but preserves the same asymptotic scaling in $m$.
• For smaller system sizes or highly frustrated geometries (e.g., triangular lattice with many interaction terms), the benefit of aTTN versus TTN/MPS is reduced, especially when memory budgets are tight.
• Selective placement of disentanglers (e.g., only on bonds associated with Hamiltonian terms) allows further trade-offs of accuracy versus computational cost.

The benchmarking analyses included GPU memory consumption and runtime as a function of $m$, system size, and number of disentanglers.

6. Prospects and Extensions

The detailed implementation in the Quantum TEA library provides a publicly available, pedagogically structured platform for reproducing and extending the aTTN methodology. Future directions identified include:

• Extension of aTTN algorithms to time-evolution, allowing for real-time or imaginary-time dynamics in higher-dimensional systems.
• Compression strategies to reduce the memory overload resulting from dressing of Hamiltonian terms by disentanglers, potentially enabling even larger simulations or higher $m$.
• Distributed and multi-GPU deployments to scale simulations further.
• Application of aTTNs to models where area law entanglement and long-range correlations prohibit the use of simpler tensor network ansätze.

The methodological framework—including the layer-wise alternation between disentangler optimization and TTN sweeps, combined with contract-then-dress measurement protocols—provides a generalizable toolkit for systematically improving the expressive power and accuracy of tensor network simulations in higher dimensions.

7. Summary Table: aTTN vs. TTN/MPS Performance Characteristics

Feature/Regime	aTTN	TTN	MPS
Area law in 2D	Captured (via D(u))	Not naturally	Not naturally
Bond dimension scaling	$\mathcal{O}(m^3)$	$\mathcal{O}(m^3)$	$\mathcal{O}(m^2)$
Near-critical entanglement	Improved accuracy	Limited	Severely limited
Memory overhead	Higher prefactor	Baseline	Lower
Implementation in QTEA	Full (mixed device)	Full (reference)	Reference routines
Observable contraction	Needs "dressing"	Standard	Standard

This characterization highlights the regimes where aTTNs, as implemented in Quantum TEA, provide a notable increase in achievable accuracy—especially in critical or high-dimensional lattice problems—relative to standard tree-based and matrix product tensor network methods (Reinić et al., 28 Jul 2025).