1import json
  2import time
  3import numpy as np
  4
  5
  6def tangent_projector(U, V, G):
  7    UU = U @ U.T
  8    VV = V @ V.T
  9    return UU @ G + G @ VV - UU @ G @ VV
 10
 11
 12def sketch_recondition(U, V, rng, oversampling=8, eps=1e-12):
 13    m, r = U.shape
 14    S = rng.normal(size=(r + oversampling, m)) / np.sqrt(r + oversampling)
 15    Y = S @ U
 16    C = Y.T @ Y + eps * np.eye(r)
 17    R = np.linalg.cholesky(C).T
 18    Ri = np.linalg.solve(R, np.eye(r))
 19    return U @ Ri, V @ R.T
 20
 21
 22def exact_recondition(U, V):
 23    Q, R = np.linalg.qr(U, mode="reduced")
 24    signs = np.where(np.diag(R) >= 0.0, 1.0, -1.0)
 25    D = np.diag(signs)
 26    return Q @ D, V @ R.T @ D
 27
 28
 29def metrics(U, V, target):
 30    W = U @ V.T
 31    return {
 32        "loss": float(0.5 * np.mean((W - target) ** 2)),
 33        "u_orth": float(np.linalg.norm(U.T @ U - np.eye(U.shape[1]))),
 34        "v_orth": float(np.linalg.norm(V.T @ V - np.eye(V.shape[1]))),
 35        "cond_u": float(np.linalg.cond(U)),
 36        "cond_v": float(np.linalg.cond(V)),
 37    }
 38
 39
 40def math_check(seed=7):
 41    rng = np.random.default_rng(seed)
 42    m, n, r = 48, 37, 6
 43    U0 = rng.normal(size=(m, r))
 44    V0 = rng.normal(size=(n, r))
 45    # Deliberately non-orthogonal basis, while retaining a well-defined product.
 46    T = np.diag(np.geomspace(1e-4, 1e4, r))
 47    U = U0 @ T
 48    V = V0 @ np.linalg.inv(T).T
 49    W = U @ V.T
 50    # Keep the exact same sketch used by the Cholesky factorization.
 51    S = rng.normal(size=(r + 8, m)) / np.sqrt(r + 8)
 52    Y = S @ U
 53    R = np.linalg.cholesky(Y.T @ Y + 1e-12 * np.eye(r)).T
 54    Us = U @ np.linalg.solve(R, np.eye(r))
 55    Vs = V @ R.T
 56    rel_invariance = np.linalg.norm(Us @ Vs.T - W) / np.linalg.norm(W)
 57    sketch_orth = np.linalg.norm((S @ Us).T @ (S @ Us) - np.eye(r))
 58    euclidean_orth = np.linalg.norm(Us.T @ Us - np.eye(r))
 59    # The displayed projector requires both factors to be Euclidean-orthonormal.
 60    Uq, _ = np.linalg.qr(rng.normal(size=(m, r)), mode="reduced")
 61    Vq, _ = np.linalg.qr(rng.normal(size=(n, r)), mode="reduced")
 62    G = rng.normal(size=(m, n))
 63    H = tangent_projector(Uq, Vq, G)
 64    residual = G - H
 65    tangent_id = np.linalg.norm(Uq.T @ residual) + np.linalg.norm(residual @ Vq)
 66    return {
 67        "factor_product_relative_error": float(rel_invariance),
 68        "sketch_gram_orth_error": float(sketch_orth),
 69        "euclidean_basis_orth_error": float(euclidean_orth),
 70        "projector_normal_residual": float(tangent_id),
 71        "projector_residual_norm": float(np.linalg.norm(residual)),
 72        "initial_u_cond": float(np.linalg.cond(U)),
 73        "reconditioned_u_cond": float(np.linalg.cond(Us)),
 74    }
 75
 76
 77def train(method, seed=123, steps=500, recondition_every=10):
 78    rng = np.random.default_rng(seed)
 79    m, n, r = 64, 52, 8
 80    # Rank-r target plus a perpendicular component, making the tangent projection meaningful.
 81    A = rng.normal(size=(m, r)); B = rng.normal(size=(n, r))
 82    target = A @ B.T / np.sqrt(r) + 0.15 * rng.normal(size=(m, n))
 83    U = rng.normal(size=(m, r))
 84    V = rng.normal(size=(n, r))
 85    # Same represented initial matrix across methods, with a well-scaled gauge.
 86    U, _ = np.linalg.qr(U, mode="reduced")
 87    V, _ = np.linalg.qr(V, mode="reduced")
 88    t0 = time.perf_counter()
 89    history = []
 90    for step in range(steps):
 91        W = U @ V.T
 92        G = (W - target) / (m * n)
 93        if method == "none":
 94            H = G
 95        else:
 96            # The idea's essential rule: standard projector, never an oblique sketch projector.
 97            Uq, Vq = exact_recondition(U, V)
 98            H = tangent_projector(Uq, Vq, G)
 99            # Factor update using the projected matrix, preserving a simple matched setup.
100            U, V = Uq, Vq
101        # A stable factor gradient for W=UV^T, with H as the dense tangent update.
102        dU = H @ V
103        dV = H.T @ U
104        lr = 0.03
105        U -= lr * dU
106        V -= lr * dV
107        if method == "sketch" and (step + 1) % recondition_every == 0:
108            U, V = sketch_recondition(U, V, rng, oversampling=8)
109        elif method == "qr" and (step + 1) % recondition_every == 0:
110            U, V = exact_recondition(U, V)
111        if step in (0, 9, 99, steps - 1):
112            history.append(metrics(U, V, target))
113    out = metrics(U, V, target)
114    out["history"] = history
115    out["seconds"] = time.perf_counter() - t0
116    return out
117
118
119def main():
120    result = {"math_check": math_check()}
121    for method in ("none", "qr", "sketch"):
122        result[method] = train(method)
123    print(json.dumps(result, indent=2))
124
125
126if __name__ == "__main__":
127    main()