1import json
  2import time
  3import numpy as np
  4
  5SEED = 194
  6rng = np.random.default_rng(SEED)
  7
  8
  9def ns_transform(M, depth):
 10    """Finite Newton-Schulz polar-like transform, scaled back to M's units."""
 11    m, n = M.shape
 12    alpha = max(1e-12, np.linalg.norm(M, 2))
 13    X = M / alpha
 14    if m >= n:
 15        I = np.eye(n)
 16        for _ in range(depth):
 17            X = 0.5 * X @ (3.0 * I - X.T @ X)
 18    else:
 19        I = np.eye(m)
 20        for _ in range(depth):
 21            X = 0.5 * (3.0 * I - X @ X.T) @ X
 22    return alpha * X
 23
 24
 25def exact_polar(M):
 26    U, _, Vt = np.linalg.svd(M, full_matrices=False)
 27    return U @ Vt
 28
 29
 30def residual(M):
 31    alpha = max(1e-12, np.linalg.norm(M, 2))
 32    X = M / alpha
 33    m, n = X.shape
 34    if m >= n:
 35        R = X.T @ X - np.eye(n)
 36    else:
 37        R = X @ X.T - np.eye(m)
 38    return np.linalg.norm(R, 'fro') / np.sqrt(min(m, n))
 39
 40
 41def duality_check():
 42    # Numerically test sup_{||X||op <= 1} <S,X> = ||S||_* using polar(S),
 43    # plus random feasible competitors that should not exceed the optimum.
 44    r = np.random.default_rng(SEED + 7)
 45    S = r.normal(size=(9, 6))
 46    P = exact_polar(S)
 47    optimum = float(np.sum(np.linalg.svd(S, compute_uv=False)))
 48    achieved = float(np.sum(S * P))
 49    competitors = []
 50    for _ in range(200):
 51        Q = r.normal(size=S.shape)
 52        qnorm = np.linalg.norm(Q, 2)
 53        competitors.append(float(np.sum(S * (Q / max(qnorm, 1e-12)))))
 54    return {'nuclear_norm': optimum, 'polar_objective': achieved,
 55            'absolute_gap': abs(achieved - optimum),
 56            'max_random_feasible_objective': max(competitors),
 57            'random_excess_over_optimum': max(competitors) - optimum}
 58
 59
 60def scalar_check():
 61    # Includes a zero, a very small singular value, and a well-scaled one.
 62    s = np.array([0.0, 1e-4, 0.03, 0.2, 0.7, 0.95])
 63    vals = []
 64    x = s.copy()
 65    for t in range(7):
 66        vals.append(x.copy())
 67        x = 0.5 * x * (3.0 - x * x)
 68    vals = np.asarray(vals)
 69    err = np.abs(1.0 - vals[:, 1:])
 70    # For a direction starting below one, normalized error contracts
 71    # rapidly once it is in the attraction basin.
 72    tail = err[1:, -1]
 73    contraction = (tail[1:] / np.maximum(tail[:-1], 1e-30)).tolist()
 74    # A finite polynomial maps zero continuously to zero, unlike polar's
 75    # nonzero-singular-value response (the intended smoothing observation).
 76    smooth_gap_at_small = float(vals[2, 1])
 77    return {
 78        'values': vals.tolist(),
 79        'small_sigma_response_t2': smooth_gap_at_small,
 80        'zero_response_all_t': vals[:, 0].tolist(),
 81        'largest_direction_errors': tail.tolist(),
 82        'successive_error_ratios': contraction,
 83        't7_small_sigma_error': float(abs(1.0 - vals[7-1, 1])),
 84    }
 85
 86
 87def make_problem(seed, m=32, n=24):
 88    r = np.random.default_rng(seed)
 89    # Ill-conditioned target: broad singular spectrum makes shallow NS visibly
 90    # different from exact polar while remaining a tiny CPU problem.
 91    U, _ = np.linalg.qr(r.normal(size=(m, n)))
 92    V, _ = np.linalg.qr(r.normal(size=(n, n)))
 93    singulars = np.geomspace(8.0, 0.02, n)
 94    A = U @ np.diag(singulars) @ V.T
 95    # A diagonal preconditioner creates nontrivial, changing momentum matrices.
 96    weights = np.geomspace(0.4, 2.5, m)[:, None] * np.geomspace(0.7, 1.8, n)[None, :]
 97    return A, weights
 98
 99
100def run_optimizer(kind, seed, steps=180):
101    A, weights = make_problem(seed)
102    r = np.random.default_rng(seed + 1000)
103    W = 0.35 * r.normal(size=A.shape)
104    M = np.zeros_like(W)
105    beta, lr = 0.90, 0.055
106    losses, residuals, smoothness = [], [], []
107    prev_dir, prev_M = None, None
108    ns_mults = 0
109    t0 = time.perf_counter()
110    for step in range(steps):
111        grad = weights * (W - A)
112        M = beta * M + (1.0 - beta) * grad
113        if kind == 'exact':
114            direction = exact_polar(M)
115            orth_res = 0.0
116        elif kind == 't5':
117            direction = ns_transform(M, 5)
118            orth_res = residual(direction)
119            ns_mults += 5
120        elif kind == 't2':
121            direction = ns_transform(M, 2)
122            orth_res = residual(direction)
123            ns_mults += 2
124        elif kind == 'adaptive':
125            depth = 2 if step < int(0.7 * steps) else 4
126            direction = ns_transform(M, depth)
127            orth_res = residual(direction)
128            ns_mults += depth
129        else:
130            raise ValueError(kind)
131        # Muon-style update, with a common scalar normalization so comparisons
132        # reflect spectral direction rather than arbitrary matrix magnitude.
133        direction = direction * (np.linalg.norm(M, 'fro') / max(np.linalg.norm(direction, 'fro'), 1e-12))
134        if prev_dir is not None:
135            smoothness.append(np.linalg.norm(direction - prev_dir, 'fro') /
136                              max(np.linalg.norm(M - prev_M, 'fro'), 1e-12))
137        W -= lr * direction
138        losses.append(float(0.5 * np.mean(weights * (W - A) ** 2)))
139        residuals.append(float(orth_res))
140        prev_dir, prev_M = direction.copy(), M.copy()
141    elapsed = time.perf_counter() - t0
142    return {
143        'final_loss': losses[-1],
144        'best_loss': min(losses),
145        'loss_at_60': losses[59],
146        'loss_at_120': losses[119],
147        'mean_update_smoothness': float(np.mean(smoothness)),
148        'median_update_smoothness': float(np.median(smoothness)),
149        'mean_reported_residual': float(np.mean(residuals)),
150        'wall_seconds': elapsed,
151        'ns_matrix_multiplications': ns_mults,
152        'loss_curve': losses,
153    }
154
155
156def main():
157    check = scalar_check()
158    methods = ['exact', 't5', 't2', 'adaptive']
159    all_runs = {}
160    for method in methods:
161        all_runs[method] = [run_optimizer(method, seed) for seed in (11, 29, 47)]
162    summary = {}
163    for method, runs in all_runs.items():
164        keys = ['final_loss', 'best_loss', 'loss_at_60', 'loss_at_120',
165                'mean_update_smoothness', 'median_update_smoothness',
166                'mean_reported_residual', 'wall_seconds', 'ns_matrix_multiplications']
167        summary[method] = {k: float(np.mean([x[k] for x in runs])) for k in keys}
168    out = {'duality_verification': duality_check(), 'scalar_verification': check, 'summary_mean_over_3_seeds': summary,
169           'runs': all_runs}
170    with open('results.json', 'w') as f:
171        json.dump(out, f, indent=2)
172    print(json.dumps({'duality_verification': out['duality_verification'], 'scalar_verification': check, 'summary_mean_over_3_seeds': summary}, indent=2))
173
174
175if __name__ == '__main__':
176    main()