File size: 2,820 Bytes
1c1d9aa | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import torch, sys, json
sys.path.insert(0, "..")
from src.moo import cagrad_update_k2, raco_update_k2, _grads_dot, _solve_lambda_k2
def test_clipping_triggers():
"""When CAGrad wants a large correction on the less-preferred objective,
RACO should clip it."""
w1, w2 = 0.8, 0.2
torch.manual_seed(42)
g1 = torch.randn(100) * 3.0
g2 = -g1 * 0.9 + torch.randn(100) * 0.3
g1_list = [g1]
g2_list = [g2]
cagrad_g = cagrad_update_k2(g1_list, g2_list, w1, w2, c=0.5)
raco_g = raco_update_k2(g1_list, g2_list, w1, w2, c=0.5)
cg = cagrad_g[0]
rg = raco_g[0]
alignment_ca_g1 = float((cg @ g1).item())
alignment_ca_g2 = float((cg @ g2).item())
alignment_ra_g1 = float((rg @ g1).item())
alignment_ra_g2 = float((rg @ g2).item())
result = {
"test": "Strong conflicting objectives, w1=0.8, w2=0.2",
"w1": w1, "w2": w2,
"g1_norm": float(g1.norm().item()),
"g2_norm": float(g2.norm().item()),
"cos_angle": float((g1 @ g2 / (g1.norm() * g2.norm())).item()),
"cagrad_alignment_g1": alignment_ca_g1,
"cagrad_alignment_g2": alignment_ca_g2,
"raco_alignment_g1": alignment_ra_g1,
"raco_alignment_g2": alignment_ra_g2,
"cagrad_more_g2_magnitude": abs(alignment_ca_g2) > abs(alignment_ra_g2),
"raco_better_respects_weights": alignment_ra_g2 > alignment_ca_g2,
}
return result
def test_theorem_convergence():
"""Test that CAGrad-Clip provides valid descent directions."""
results = {}
for c in [0.1, 0.3, 0.5, 0.7]:
w1 = 0.7
w2 = 0.3
torch.manual_seed(123)
g1 = torch.randn(50) * 2.0
g2 = -g1 * 0.7 + torch.randn(50) * 0.5
g1_list, g2_list = [g1], [g2]
try:
rg = raco_update_k2(g1_list, g2_list, w1, w2, c=c)
g_update = rg[0]
g0 = g1 * w1 + g2 * w2
descent_ok = float((g_update @ g_update).item()) > 0
results[f"c={c}"] = {
"valid": descent_ok,
"g0_norm": float(g0.norm().item()),
"g_update_norm": float(g_update.norm().item()),
}
except Exception as e:
results[f"c={c}"] = {"error": str(e)}
return results
if __name__ == "__main__":
all_tests = {}
all_tests["test_clipping"] = test_clipping_triggers()
all_tests["test_convergence"] = test_theorem_convergence()
print(json.dumps(all_tests, indent=2))
with open("/Users/equan_p/Developer/playground/ICML-2/repro_raco/outputs/tier1/test_results.json", "w") as f:
json.dump(all_tests, f, indent=2)
all_pass = all(
t.get("raco_better_respects_weights", True) or t.get("valid", True)
for t in all_tests.values()
)
print(f"\nAll tests passed: {all_pass}")
|