import json import time import numpy as np SEED = 194 rng = np.random.default_rng(SEED) def ns_transform(M, depth): """Finite Newton-Schulz polar-like transform, scaled back to M's units.""" m, n = M.shape alpha = max(1e-12, np.linalg.norm(M, 2)) X = M / alpha if m >= n: I = np.eye(n) for _ in range(depth): X = 0.5 * X @ (3.0 * I - X.T @ X) else: I = np.eye(m) for _ in range(depth): X = 0.5 * (3.0 * I - X @ X.T) @ X return alpha * X def exact_polar(M): U, _, Vt = np.linalg.svd(M, full_matrices=False) return U @ Vt def residual(M): alpha = max(1e-12, np.linalg.norm(M, 2)) X = M / alpha m, n = X.shape if m >= n: R = X.T @ X - np.eye(n) else: R = X @ X.T - np.eye(m) return np.linalg.norm(R, 'fro') / np.sqrt(min(m, n)) def duality_check(): # Numerically test sup_{||X||op <= 1} = ||S||_* using polar(S), # plus random feasible competitors that should not exceed the optimum. r = np.random.default_rng(SEED + 7) S = r.normal(size=(9, 6)) P = exact_polar(S) optimum = float(np.sum(np.linalg.svd(S, compute_uv=False))) achieved = float(np.sum(S * P)) competitors = [] for _ in range(200): Q = r.normal(size=S.shape) qnorm = np.linalg.norm(Q, 2) competitors.append(float(np.sum(S * (Q / max(qnorm, 1e-12))))) return {'nuclear_norm': optimum, 'polar_objective': achieved, 'absolute_gap': abs(achieved - optimum), 'max_random_feasible_objective': max(competitors), 'random_excess_over_optimum': max(competitors) - optimum} def scalar_check(): # Includes a zero, a very small singular value, and a well-scaled one. s = np.array([0.0, 1e-4, 0.03, 0.2, 0.7, 0.95]) vals = [] x = s.copy() for t in range(7): vals.append(x.copy()) x = 0.5 * x * (3.0 - x * x) vals = np.asarray(vals) err = np.abs(1.0 - vals[:, 1:]) # For a direction starting below one, normalized error contracts # rapidly once it is in the attraction basin. tail = err[1:, -1] contraction = (tail[1:] / np.maximum(tail[:-1], 1e-30)).tolist() # A finite polynomial maps zero continuously to zero, unlike polar's # nonzero-singular-value response (the intended smoothing observation). smooth_gap_at_small = float(vals[2, 1]) return { 'values': vals.tolist(), 'small_sigma_response_t2': smooth_gap_at_small, 'zero_response_all_t': vals[:, 0].tolist(), 'largest_direction_errors': tail.tolist(), 'successive_error_ratios': contraction, 't7_small_sigma_error': float(abs(1.0 - vals[7-1, 1])), } def make_problem(seed, m=32, n=24): r = np.random.default_rng(seed) # Ill-conditioned target: broad singular spectrum makes shallow NS visibly # different from exact polar while remaining a tiny CPU problem. U, _ = np.linalg.qr(r.normal(size=(m, n))) V, _ = np.linalg.qr(r.normal(size=(n, n))) singulars = np.geomspace(8.0, 0.02, n) A = U @ np.diag(singulars) @ V.T # A diagonal preconditioner creates nontrivial, changing momentum matrices. weights = np.geomspace(0.4, 2.5, m)[:, None] * np.geomspace(0.7, 1.8, n)[None, :] return A, weights def run_optimizer(kind, seed, steps=180): A, weights = make_problem(seed) r = np.random.default_rng(seed + 1000) W = 0.35 * r.normal(size=A.shape) M = np.zeros_like(W) beta, lr = 0.90, 0.055 losses, residuals, smoothness = [], [], [] prev_dir, prev_M = None, None ns_mults = 0 t0 = time.perf_counter() for step in range(steps): grad = weights * (W - A) M = beta * M + (1.0 - beta) * grad if kind == 'exact': direction = exact_polar(M) orth_res = 0.0 elif kind == 't5': direction = ns_transform(M, 5) orth_res = residual(direction) ns_mults += 5 elif kind == 't2': direction = ns_transform(M, 2) orth_res = residual(direction) ns_mults += 2 elif kind == 'adaptive': depth = 2 if step < int(0.7 * steps) else 4 direction = ns_transform(M, depth) orth_res = residual(direction) ns_mults += depth else: raise ValueError(kind) # Muon-style update, with a common scalar normalization so comparisons # reflect spectral direction rather than arbitrary matrix magnitude. direction = direction * (np.linalg.norm(M, 'fro') / max(np.linalg.norm(direction, 'fro'), 1e-12)) if prev_dir is not None: smoothness.append(np.linalg.norm(direction - prev_dir, 'fro') / max(np.linalg.norm(M - prev_M, 'fro'), 1e-12)) W -= lr * direction losses.append(float(0.5 * np.mean(weights * (W - A) ** 2))) residuals.append(float(orth_res)) prev_dir, prev_M = direction.copy(), M.copy() elapsed = time.perf_counter() - t0 return { 'final_loss': losses[-1], 'best_loss': min(losses), 'loss_at_60': losses[59], 'loss_at_120': losses[119], 'mean_update_smoothness': float(np.mean(smoothness)), 'median_update_smoothness': float(np.median(smoothness)), 'mean_reported_residual': float(np.mean(residuals)), 'wall_seconds': elapsed, 'ns_matrix_multiplications': ns_mults, 'loss_curve': losses, } def main(): check = scalar_check() methods = ['exact', 't5', 't2', 'adaptive'] all_runs = {} for method in methods: all_runs[method] = [run_optimizer(method, seed) for seed in (11, 29, 47)] summary = {} for method, runs in all_runs.items(): keys = ['final_loss', 'best_loss', 'loss_at_60', 'loss_at_120', 'mean_update_smoothness', 'median_update_smoothness', 'mean_reported_residual', 'wall_seconds', 'ns_matrix_multiplications'] summary[method] = {k: float(np.mean([x[k] for x in runs])) for k in keys} out = {'duality_verification': duality_check(), 'scalar_verification': check, 'summary_mean_over_3_seeds': summary, 'runs': all_runs} with open('results.json', 'w') as f: json.dump(out, f, indent=2) print(json.dumps({'duality_verification': out['duality_verification'], 'scalar_verification': check, 'summary_mean_over_3_seeds': summary}, indent=2)) if __name__ == '__main__': main()