Skip to content

Commit bd78e23

Browse files
CopilotrobertodrPanadestein
authored
feat: ✨ oversample SRC sketches before exact SVD truncation (#63)
* Initial plan * feat: ✨ oversample SRC before exact SVD truncation Assisted-by: GitHub Copilot:unknown Co-authored-by: robertodr <3708689+robertodr@users.noreply.github.com> * fix: 🐛 release discarded SVD storage Assisted-by: GitHub Copilot:unknown Co-authored-by: robertodr <3708689+robertodr@users.noreply.github.com> * perf(sweep): ⚡️ reduce the final truncation to a small SVD after a QR Assisted-by: GitHub Copilot:claude-sonnet-5.5 * feat: ✨ make oversampling opt-in with a default of 1.0 Assisted-by: GitHub Copilot:claude-opus-5.5 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: robertodr <3708689+robertodr@users.noreply.github.com> Co-authored-by: Roberto Di Remigio Eikås <robertodr@users.noreply.github.com> Co-authored-by: Ramón L. Panadés-Barrueta <rpana92@gmail.com>
1 parent a996df3 commit bd78e23

11 files changed

Lines changed: 431 additions & 62 deletions

File tree

‎README.md‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,22 @@ Whether a train is an MPS or an MPO is inferred from the rank of its first site
3333

3434
**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, an exact SVD-based fallback is dispatched, with a warning.
3535

36+
### Oversampling and Truncation
37+
38+
By default, `apply`, `compress` and `src` sketch at `chi_out` and return
39+
**right-canonical** output. For better accuracy at the same `chi_out`, pass
40+
`oversample`: a finite factor of at least `1.0` sketches at
41+
`ceil(oversample * chi_out)`, and `None` at the paper's
42+
`max(ceil(1.5 * chi_out), chi_out + 10)`. A final left-to-right SVD sweep then
43+
truncates to `chi_out`, and the result is **left-canonical**, with the norm on
44+
the last site. When positive, `cutoff` discards Schmidt values below
45+
`cutoff * sigma_max` at each cut in this final pass. Two-site trains keep the
46+
exact, right-canonical fallback and ignore both options.
47+
48+
Larger sketches improve accuracy but increase time and memory costs, often
49+
several-fold; compare factors at the same `chi_out` when choosing an
50+
accuracy/runtime trade-off.
51+
3652
### Tensor Indexing Conventions
3753

3854
The array layout follows the default `quimb` tensor indexing conventions, so results round-trip through [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) without any permutation:
@@ -81,7 +97,7 @@ bra = src([t.conj() for t in psi], A, B, chi_out=64)
8197

8298
The result then pairs with a ket by plain contraction, with no further conjugation.
8399

84-
**Cost.** The per-site cost of the sweep grows with `chi_out**2` times the product of the bond dimensions of the layers. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. See [the stack depth benchmarks](benches/stack/README.md) for measurements.
100+
**Cost.** The per-site cost of the sweep grows with the sketch size squared times the product of the bond dimensions of the layers. Compressing a whole stack at once pays off for shallow stacks of thin layers, such as two or three Trotter layers; apply anything else pairwise. See [the stack depth benchmarks](benches/stack/README.md) for measurements.
85101

86102
### GPU Support
87103

‎docs/content/docs/concepts/algorithm.mdx‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ $A$ and an MPS $|\psi\rangle$, without ever forming the exact product. The exact
99
product of an MPO with bond dimension $\chi_A$ and an MPS with bond dimension
1010
$\chi_\psi$ has bond dimension $\chi_A \chi_\psi$; compressing it afterwards, as
1111
in the usual contract-then-truncate approach, pays for that intermediate size.
12-
SRC targets the output bond dimension $\chi_{\text{out}}$ directly.
12+
SRC targets the output bond dimension $\chi_{\text{out}}$ directly, or an
13+
oversampled sketch size $\chi_{\text{os}}$ followed by a final SVD truncation to
14+
$\chi_{\text{out}}$.
1315

1416
## Randomized range finding
1517

@@ -22,12 +24,20 @@ probability, and the cost is dominated by the product $M\Omega$.
2224
SRC applies this idea one site at a time along the chain, with the matrix $M$
2325
being an unfolding of the (never formed) product network.
2426

25-
## Two sweeps
27+
## Sketching and truncation
28+
29+
By default, $\chi_{\text{os}} = \chi_{\text{out}}$ and only the first two sweeps
30+
run: the output is right-canonical, with bond dimension at most
31+
$\chi_{\text{out}}$. An `oversample` factor sets
32+
$\chi_{\text{os}} = \lceil \text{oversample}\,\chi_{\text{out}}\rceil$, and
33+
`oversample=None` the paper's
34+
$\max(\lceil 1.5\chi_{\text{out}}\rceil, \chi_{\text{out}} + 10)$; the third
35+
sweep then truncates to $\chi_{\text{out}}$.
2636

2737
1. **Left to right: sketching.** At every site the open physical legs are
2838
contracted with a Gaussian tensor $\Omega_j$ of sketch size
29-
$\chi_{\text{out}}$. Accumulating these from the left gives the sketched
30-
environments $C_j$, each of size $\chi_{\text{out}}$ times the product of the
39+
$\chi_{\text{os}}$. Accumulating these from the left gives the sketched
40+
environments $C_j$, each of size $\chi_{\text{os}}$ times the product of the
3141
input bond dimensions. The same sketch index is shared by every site, which
3242
makes the sketch a Khatri-Rao product rather than a dense Gaussian matrix of
3343
the full Hilbert space.
@@ -37,16 +47,20 @@ being an unfolding of the (never formed) product network.
3747
tensor as an isometry. The isometry is then folded into $S$ for the next site.
3848
The first site absorbs the remaining weight.
3949

40-
The output is therefore in right-canonical form, with bond dimension at most
41-
$\chi_{\text{out}}$. Without a [cutoff](/features/cutoff) every inner bond is
42-
$\chi_{\text{out}}$ unless the physical dimensions near the ends of the chain
43-
force it lower.
50+
3. **Left to right: exact truncation (oversampling only).** The intermediate is
51+
already right-canonical, so no re-orthogonalisation is needed. An SVD of each
52+
centre tensor gives the Schmidt values at that cut. Keep at most
53+
$\chi_{\text{out}}$ values, apply the relative [cutoff](/features/cutoff), and
54+
absorb the remaining weight into the next site. The output is
55+
**left-canonical**, with the norm on the last site. The discarded weights add
56+
in quadrature, giving the usual TT-SVD quasi-optimality bound for this
57+
intermediate.
4458

4559
## Cost
4660

4761
For an MPO with bond dimension $\chi_A$ applied to an MPS with bond dimension
4862
$\chi_\psi$, the dominant per-site cost grows like
49-
$\chi_{\text{out}}^2 \chi_A \chi_\psi$, compared with the $(\chi_A\chi_\psi)^3$ of
63+
$\chi_{\text{os}}^2 \chi_A \chi_\psi$, compared with the $(\chi_A\chi_\psi)^3$ of
5064
an exact contraction followed by an SVD-based truncation. The same kernel handles
5165
MPO-MPO products and plain compression; for a [stack](/features/stacks) of trains,
5266
the product of all their bond dimensions replaces $\chi_A\chi_\psi$.
Lines changed: 38 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
---
22
title: Adaptive truncation
3-
description: Trim bonds below chi_out with a relative singular-value cutoff.
3+
description: Trim bonds below chi_out with a relative cutoff, and oversample the sketch.
44
---
55

6-
By default every inner bond of the output has dimension `chi_out`, even if the
7-
product has lower rank across that cut. Set `cutoff` to trim each bond to its
8-
effective rank:
6+
`chi_out` is the maximum output bond dimension. Set `cutoff` to trim each bond
7+
further using a relative singular-value threshold:
98

109
```python
1110
import quimb.tensor as qtn
@@ -22,12 +21,41 @@ print([a.shape[0] for a in fixed[1:]])
2221
print([a.shape[0] for a in trimmed[1:]])
2322
```
2423

25-
During the right-to-left sweep, the QR factorisation at each site is followed by
26-
an SVD of its small `(chi_out, chi_out)` triangular factor, and singular values
27-
below `cutoff * sigma_max` are discarded. The extra cost is negligible next to the
28-
contractions.
24+
By default, the sketch size is `chi_out`. During the right-to-left sweep, the QR
25+
factorisation at each site is followed by an SVD of its small triangular factor,
26+
and singular values below `cutoff * sigma_max` are discarded. The output is
27+
**right-canonical**.
28+
29+
## Oversampling
30+
31+
Sketching at exactly `chi_out` lets the error of the randomized range finder
32+
dominate. Set `oversample` to a finite factor of at least `1.0` to sketch at
33+
`ceil(oversample * chi_out)`, or to `None` for the paper's recommendation,
34+
`max(ceil(1.5 * chi_out), chi_out + 10)`:
35+
36+
```python
37+
import quimb.tensor as qtn
38+
39+
from src_method import apply
40+
41+
H = qtn.MPO_rand(8, bond_dim=4, phys_dim=2, seed=0)
42+
psi = qtn.MPS_rand_state(8, bond_dim=16, seed=0)
43+
44+
out = apply(H.arrays, psi.arrays, chi_out=32, oversample=None, seed=1)
45+
```
46+
47+
A final left-to-right SVD sweep then truncates the oversampled, right-canonical
48+
intermediate to `chi_out`, discarding singular values below `cutoff * sigma_max`
49+
at each cut. These are the Schmidt values of the intermediate train as truncated
50+
so far, not of the random sketch matrix. The result is **left-canonical**, with
51+
the norm on the last site.
52+
53+
Oversampling improves accuracy at a fixed `chi_out` but costs more: the sweep
54+
runs at the larger sketch size, `cutoff` no longer trims bonds during the sweep,
55+
and the final pass adds one SVD per site. With `oversample=None`, expect a run
56+
time several times that of the default.
2957

3058
- `cutoff` must lie in `[0.0, 1.0)`; `0.0` (the default) disables trimming.
3159
- `chi_out` remains the upper bound: the cutoff only ever lowers a bond.
32-
- Two-site trains use the exact fallback, which ignores `cutoff` and truncates to
33-
`chi_out`.
60+
- Two-site trains use the exact fallback, which ignores `cutoff` and `oversample`,
61+
truncates to `chi_out` and returns right-canonical output.

‎docs/content/docs/features/stacks.mdx‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,13 @@ form.
2929

3030
## When a single sweep pays off
3131

32-
The per-site cost of the sweep grows with $\chi_{\text{out}}^2$ times the
32+
The per-site cost of the sweep grows with the sketch size squared times the
3333
*product* of the bond dimensions of all layers, so it grows exponentially with the
3434
depth of the stack. In exchange, a single sweep truncates once instead of once per
3535
layer. In practice that buys at most a moderate gain in accuracy: the error of
3636
the randomized sketch dominates the error compounded across layers.
37+
[Oversampling](/features/cutoff#oversampling) reduces that error for both the
38+
one-shot and pairwise approaches.
3739

3840
Compressing a whole stack at once pays off for shallow stacks of thin layers, such
3941
as two or three Trotter layers; apply anything else pairwise. The [stack

‎src/src_method/_sweep.py‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
release_cached_blocks,
5050
to_numpy,
5151
truncated_qr,
52+
use_device,
5253
)
5354

5455
if TYPE_CHECKING:
@@ -164,6 +165,7 @@ def sweep(
164165
*,
165166
device: str = "cpu",
166167
cutoff: float = 0.0,
168+
truncate_to: int | None = None,
167169
dtype: DTypeLike,
168170
resources: Resources | None = None,
169171
) -> list[np.ndarray]:
@@ -175,18 +177,21 @@ def sweep(
175177
array-likes with ``shape``, ``dtype``, ``ndim`` and ``np.asarray``
176178
support; each is read only when the sweep reaches it.
177179
kind: The kind of the contracted train.
178-
chi_out: The sketch size, which is the maximum output bond dimension.
180+
chi_out: The sketch size, which bounds the intermediate bond dimension.
179181
prng: The generator for the Gaussian sketches, always host-side so that a
180182
seed gives the same draws on every device.
181183
device: ``"cpu"`` or ``"gpu"``, as in `src_method.stack.src`.
182184
cutoff: Relative singular-value cutoff for adaptive bond truncation.
185+
truncate_to: If set, truncate the right-canonical intermediate to this
186+
bond dimension with a final left-to-right SVD pass, applying
187+
``cutoff`` there instead of to the sketch.
183188
dtype: The data type of the sketches.
184189
resources: Memory budgets, scratch space and devices; detected when
185190
``None``.
186191
187192
Returns:
188-
The site arrays of the compressed train in right-canonical form, as numpy
189-
arrays.
193+
The site arrays of the compressed train as numpy arrays, right-canonical
194+
unless ``truncate_to`` is set, in which case they are left-canonical.
190195
"""
191196
shapes = padded_shapes(layers)
192197
# The dtype of the environments and the output: the sketches promoted by the
@@ -220,7 +225,7 @@ def sweep(
220225
budgets,
221226
shapes,
222227
chi_out=chi_out,
223-
cutoff=cutoff,
228+
cutoff=cutoff if truncate_to is None else 0.0,
224229
sketches=_Sketches(prng, chi_out, dtype, group.size),
225230
work=work,
226231
depth=len(layers),
@@ -230,9 +235,44 @@ def sweep(
230235
) as source:
231236
eta = group.run(lambda rank: sweep_.run(rank, source))[0]
232237
assert eta is not None # noqa: S101 (rank 0 returns the cores)
238+
if truncate_to is not None:
239+
xp = group.xp
240+
with (
241+
use_device(xp, group.devices[0]),
242+
device_pool_limit(xp, budgets.device_cap),
243+
):
244+
eta = _truncate(eta, truncate_to, cutoff, group)
233245
return [to_numpy(site) for site in unpad(eta, kind)]
234246

235247

248+
def _truncate(
249+
sites: list[np.ndarray], chi_out: int, cutoff: float, group: DeviceGroup
250+
) -> list[np.ndarray]:
251+
"""Truncate a padded right-canonical train, staging one centre at a time."""
252+
xp = group.xp
253+
centre = xp.asarray(sites[0])
254+
for j in range(len(sites) - 1):
255+
left, right, up, down = centre.shape
256+
matrix = centre.transpose(0, 2, 3, 1).reshape(left * up * down, right)
257+
# A tall matrix is reduced by QR first, so the SVD sees a small square. It
258+
# runs on the host, as in `truncated_qr`: device SVDs are slow at this size.
259+
Q, R = xp.linalg.qr(matrix) if matrix.shape[0] > right else (None, matrix)
260+
U, S, Vh = np.linalg.svd(to_numpy(R), full_matrices=False)
261+
rank = min(chi_out, S.size)
262+
if cutoff > 0:
263+
rank = min(rank, max(1, int((cutoff * S[0] <= S).sum())))
264+
U = xp.asarray(U[:, :rank])
265+
U = U if Q is None else Q @ U
266+
sites[j] = to_numpy(
267+
xp.ascontiguousarray(U.reshape(left, up, down, rank).transpose(0, 3, 1, 2))
268+
)
269+
centre = xp.tensordot(
270+
xp.asarray(S[:rank, None] * Vh[:rank]), xp.asarray(sites[j + 1]), axes=1
271+
)
272+
sites[-1] = to_numpy(centre)
273+
return sites
274+
275+
236276
class _Sketches:
237277
"""The Gaussian tensors of the left-to-right pass, shared by the devices.
238278

‎src/src_method/_validation.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from __future__ import annotations
88

99
import numbers
10+
from math import ceil, isfinite
1011
from typing import TYPE_CHECKING
1112

1213
import numpy as np
@@ -16,7 +17,39 @@
1617

1718
from ._tensor_train import Site
1819

19-
__all__ = ["validate_chi_out", "validate_cutoff", "validate_open_boundary"]
20+
__all__ = [
21+
"sketch_size",
22+
"validate_chi_out",
23+
"validate_cutoff",
24+
"validate_open_boundary",
25+
]
26+
27+
28+
def sketch_size(chi_out: int, oversample: float | None) -> int:
29+
"""Resolve the sketch size from an optional multiplicative factor.
30+
31+
Args:
32+
chi_out: The validated output bond dimension.
33+
oversample: A finite factor of at least one, or ``None`` for the paper's
34+
recommendation.
35+
36+
Returns:
37+
The oversampled bond dimension.
38+
39+
Raises:
40+
ValueError: If ``oversample`` is not a finite real number of at least one.
41+
"""
42+
if oversample is None:
43+
return max((3 * chi_out + 1) // 2, chi_out + 10)
44+
if (
45+
isinstance(oversample, bool)
46+
or not isinstance(oversample, numbers.Real)
47+
or not isfinite(oversample)
48+
or oversample < 1.0
49+
):
50+
msg = f"oversample must be a finite factor >= 1.0, got {oversample!r}."
51+
raise ValueError(msg)
52+
return ceil(oversample * chi_out)
2053

2154

2255
def validate_chi_out(chi_out: object) -> int:

‎src/src_method/apply.py‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def apply(
2929
chi_out: int | np.integer,
3030
*,
3131
cutoff: float = 0.0,
32+
oversample: float | None = 1.0,
3233
dtype: DTypeLike | None = None,
3334
seed: int | None = None,
3435
device: str = "cpu",
@@ -49,11 +50,12 @@ def apply(
4950
right_tensor: The site arrays of the right tensor network (MPO or MPS).
5051
chi_out: The desired maximum bond dimension of the output tensor network.
5152
cutoff: Relative singular-value cutoff for adaptive bond truncation.
52-
When positive, bonds are trimmed to their effective rank by
53-
discarding singular values below ``cutoff * sigma_max`` at each
54-
site during the right-to-left sweep. The SVD operates on the
55-
small ``(chi_out, chi_out)`` R factor from QR, so overhead is
56-
minimal. Set to 0.0 (default) to keep all bonds at chi_out.
53+
Discards singular values below ``cutoff * sigma_max``; see
54+
`src_method.stack.src`. Set to 0.0 (default) to disable cutoff
55+
trimming. Ignored for two sites.
56+
oversample: Multiplicative sketch-size factor, finite and at least 1.0,
57+
or ``None`` for the paper's recommendation; see
58+
`src_method.stack.src`. Defaults to 1.0, no oversampling.
5759
dtype: Data type of the random sketches. Defaults to the promoted
5860
floating dtype of the inputs; an explicit dtype can promote the result.
5961
seed: An optional seed for the random number generator.
@@ -62,13 +64,16 @@ def apply(
6264
resources: Memory budgets and scratch space; see `src_method.stack.src`.
6365
6466
Returns:
65-
The site arrays of the compressed tensor network (MPS or MPO).
67+
The site arrays of the compressed tensor network (MPS or MPO), in
68+
right-canonical form, or left-canonical form when oversampling (three
69+
sites or more).
6670
6771
Raises:
6872
TypeError: If ``chi_out`` is not an integer or the combination of input
6973
tensor types is unsupported.
7074
ValueError: If ``chi_out`` is not positive, if ``cutoff`` is not
71-
in ``[0.0, 1.0)``, if a train is not open-boundary, if the two trains
75+
in ``[0.0, 1.0)``, if ``oversample`` is not finite and at least 1.0,
76+
if a train is not open-boundary, if the two trains
7277
differ in length or in the physical dimensions they join, if a
7378
sub-three-site train is not exactly two sites, or if ``device`` is not
7479
recognised.
@@ -88,6 +93,7 @@ def apply(
8893
right_tensor,
8994
chi_out=chi_out,
9095
cutoff=cutoff,
96+
oversample=oversample,
9197
dtype=dtype,
9298
seed=seed,
9399
device=device,

0 commit comments

Comments
 (0)