File size: 2,808 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 | import torch
import numpy as np
from .moo import cagrad_update_k2, raco_update_k2
class ConflictingQuadratics:
def __init__(self, dim=100, conflict_angle=np.pi * 0.6, seed=42):
rng = np.random.RandomState(seed)
self.dim = dim
self.A1 = rng.randn(dim, dim).astype(np.float32)
self.A1 = self.A1.T @ self.A1 / dim
self.b1 = rng.randn(dim).astype(np.float32) * 0.1
self.A2 = rng.randn(dim, dim).astype(np.float32)
self.A2 = self.A2.T @ self.A2 / dim
self.b2 = rng.randn(dim).astype(np.float32) * 0.1
v = rng.randn(dim).astype(np.float32)
v = v / np.linalg.norm(v)
u = rng.randn(dim).astype(np.float32)
u = u - np.dot(u, v) * v
u = u / np.linalg.norm(u)
g1_0 = v
g2_0 = v * np.cos(conflict_angle) + u * np.sin(conflict_angle)
scale = np.linalg.norm(self.A1 @ self.b1 + self.b1)
self.b1 = self.b1 - g1_0 * scale
self.b2 = self.b2 - g2_0 * scale
def run_synthetic_comparison(dim=100, conflict_angle=np.pi*0.6, steps=200, lr=0.01, seed=42):
import math
prob = ConflictingQuadratics(dim=dim, conflict_angle=conflict_angle, seed=seed)
A1_t = torch.from_numpy(prob.A1).float()
b1_t = torch.from_numpy(prob.b1).float()
A2_t = torch.from_numpy(prob.A2).float()
b2_t = torch.from_numpy(prob.b2).float()
theta0 = torch.from_numpy(np.random.RandomState(seed + 1).randn(dim).astype(np.float32) * 0.5)
methods = {}
for name, use_clip, w1, w2 in [
("DPO_LW", False, 0.7, 0.3),
("CAGrad", False, 0.7, 0.3),
("RACO", True, 0.7, 0.3),
]:
theta = theta0.clone().requires_grad_(True)
losses = {0: [], 1: []}
opt = torch.optim.SGD([theta], lr=lr)
for t in range(steps):
opt.zero_grad()
l1 = 0.5 * (theta @ A1_t @ theta) + b1_t @ theta
l2 = 0.5 * (theta @ A2_t @ theta) + b2_t @ theta
g1 = torch.autograd.grad(l1, theta, retain_graph=True)[0]
g2 = torch.autograd.grad(l2, theta, retain_graph=True)[0]
g1_list = [g1]
g2_list = [g2]
if name == "DPO_LW":
g_update = g1 * w1 + g2 * w2
elif name == "CAGrad":
gl = cagrad_update_k2(g1_list, g2_list, w1, w2, c=0.4)
g_update = gl[0]
elif name == "RACO":
gl = raco_update_k2(g1_list, g2_list, w1, w2, c=0.4)
g_update = gl[0]
theta.grad = g_update.detach() if hasattr(g_update, 'detach') else g_update
opt.step()
losses[0].append(float(l1.detach()))
losses[1].append(float(l2.detach()))
methods[name] = {"theta": theta.detach().numpy().copy(), "losses": losses}
return methods
|