repro-a-fully-first-order-layer-for-differentiable-optimization / source_current /baselines /AltDiff_socp.py
| import torch | |
| from torch.autograd import Function | |
| def _safe_norm(x, delta=1e-12): | |
| return torch.sqrt(torch.dot(x, x) + x.new_tensor(delta)) | |
| def _newton_x(P, q, A, b, G, h, s, lam, nu, t, mu, rho, newton_max=30, newton_tol=1e-10, reg=1e-9): | |
| n = P.shape[0] | |
| I = torch.eye(n, device=P.device, dtype=P.dtype) | |
| Hlin = P + rho * (A.T @ A) + rho * (G.T @ G) | |
| rhs_lin = q + A.T @ lam - rho * (A.T @ b) + G.T @ nu + rho * (G.T @ (s - h)) | |
| x = torch.zeros(n, device=P.device, dtype=P.dtype) | |
| for _ in range(newton_max): | |
| r = _safe_norm(x) | |
| u = x / r | |
| c = r - 1.0 + t | |
| alpha = mu + rho * c | |
| F = Hlin @ x + rhs_lin + alpha * u | |
| if torch.linalg.norm(F) <= newton_tol * (1.0 + torch.linalg.norm(rhs_lin)): | |
| break | |
| M = Hlin + rho * (u[:, None] @ u[None, :]) + alpha * ((1.0 / r) * I - (1.0 / (r**3)) * (x[:, None] @ x[None, :])) | |
| M = M + reg * I | |
| dx = torch.linalg.solve(M, -F) | |
| step = 1.0 | |
| x_new = x + step * dx | |
| if not torch.isfinite(x_new).all(): | |
| step = 0.5 | |
| x_new = x + step * dx | |
| x = x_new | |
| r = _safe_norm(x) | |
| u = x / r | |
| c = r - 1.0 + t | |
| alpha = mu + rho * c | |
| M = Hlin + rho * (u[:, None] @ u[None, :]) + alpha * ((1.0 / r) * I - (1.0 / (r**3)) * (x[:, None] @ x[None, :])) | |
| M = M + reg * I | |
| return x, M, u | |
| def alt_diff_qp_ball(P, q, A, b, G, h, rho=50.0, eps=1e-8, max_iter=2000): | |
| if rho <= 0: | |
| raise ValueError("rho must be > 0") | |
| rho = torch.tensor(float(rho), device=q.device, dtype=q.dtype) | |
| n = P.shape[0] | |
| m = b.shape[0] | |
| p = h.shape[0] | |
| x = torch.zeros(n, device=q.device, dtype=q.dtype) | |
| s = torch.zeros(p, device=q.device, dtype=q.dtype) | |
| t = torch.zeros((), device=q.device, dtype=q.dtype) # scalar slack | |
| lam = torch.zeros(m, device=q.device, dtype=q.dtype) | |
| nu = torch.zeros(p, device=q.device, dtype=q.dtype) | |
| mu = torch.zeros((), device=q.device, dtype=q.dtype) # scalar multiplier | |
| Jx = torch.zeros((n, n), device=q.device, dtype=q.dtype) | |
| Js = torch.zeros((p, n), device=q.device, dtype=q.dtype) | |
| Jt = torch.zeros((1, n), device=q.device, dtype=q.dtype) | |
| Jlam = torch.zeros((m, n), device=q.device, dtype=q.dtype) | |
| Jnu = torch.zeros((p, n), device=q.device, dtype=q.dtype) | |
| Jmu = torch.zeros((1, n), device=q.device, dtype=q.dtype) | |
| I = torch.eye(n, device=q.device, dtype=q.dtype) | |
| for _ in range(max_iter): | |
| x, M, u = _newton_x(P, q, A, b, G, h, s, lam, nu, t, mu, rho) | |
| term = I | |
| if m: | |
| term = term + A.T @ Jlam | |
| if p: | |
| term = term + G.T @ Jnu + rho * (G.T @ Js) | |
| term = term + (u[:, None] @ Jmu) + rho * (u[:, None] @ Jt) | |
| Jx = torch.linalg.solve(M, -term) | |
| # s-step | |
| if p: | |
| pre = -(nu / rho) - (G @ x - h) | |
| s_new = torch.relu(pre) | |
| gate_s = (s_new > 0).to(q.dtype) | |
| Js = gate_s[:, None] * (-(1.0 / rho) * Jnu - (G @ Jx)) | |
| s = s_new | |
| # t-step (scalar) | |
| r = _safe_norm(x) | |
| g = 1.0 - r - (mu / rho) | |
| t_new = torch.relu(g) | |
| gate_t = (t_new > 0).to(q.dtype) # scalar 0/1 | |
| Jt = gate_t * (-(u[None, :] @ Jx) - (1.0 / rho) * Jmu) | |
| t = t_new | |
| # dual updates | |
| if m: | |
| lam = lam + rho * (A @ x - b) | |
| Jlam = Jlam + rho * (A @ Jx) | |
| if p: | |
| nu = nu + rho * (G @ x + s - h) | |
| Jnu = Jnu + rho * ((G @ Jx) + Js) | |
| mu = mu + rho * (r - 1.0 + t) | |
| Jmu = Jmu + rho * ((u[None, :] @ Jx) + Jt) | |
| # stopping (primal residuals) | |
| r_eq = torch.linalg.norm(A @ x - b) if m else x.new_tensor(0.0) | |
| r_in = torch.linalg.norm(G @ x + s - h) if p else x.new_tensor(0.0) | |
| r_ball = torch.abs(r - 1.0 + t) | |
| if (r_eq + r_in + r_ball) <= eps: | |
| break | |
| return x, Jx | |
| def AltDiffLayer(eps=1e-8, max_iter=2000, rho=50.0): | |
| class L(Function): | |
| def forward(ctx, P, q, G=None, h=None, A=None, b=None): | |
| def bat(x, d): | |
| if x is None: return None | |
| return x.unsqueeze(0) if x.dim() == d else x | |
| P = bat(P, 2); q = bat(q, 1) | |
| B, n = q.shape | |
| A = bat(A, 2) if A is not None else None | |
| b = bat(b, 1) if b is not None else None | |
| G = bat(G, 2) if G is not None else None | |
| h = bat(h, 1) if h is not None else None | |
| if A is None or A.numel() == 0: A = torch.empty((1, 0, n), device=q.device, dtype=q.dtype) | |
| if b is None or b.numel() == 0: b = torch.empty((1, 0), device=q.device, dtype=q.dtype) | |
| if G is None or G.numel() == 0: G = torch.empty((1, 0, n), device=q.device, dtype=q.dtype) | |
| if h is None or h.numel() == 0: h = torch.empty((1, 0), device=q.device, dtype=q.dtype) | |
| if A.shape[0] == 1 and B > 1: A = A.expand(B, -1, -1) | |
| if b.shape[0] == 1 and B > 1: b = b.expand(B, -1) | |
| if G.shape[0] == 1 and B > 1: G = G.expand(B, -1, -1) | |
| if h.shape[0] == 1 and B > 1: h = h.expand(B, -1) | |
| P_det = P.detach(); q_det = q.detach() | |
| A_det = A.detach(); b_det = b.detach(); G_det = G.detach(); h_det = h.detach() | |
| xs, Js = [], [] | |
| for i in range(B): | |
| xi, Ji = alt_diff_qp_ball(P_det[i], q_det[i], A_det[i], b_det[i], G_det[i], h_det[i], | |
| rho=rho, eps=eps, max_iter=max_iter) | |
| xs.append(xi); Js.append(Ji) | |
| X = torch.stack(xs, 0) | |
| J = torch.stack(Js, 0) | |
| ctx.save_for_backward(J) | |
| ctx.unbatched = (q.dim() == 1 and P.dim() == 2) | |
| return X[0] if ctx.unbatched else X | |
| def backward(ctx, grad_output): | |
| (J,) = ctx.saved_tensors | |
| g = grad_output.unsqueeze(0) if grad_output.dim() == 1 else grad_output | |
| grad_q = torch.zeros_like(g) | |
| for i in range(g.shape[0]): | |
| grad_q[i] = J[i].T.mv(g[i]) | |
| if grad_output.dim() == 1: | |
| grad_q = grad_q[0] | |
| return None, grad_q, None, None, None, None | |
| return L.apply | |