Spaces:
Sleeping
Sleeping
File size: 3,426 Bytes
209ba5e | 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 85 86 87 88 | # Copyright (C) 2026 Hengzhe Zhao. All rights reserved.
# Licensed under AGPL-3.0 or a commercial license. See LICENSE.
"""Tensor-native optimization, shared by CL, MXL and GMNL.
The parameter layout and likelihood belong to the calling model. Scaling by
the number of choice tasks keeps the stopping tolerance independent of N*T.
"""
from dataclasses import dataclass
from typing import Callable
import time
import numpy as np
import torch
@dataclass
class OptimizationResult:
x: np.ndarray
fun: float
success: bool
message: str
nit: int
runtime_seconds: float
def synchronize(device: torch.device) -> None:
if device.type == 'mps':
torch.mps.synchronize()
elif device.type == 'cuda':
torch.cuda.synchronize(device)
def fit_lbfgs(objective: Callable, initial, *, device, maxiter: int,
loss_scale: int = 1) -> OptimizationResult:
if maxiter < 1 or loss_scale < 1:
raise ValueError('maxiter and loss_scale must be positive.')
device = torch.device(device)
theta = torch.tensor(initial, dtype=torch.float32, device=device, requires_grad=True)
if theta.ndim != 1 or not torch.isfinite(theta).all():
raise ValueError('Initial parameters must be a finite one-dimensional vector.')
# A per-task tolerance compatible with FP32 objective resolution.
# This applies to the web prototype, not the controlled paper benchmark.
tolerance = 1e-4
optimizer = torch.optim.LBFGS(
[theta], max_iter=maxiter, history_size=100, line_search_fn='strong_wolfe',
tolerance_grad=tolerance, tolerance_change=1e-9,
)
def closure():
optimizer.zero_grad()
loss = objective(theta) / loss_scale
loss.backward()
if not torch.isfinite(loss) or not torch.isfinite(theta.grad).all():
raise FloatingPointError('Non-finite likelihood or gradient during estimation.')
return loss
synchronize(device)
start = time.perf_counter()
error = None
try:
optimizer.step(closure)
except FloatingPointError as exc:
error = str(exc)
synchronize(device)
runtime = time.perf_counter() - start
# Evaluate the returned point, not the last line-search trial.
optimizer.zero_grad()
loss = objective(theta)
(loss / loss_scale).backward()
grad = float(theta.grad.detach().abs().max().cpu())
value = float(loss.detach().cpu())
state = optimizer.state[theta]
iterations = int(state.get('n_iter', 0))
valid = np.isfinite(value) and np.isfinite(grad)
success = bool(valid and error is None and grad <= tolerance)
if error or not valid:
message = error or 'Non-finite likelihood or gradient at the returned point.'
elif success:
message = f'L-BFGS gradient tolerance reached (per-task gradient {grad:.2g}).'
elif iterations >= maxiter:
message = f'Iteration limit reached; per-task gradient {grad:.2g} exceeds {tolerance:g}.'
elif state.get('func_evals', 0) >= optimizer.param_groups[0]['max_eval']:
message = f'Evaluation limit reached; per-task gradient {grad:.2g}.'
else:
message = f'L-BFGS stopped with small steps or objective changes; per-task gradient {grad:.2g} exceeds {tolerance:g}.'
return OptimizationResult(theta.detach().cpu().numpy().astype(np.float64),
value, success, message, iterations, runtime)
|