Fejer reflection accelerator for fixed-point layers / fejer_experiment.py

Failed on benchmark

Raw ⬇ ZIP
 1import json, math, time
 2import numpy as np
 3
 4
 5def rot(theta):
 6    return np.array([[math.cos(theta), -math.sin(theta)],
 7                     [math.sin(theta), math.cos(theta)]], dtype=float)
 8
 9
10def J_from_R(R):
11    return (np.eye(R.shape[0]) + R) / 2.0
12
13
14def fejer(J, y0, K):
15    z = y0.copy()
16    a = y0.copy()
17    for _ in range(K):
18        z = 2.0 * (J @ z) - z
19        a += z
20    return a / (K + 1)
21
22
23def ppm(J, y0, n):
24    y = y0.copy()
25    for _ in range(n):
26        y = J @ y
27    return y
28
29
30def residual(J, y):
31    return np.linalg.norm(J @ y - y)
32
33
34def run():
35    y0 = np.array([1.0, 0.0])
36    rows = []
37    # Prediction 1: worst-case over rotation angle equals 1/(K+1).
38    for K in [1, 3, 7, 15, 31]:
39        th = np.linspace(0, 2*np.pi, 200001)
40        # exact closed form for the residual of the averaged reflection
41        vals = np.abs(np.sin((K + 1) * th / 2.0)) / (K + 1)
42        observed = float(vals.max())
43        predicted = 1.0 / (K + 1)
44        rows.append({'test':'worst_case_bound','K':K,'observed':observed,
45                     'predicted':predicted,'relative_error':abs(observed-predicted)/predicted})
46    # Prediction 2: transition at (K+1) theta/2 = pi/2, i.e. theta*=pi/(K+1).
47    # Sweep first crossing of 90% of worst case for K=15.
48    K = 15
49    angles = np.linspace(1e-7, np.pi, 200000)
50    vals = np.abs(np.sin((K + 1) * angles / 2.0)) / (K + 1)
51    target = .9 / (K + 1)
52    idx = np.where(vals >= target)[0][0]
53    observed_theta = float(angles[idx])
54    predicted_theta = math.asin(.9) * 2.0 / (K + 1)
55    rows.append({'test':'transition_90pct','K':K,'observed_theta':observed_theta,
56                 'predicted_theta':predicted_theta,'observed_Ktheta':observed_theta*K,
57                 'predicted_Ktheta':predicted_theta*K,
58                 'relative_error':abs(observed_theta-predicted_theta)/predicted_theta})
59    # Prediction 3: low-frequency residual is approximately theta/2 and does not
60    # improve with K while K theta << 1; report ratios over a parameter sweep.
61    low_rows = []
62    for K in [3, 7, 15, 31]:
63        theta = 0.01 / (K + 1)  # safely in the low-frequency regime
64        R = rot(theta); J = J_from_R(R)
65        got = residual(J, fejer(J, y0, K))
66        predicted = theta / 2.0
67        low_rows.append({'K':K,'theta':theta,'observed':got,'predicted_theta_over_2':predicted,
68                         'ratio':got/predicted})
69    # Direct numerical identity check on several angles and K.
70    identity_err = 0.0
71    for theta in np.linspace(0.03, 5.9, 12):
72        R = rot(float(theta)); J = J_from_R(R)
73        for K in [0,1,2,5,11]:
74            yh = fejer(J, y0, K)
75            rhs = (np.linalg.matrix_power(R, K+1) @ y0 - y0) / (2*(K+1))
76            identity_err = max(identity_err, np.linalg.norm((J @ yh-yh)-rhs))
77    # Equal oracle evaluations: Fejer K uses K J calls; PPM uses K calls.
78    comparison = []
79    theta = 0.35
80    R = rot(theta); J = J_from_R(R)
81    for K in [3,7,15,31]:
82        yf = fejer(J,y0,K); yp = ppm(J,y0,K)
83        comparison.append({'evaluations':K,'fejer_residual':residual(J,yf),
84                           'ppm_residual':residual(J,yp),
85                           'ratio_fejer_over_ppm':residual(J,yf)/residual(J,yp)})
86    out = {'identity_max_abs_error':identity_err,'predictions':rows,
87           'low_frequency':low_rows,'equal_eval_rotation_theta':theta,
88           'comparison':comparison,
89           'notes':'R is a rotation, hence nonexpansive; J=(I+R)/2 is an exact resolvent of a linear maximal monotone skew operator when -1 is absent.'}
90    print(json.dumps(out, indent=2))
91
92if __name__ == '__main__':
93    run()