fit: adaptive low-rank initialization for variational (MPO·MPO) contraction

[SamuelBadr]

Problem

ContractOptions::fit() (and the TreeTN contract_fit) always started the
variational state with a topology-preserving zip-up contraction of A·B.
Without an explicit truncation policy that initializer materializes the exact
product — output bonds up to χ_A·χ_B per edge, plus intermediates of the
same scale — before any sweep ever runs. For large QTT/MPO bonds this

1. forces a max_bond_dim on the user just to keep initialization from
   blowing up, and
2. defeats the purpose of a variational (compressing) fit: the product bond
   is constructed first and only then refined.

The requested semantics are C₀ small → two-site FIT sweeps + tolerance → C ≈ AB
with adaptive rank growth and no required rank cap, and never an
explicit χ_A·χ_B product bond as an intermediate.

What I found in the current implementation

• The FIT sweeps themselves never construct a χ_A·χ_B product bond.
  Environments env[(from,to)] carry the A-, B- and C-link indices separately
  (χ_A × χ_B × χ_C) (explicitly allowed). In the two-site update every A/B
  bond is contracted away into the environments before factorization, so the
  local SVD matrix sides are built purely from C-bonds and output sites — no
  χ_A·χ_B term appears in the factorization.
• Adaptive rank growth already works in the sweeps: with
  max_bond_dim = None and an svd_policy, the bond cap is None, so ranks
  grow freely per the tolerance; with a cap the cap wins; with neither, existing
  ranks are preserved.
• The zip-up initializer is the sole offender. At each edge it builds R
  intermediates whose element count scales as rank · χ_A · χ_B and, with no
  truncation, emits an output bond of dimension up to χ_A·χ_B; the root node
  materializes the full exact contraction. So the product bond is constructed
  a priori exactly when the user leaves max_bond_dim unset.

This matches upstream behavior: FastMPOContractions.jl defaults to zip-up init
(init === nothing), while ITensorMPS.jl's fit starts from a small random
output-site MPS and grows with maxdim = typemax(Int) + cutoff.

Design chosen

1. Split initialization from the sweep engine in tensor4all-treetn:
   • contract_fit_from_initial(tn_a, tn_b, center, options, tn_c) — the
     generic two-site sweep engine over a caller-provided C₀ (the
     initialization, zip-up or otherwise, is now just how you get C₀).
   • low_rank_initializer_tree_tn(tn_a, tn_b, center, bond_dim, seed) —
     builds a deterministic small-rank random C₀ carrying the surviving
     output site indices, with every bond at bond_dim (default 1). It never
     contracts A and B, so no χ_A·χ_B-sized tensor is ever formed. Concrete on
     IdxTensor (f64/Complex64), which is where seeded random construction
     lives (the generic tensor layer is scalar-erased).
2. New public default at the tensor4all-itensorlike layer:
   • FitInitializer { ZipUp, LowRankRandom { bond_dim, seed } } added to
     ContractOptions (with_initializer / initializer()).
   • ContractOptions::fit() defaults to
     LowRankRandom { bond_dim: 1, seed: None } (deterministic seed). So
     fit() + with_svd_policy(tol) and max_bond_dim = None now genuinely
     means unconstrained adaptive rank growth from a small start, with no
     product-bond intermediate.
   • FitInitializer::ZipUp restores the previous exact-start behavior.
   • The generic treetn::contraction::contract(…, Fit) dispatcher is
     unchanged (still zip-up init), because a generic random C₀ cannot be built
     for an arbitrary TensorLike; callers who need a custom start use
     contract_fit_from_initial.

Alternatives considered (and why rejected)

• Keep zip-up as the default, add low-rank as opt-in. Rejected: the
  task's core requirement is that adaptive FIT never constructs the product
  bond; leaving zip-up as the default would keep violating it for every default
  call. The repo is early-stage with no back-compat burden.
• Generic random construction in the tensor traits. Rejected: the tensor
  layer is scalar-erased (TensorLike has no associated scalar type), so a
  meaningful random tensor for arbitrary T is not expressible; IdxTensor
  is the right seam.
• Deterministic identity/product-like init. Works for correctness but is
  degenerate for products (symmetry traps); seeded random is the standard,
  reproducible choice (ITensorMPS-style) and is what we ship.
• max_bond_dim = None = "grow freely even without a tolerance"
  (one-site/two-site enrichment with no criterion). Rejected: there is no
  sensible growth criterion without a tolerance or cap; the exact-product
  behavior only survives under explicit ZipUp init.

Exact rank-growth semantics

• max_bond_dim = Some(cap) → every sweep bond capped at cap.
• max_bond_dim = None + svd_policy → unconstrained adaptive growth per the
  tolerance (no user rank cap required).
• max_bond_dim = None and no policy → existing (initializer-provided) ranks
  are preserved.
• nfullsweeps = 0 → returns the initializer unchanged.

Does any χ_A·χ_B-sized intermediate remain?

• No χ_A·χ_B output/product bond is created anywhere in the
  low-rank-initialized path: not in initialization (bonds start at 1) and not
  in the sweeps (A/B bonds only appear inside environments as separate
  (χ_A, χ_B, χ_C) indices, which is explicitly allowed).
• χ_A·χ_B transient contraction cost still appears in environment
  construction (env of shape χ_A×χ_B×χ_C is contracted with the local
  tensors) — this is intrinsic to any two-site tree fit and is cost, not a
  materialized product bond.
• Under explicit FitInitializer::ZipUp, zip-up still forms the exact product
  bond — by explicit opt-in only.

Tests

• test_low_rank_initializer_stays_small_no_product_bond — structural
  regression: with χ_A = 6, χ_B = 7 the initializer must keep every bond at
  bond_dim (1/3); would fail if any χ_A·χ_B bond were created.
• test_low_rank_initializer_carries_surviving_sites / _is_deterministic /
  _requires_bond_dim_at_least_one.
• test_contract_fit_from_rank1_grows_adaptively — rank-1 C₀ + tolerance,
  max_bond_dim = None → bond grows beyond 1 and matches the exact product.
• adaptive_fit_grows_rank_without_cap, adaptive_fit_tolerance_controls_rank_and_error,
  adaptive_fit_avoids_product_bond_intermediate (χ_A·χ_B = 144, compressible
  to rank 12), adaptive_fit_respects_explicit_cap.
• Correctness: chains (lengths 2–5, phys 2–3, bonds 2–3), a branched star
  tree, and Complex64, all matching the exact product to ≤ 1e-6.
• Existing zip-up-based tests (fit_bond_capping, bug_fit_elementwise,
  test_contract_fit_nhalfsweeps_zero_ok) were pinned to
  FitInitializer::ZipUp to preserve their original intent under the new
  default.

Benchmark

benchmarks/results/2026-08-19-adaptive-fit-lowrank-initializer.md:
χ_A = χ_B = 24–48, nsweeps = 3, tolerance 1e-8, compressible product (B is
the identity padded to bond χ_B):

χ	zipup exact	fit + zipup init	fit + low-rank init
24	50 ms	191 ms	120 ms (err 2.4e-8)
32	75 ms	399 ms	227 ms
48	326 ms	2.43 s	1.32 s

Low-rank-initialized fit is ~1.6–1.8× faster at equal accuracy and never
constructs the χ² intermediate.

API compatibility

• New public API: treetn::{contract_fit_from_initial, low_rank_initializer_tree_tn}
  (+ re-exports), itensorlike::FitInitializer, ContractOptions::with_initializer
  / initializer(), DEFAULT_FIT_INITIALIZER_SEED.
• Behavior change: ContractOptions::fit() with no explicit initializer now
  starts from a deterministic rank-1 random state instead of the exact zip-up
  state. FitInitializer::ZipUp restores the old behavior. The repo's
  "early development, no backward compatibility" policy applies; the affected
  tests were updated deliberately.
• tensor4all-itensorlike's contract() drives Fit directly; the treetn
  generic contraction::contract(…, Fit) default is unchanged (zip-up init).
• C API: left untouched (no ABI change; downstream can reach the adaptive
  path through the same ContractOptions surface once exposed if desired).

Validation

• cargo fmt --all -- --check ✓
• cargo test --release --workspace (excluding tensor4all-hdf5 and
  docs/book-tests, which need a system HDF5 that is absent in this
  environment — a pre-existing limitation) → 3789 passed, 0 failed.
• cargo test --release --doc for changed crates ✓ (treetn 121, itensorlike 28).
• cargo clippy on the changed crates: no warnings in changed files. Note:
  cargo clippy --workspace -- -D warnings currently fails on pristine
  main too (45 pre-existing nonminimal_bool lints in tensor4all-core
  under rustc 1.96.1), unrelated to this PR.
• python3 scripts/repository-rules-review.py --base origin/main --worktree --dry-run → pass, no findings.

Remaining limitations / future work

• The low-rank initializer is concrete on IdxTensor (f64/Complex64); other
  dtypes fall back to the zip-up/Provided paths. A generic random
  construction would require an associated scalar type on TensorLike.
• max_bond_dim = None + no tolerance preserves the initial rank; a caller who
  wants compression must supply an svd_policy. This is documented.
• Rank-revealing / AMEn-style enrichment and a randomized (SRC-style) variant
  remain future work (see Feature request: successive randomized compression (SRC) for MPO-MPS / MPO-MPO contraction #563).

[github-actions]

Repository rules review

Repository rules review (9e9aedaebe0d3918b34dd399ff0981e337f3835b...06c0ea86544f01d3c286e73190432d123420b28f)
Verdict: pass
Findings:
- [warn] llm-skipped (External LLM Review) <unknown>: External LLM review was skipped
  External LLM review is permanently disabled in this repository; deterministic rule checks run instead.

[shinaoka]

Is this still alive?