File size: 2,228 Bytes
49dddf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # Copyright 2026 Angshul Majumdar. Licensed under the Apache License, Version 2.0.
"""Rank-one CP approximation by higher-order power iteration (HOPM)."""
import torch
_L = "abcdefgh"
def contract_except(T, us, n):
N = T.dim()
idx = _L[:N]
spec = idx + "," + ",".join(idx[m] for m in range(N) if m != n) + "->" + idx[n]
return torch.einsum(spec, T, *[us[m] for m in range(N) if m != n])
def cp_value(T, us):
idx = _L[: T.dim()]
return torch.einsum(idx + "," + ",".join(idx) + "->", T, *us)
def outer(us):
out = us[0]
for u in us[1:]:
out = out[..., None] * u
return out
@torch.no_grad()
def rank1_cp(T, iters=30, tol=1e-6):
"""Unit factors u_1..u_N (approximately) maximising sigma = <T, u_1 o ... o u_N>.
T has shape (leading..., d_in, d_out) with one or two leading modes. Initialised from the top
singular pair of the largest (d_in x d_out) slice (deterministic), then HOPM sweeps."""
N = T.dim()
lead_shape = T.shape[:-2]
slices = T.reshape(-1, T.shape[-2], T.shape[-1])
j = int(slices.flatten(1).norm(dim=1).argmax())
M = slices[j]
small = M @ M.T if M.shape[0] <= M.shape[1] else M.T @ M # exact top pair via Gram eigvec
_, vecs = torch.linalg.eigh(small.double())
v = vecs[:, -1].to(T.dtype)
if M.shape[0] <= M.shape[1]:
u_in, u_out = v, M.T @ v
else:
u_in, u_out = M @ v, v
u_in, u_out = u_in / u_in.norm().clamp_min(1e-30), u_out / u_out.norm().clamp_min(1e-30)
core = torch.einsum("sij,i,j->s", slices, u_in, u_out).reshape(lead_shape)
if core.dim() == 1:
us = [core / core.norm().clamp_min(1e-30)]
else:
Uc, _, Vc = torch.linalg.svd(core, full_matrices=False)
us = [Uc[:, 0], Vc[0]]
us = us + [u_in, u_out]
sigma = cp_value(T, us)
for _ in range(iters):
for n in range(N):
w = contract_except(T, us, n)
us[n] = w / w.norm().clamp_min(1e-30)
new = cp_value(T, us)
done = abs(float(new) - float(sigma)) <= tol * max(abs(float(new)), 1e-30)
sigma = new
if done:
break
if sigma < 0:
us[-1] = -us[-1]
sigma = -sigma
return us, float(sigma)
|