"""FiniteModel-related helpers and custom OT solvers built on FiniteModel/OT.""" import logging import math import torch import torch.nn.functional as F from torch import nn from optimal_transport.ot import OT from tools.feedback import logger from models.helpers import FixedFirstIntercept class TwoStageOptimizer(torch.optim.Optimizer): """ Wrapper that runs Adam with a large learning rate until the gradient norm exceeds a threshold (or a maximum warmup step count), then switches to SGD with a steady learning rate for the remainder of training. """ def __init__( self, params, steady_lr: float, warmup_lr: float, betas=(0.5, 0.9), grad_threshold: float = 0.5, max_warm_steps: int = 200, steady_momentum: float = 0.0, steady_weight_decay: float = 0.0, ): if warmup_lr <= 0 or steady_lr <= 0: raise ValueError("Learning rates must be positive.") if grad_threshold <= 0: raise ValueError("grad_threshold must be positive.") if max_warm_steps <= 0: raise ValueError("max_warm_steps must be positive.") if steady_momentum < 0: raise ValueError("steady_momentum must be non-negative.") if steady_weight_decay < 0: raise ValueError("steady_weight_decay must be non-negative.") params = list(params) if len(params) == 0: raise ValueError("Optimizer received an empty parameter list.") self._adam = torch.optim.Adam(params, lr=warmup_lr, betas=betas) self._sgd = torch.optim.SGD( params, lr=steady_lr, momentum=steady_momentum, weight_decay=steady_weight_decay, ) self.warm_lr = warmup_lr self.steady_lr = steady_lr self.grad_threshold = grad_threshold self.max_warm_steps = max_warm_steps self.steady_momentum = steady_momentum self.steady_weight_decay = steady_weight_decay self._use_warm_lr = True self._warm_steps = 0 self.param_groups = self._adam.param_groups self.state = self._adam.state def zero_grad(self): """Zero gradients for both Adam and SGD phases.""" self._adam.zero_grad() self._sgd.zero_grad() def state_dict(self): """Return optimizer state for checkpointing across warm-to-steady transition.""" return { "adam_state": self._adam.state_dict(), "sgd_state": self._sgd.state_dict(), "warm_lr": self.warm_lr, "steady_lr": self.steady_lr, "grad_threshold": self.grad_threshold, "max_warm_steps": self.max_warm_steps, "steady_momentum": self.steady_momentum, "steady_weight_decay": self.steady_weight_decay, "use_warm_lr": self._use_warm_lr, "warm_steps": self._warm_steps, } def load_state_dict(self, state_dict): """Restore optimizer state and adjust current phase accordingly.""" self._adam.load_state_dict(state_dict["adam_state"]) self._sgd.load_state_dict(state_dict["sgd_state"]) self.warm_lr = state_dict["warm_lr"] self.steady_lr = state_dict["steady_lr"] self.grad_threshold = state_dict["grad_threshold"] self.max_warm_steps = state_dict["max_warm_steps"] self.steady_momentum = state_dict.get("steady_momentum", self.steady_momentum) self.steady_weight_decay = state_dict.get("steady_weight_decay", self.steady_weight_decay) self._use_warm_lr = state_dict["use_warm_lr"] self._warm_steps = state_dict["warm_steps"] self._set_lr(self.warm_lr if self._use_warm_lr else self.steady_lr) self.param_groups = self._adam.param_groups if self._use_warm_lr else self._sgd.param_groups def _set_lr(self, lr): """Apply same learning rate to both internal optimizers.""" for group in self._adam.param_groups: group["lr"] = lr for group in self._sgd.param_groups: group["lr"] = lr def step(self, closure=None, grad_norm: float | None = None): """Perform a parameter update and switch to SGD once warmup criteria are met.""" if self._use_warm_lr: trigger = False if grad_norm is not None and grad_norm >= self.grad_threshold: trigger = True elif self._warm_steps >= self.max_warm_steps: trigger = True if trigger: self._use_warm_lr = False self._set_lr(self.steady_lr) self.param_groups = self._sgd.param_groups self.state = self._sgd.state logger.info( "[TwoStageOptimizer] Switching to SGD: grad_norm=%.3f, steps=%d", grad_norm if grad_norm is not None else float("nan"), self._warm_steps, ) self._warm_steps += 1 if self._use_warm_lr: self._adam.step(closure=closure) else: self._sgd.step(closure=closure) class FCOTSeparable(OT): """ Finitely-concave OT solver for separable costs using FiniteSeparableModel. This is a specialized version of FCOT that exploits cost separability: c(x,y) = sum_d c_d(x_d, y_d) Key advantages over standard FCOT: - Exponentially fewer parameters: O(d·|Y_0|) vs O(|Y_0|^d) - Exact discrete transforms (no numerical optimization needed) - Always converges since transforms are solved exactly on grids Parameterization: - Uses FiniteSeparableModel with discretized grids on [-R, R] - Kernel is the 1D version of the cost: kernel(x,y) = c(x,y) for scalars - Mode is "concave" for c-concave potential representation Dual objective: D = E_x[u(x)] + E_y[u^c(y)] Example: >>> from models import FiniteSeparableModel >>> >>> # For L2^2 cost: c(x,y) = ||x-y||^2 = sum_d (x_d - y_d)^2 >>> def kernel_1d(x, y): ... return -(x - y)**2 # Negative for c-concave >>> >>> # Create separable model >>> model = FiniteSeparableModel( ... kernel=kernel_1d, ... num_dims=2, ... radius=5.0, ... y_accuracy=0.1, ... x_accuracy=0.1, ... mode="concave" ... ) >>> >>> # Create solver >>> fcot_sep = FCOTSeparable( ... input_dim=2, ... model=model, ... inverse_kx=lambda x, p: x - 0.5 * p # For L2^2 ... ) >>> >>> # Or use factory method >>> fcot_sep = FCOTSeparable.initialize_right_architecture( ... dim=2, ... radius=5.0, ... n_params=100, # 50 Y-grid points per dimension ... x_accuracy=0.1, ... kernel_1d=kernel_1d, ... inverse_kx=lambda x, p: x - 0.5 * p ... ) """ @staticmethod def initialize_right_architecture( dim: int, radius: float, n_params: int, x_accuracy: float, kernel_1d, # 1D kernel function c_d(x, y) for scalars inverse_kx, # Inverse gradient for the full cost outer_lr: float = 1e-3, betas=(0.5, 0.9), device: str = "cpu", temp: float = 50.0, epsilon: float = 1e-4, cache_gradients: bool = False, warmup_lr: float | None = None, warmup_grad_threshold: float = 0.5, warmup_max_steps: int = 200, sgd_momentum: float = 0.0, sgd_weight_decay: float = 0.0, temp_min: float | None = None, temp_max: float | None = None, temp_warmup_iters: int | None = None, reactivate_every: int | None = None, reactivate_eps: float = 1e-2, full_refresh_every: int | None = None, **model_kwargs ): """ Factory for building a grid-based separable FCOT solver. Parameters ---------- dim : int Input/output dimensionality. radius : float Bound on coordinate magnitude (domain = [-radius, radius]^d). n_params : int Total number of intercepts (must be divisible by `dim`). x_accuracy : float Spacing for the X grid used in infimal convolution transforms. kernel_1d : callable 1D kernel function `c_d(x, y)` describing separable costs. inverse_kx : callable Analytic inverse gradient mapping ∇_x c(x, y) → y used in Monge map. cache_gradients : bool Whether to precompute derivatives on the grid for faster transforms. Other params : See FCOTSeparable.__init__ for further hyperparameters. Returns ------- FCOTSeparable Instantiated solver configured for the requested grid resolution. Notes ----- Parameter budget satisfies `n_params = dim × |Y_0|`. """ if n_params % dim != 0: raise ValueError(f"n_params={n_params} must be divisible by dim={dim}.") ny = n_params // dim if ny < 2: raise ValueError("Need at least two Y-grid points (n_params/dim >= 2).") y_accuracy = (2 * radius) / (ny - 1) actual_params = ny * dim logger.info(f"[FCOT-SEP ARCH] dim={dim}, radius={radius}, ny={ny}") logger.info(f" x_accuracy={x_accuracy:.4f}") logger.info(f" y_accuracy={y_accuracy:.4f}") logger.info(f" actual_params={actual_params}") logger.info(f" cache_gradients={cache_gradients}") # Build FiniteSeparableModel model_kwargs = dict(model_kwargs) # Backwards compatibility if 'lr' in model_kwargs and outer_lr == 1e-3: outer_lr = model_kwargs.pop('lr') logger.warning("[DEPRECATION] 'lr' argument is deprecated; use 'outer_lr' instead.") model = FiniteSeparableModel( kernel=kernel_1d, num_dims=dim, radius=radius, y_accuracy=y_accuracy, x_accuracy=x_accuracy, mode="concave", temp=temp, epsilon=epsilon, cache_gradients=cache_gradients, **model_kwargs, ).to(device) return FCOTSeparable( input_dim=dim, model=model, inverse_kx=inverse_kx, outer_lr=outer_lr, betas=betas, device=device, warmup_lr=warmup_lr, warmup_grad_threshold=warmup_grad_threshold, warmup_max_steps=warmup_max_steps, sgd_momentum=sgd_momentum, sgd_weight_decay=sgd_weight_decay, temp_min=temp_min, temp_max=temp_max, temp_warmup_iters=temp_warmup_iters, reactivate_every=reactivate_every, reactivate_eps=reactivate_eps, full_refresh_every=full_refresh_every, ) def __init__( self, input_dim: int, model: nn.Module, # FiniteSeparableModel(mode="concave") inverse_kx, # (x, p) ↦ y solving ∇_x k(x,y) = p outer_lr: float = 1e-3, lr: float | None = None, # deprecated alias betas=(0.5, 0.9), device: str = "cpu", warmup_lr: float | None = None, warmup_grad_threshold: float = 0.5, warmup_max_steps: int = 200, sgd_momentum: float = 0.0, sgd_weight_decay: float = 0.0, temp_min: float | None = None, temp_max: float | None = None, temp_warmup_iters: int | None = None, reactivate_every: int | None = None, reactivate_eps: float = 1e-2, full_refresh_every: int | None = None, ): """ Initialize separable FCOT solver over a finite grid. Parameters ---------- input_dim : int Dimensionality of source/target space. model : nn.Module FiniteSeparableModel providing u/u^c transforms on a grid. inverse_kx : callable Mapping (x, ∇u(x)) → y used to evaluate transport map. outer_lr : float Learning rate for the outer Adam optimizer. lr : float | None Deprecated alias for `outer_lr`. betas : Tuple[float, float] Adam beta coefficients. temp_min, temp_max, temp_warmup_iters : optional float Controls adaptive temperature scheduling for discrete transforms. reactivate_every, reactivate_eps, full_refresh_every : int/float Hooks that keep inactive intercepts viable by reactivating/refreshing them. """ # Handle deprecated lr parameter before delegating to OT.__init__ if lr is not None: logger.warning("[DEPRECATION] 'lr' argument is deprecated; use 'outer_lr' instead.") outer_lr = lr super().__init__(outer_lr=outer_lr, inner_lr=None) self.input_dim = input_dim self.device = torch.device(device) self.model = model.to(self.device) self.inverse_kx = inverse_kx self.outer_lr = outer_lr self.betas = betas self.inner_steps = 0 self.inner_optimizer = None self.inner_tol = None self.inner_lam = None # Flag to skip generic warm-up in base OT (transforms are exact) self._skip_warmup = True self._temp_min = temp_min self._temp_max = temp_max self._temp_warmup_iters = temp_warmup_iters self._temp_step = 0 self._grad_activity_eps = 1e-8 self._reactivate_every = reactivate_every if reactivate_every and reactivate_every > 0 else None self._reactivate_eps = reactivate_eps self._full_refresh_every = full_refresh_every if full_refresh_every and full_refresh_every > 0 else None self._warned_refresh_unavailable = False self._step_count = 0 # Outer optimizer for model parameters (intercepts) # Warm-start schedule (TwoStageOptimizer) is no longer used. if warmup_lr is not None: logger.warning( "[DEPRECATION] warmup_* arguments are ignored; using Adam with lr=outer_lr." ) self.optimizer = torch.optim.Adam( self.model.parameters(), lr=outer_lr, betas=betas, ) # No inner optimization needed for separable case self.raise_on_inner_divergence = False def _compute_dual_objective(self, X, Y, sample_idx=None): """ Compute the Kantorovich dual objective: D = E_x[u(x)] + E_y[u^c(y)] For separable costs, transforms are computed exactly on grids. """ # First term: E_x[u(x)] _, u_X = self.model.forward(X, selection_mode="soft") term1 = u_X.mean() # Second term: E_y[u^c(y)] # Note: sample_idx not used for separable model (no warm starts needed) _, u_c_Y, converged = self.model.inf_transform(Y) term2 = u_c_Y.mean() # Dual objective to maximize dual_obj = term1 + term2 # Separable model always converges (exact discrete optimization) inner_converged = converged return dual_obj, inner_converged def _dual_objective(self, x_batch, y_batch, idx_y=None, active_inner_steps=None): """ Evaluate dual objective D = E_x[u(X)] + E_y[u^c(Y)] on a minibatch. Returns (D, u_mean, uc_mean, converged) for reporting. """ _, u_vals = self.model.forward(x_batch, selection_mode="soft") u_mean = u_vals.mean() _, uc_vals, converged = self.model.inf_transform(y_batch) uc_mean = uc_vals.mean() D = u_mean + uc_mean return D, u_mean, uc_mean, converged def step(self, x_batch, y_batch, idx_y=None, active_inner_steps=None): """ Perform a gradient step on the separable dual objective with diagnostics. Returns dict with dual objective, mean potentials, convergence status, gradient norm, and intercept activity statistics for logging. """ self._maybe_update_temperature() self.optimizer.zero_grad() D, u_mean, uc_mean, converged = self._dual_objective(x_batch, y_batch, idx_y, active_inner_steps) #TODO fix ths adhoc scaling (-100 * D).backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) grad_norm_sq = 0.0 for p in self.model.parameters(): if p.grad is not None: grad_norm_sq += p.grad.norm().item() ** 2 grad_norm = grad_norm_sq ** 0.5 step_kwargs = {} if isinstance(self.optimizer, TwoStageOptimizer): step_kwargs["grad_norm"] = grad_norm self.optimizer.step(**step_kwargs) self._step_count += 1 self._maybe_reactivate(x_batch) self._maybe_full_refresh() activity = None active_count = None total_count = None intercepts_param = getattr(self.model, "intercepts_param", None) theta_grad = None if intercepts_param is not None and hasattr(intercepts_param, "theta"): theta_grad = intercepts_param.theta.grad if theta_grad is not None: grad_vals = torch.zeros_like(self.model.intercepts, device=theta_grad.device, dtype=theta_grad.dtype) grad_vals[1:, :] = theta_grad.detach() active = (grad_vals.abs() > self._grad_activity_eps).sum().item() total = grad_vals.numel() activity = active / total if total > 0 else None active_count = active total_count = total return { "dual": float(D.detach().item()), "u_mean": float(u_mean.detach().item()), "uc_mean": float(uc_mean.detach().item()), "inner_converged": converged, "grad_norm": grad_norm, "intercept_active_frac": activity, "intercept_active_count": active_count, "intercept_total": total_count, } def _maybe_reactivate(self, x_batch): """ Periodically bump inactive intercepts using the kernel transform to avoid stagnation. Guidelines: - Enabled by setting reactivate_every>0 at init. - Uses gradient activity below _grad_activity_eps to decide "inactive". - Lifts inactive intercepts toward a small margin (eps) above the current kernel lower envelope for the given minibatch. """ if self._reactivate_every is None: return if self._step_count % self._reactivate_every != 0: return intercepts_param = getattr(self.model, "intercepts_param", None) theta_grad = None if intercepts_param is not None and hasattr(intercepts_param, "theta"): theta_grad = intercepts_param.theta.grad if theta_grad is None: return intercepts = self.model.intercepts full_grad = torch.zeros_like(intercepts, device=theta_grad.device, dtype=theta_grad.dtype) full_grad[1:, :] = theta_grad.detach() inactive_mask = full_grad.abs() <= self._grad_activity_eps if inactive_mask.sum().item() == 0: return x_batch = x_batch.detach().to(device=self.device, dtype=intercepts.dtype) if x_batch.dim() != 2: return y_grid = self.model.Y_grid.to(device=self.device, dtype=intercepts.dtype) eps = float(self._reactivate_eps) adjusted = 0 with torch.no_grad(): for dim in range(self.model.num_dims): mask_dim = inactive_mask[:, dim] if not mask_dim.any(): continue x_vals = x_batch[:, dim] kernel_vals = self.model.kernel_fn( x_vals.unsqueeze(-1), y_grid.unsqueeze(0), ) column = intercepts[:, dim] scores = kernel_vals - column.unsqueeze(0) best_scores = scores.min(dim=1).values target = (kernel_vals - best_scores.unsqueeze(1) + eps).min(dim=0).values updated_column = torch.where(mask_dim, torch.maximum(column, target), column) if torch.any(updated_column != column): if intercepts_param is not None and hasattr(intercepts_param, "set_column_from_raw_"): intercepts_param.set_column_from_raw_(dim, updated_column) else: intercepts[:, dim].copy_(updated_column) adjusted += int(mask_dim.sum().item()) if adjusted > 0: logger.info( "[FCOT-SEP] Reactivated %d intercepts (eps=%g) at step %d", adjusted, self._reactivate_eps, self._step_count, ) def _maybe_full_refresh(self): """ Periodically refresh intercepts/optimizer momentum and inject jitter. Enabled when full_refresh_every>0; calls refresh_intercepts_via_transform, resets optimizer momentum on the intercept parameter, and adds tiny noise to break ties. """ if self._full_refresh_every is None: return if self._step_count % self._full_refresh_every != 0: return # Do full refresh + momentum reset + jitter self._refresh_reset_jitter() ############################################################################### # Helpers: reset optimizer momentum + jitter after refresh ############################################################################### @torch.no_grad() def _reset_intercept_momentum(self): """ Reset momentum / Adam state ONLY for intercept parameters. Works for both Adam and TwoStageOptimizer. """ if not hasattr(self, "optimizer") or self.optimizer is None: return # Identify the underlying parameter that stores intercepts. intercepts_param = getattr(self.model, "intercepts_param", None) target_param = None if intercepts_param is not None and hasattr(intercepts_param, "theta"): target_param = intercepts_param.theta else: target_param = getattr(self.model, "intercepts", None) if target_param is None: return # Find optimizer states that correspond to the intercept parameter for group in self.optimizer.param_groups: for p in group["params"]: if p is target_param: state = self.optimizer.state.get(p, None) if state is not None: state.clear() # wipes momentum, exp_avg, exp_avg_sq, step return @torch.no_grad() def _jitter_intercepts(self, strength: float = None): """ Inject tiny Gaussian noise to break ties after b <- b^cc. """ intercepts_param = getattr(self.model, "intercepts_param", None) target_param = None if intercepts_param is not None and hasattr(intercepts_param, "theta"): target_param = intercepts_param.theta else: target_param = getattr(self.model, "intercepts", None) if target_param is None: return # default strength relative to grid resolution if strength is None: strength = 1e-3 * float(self.model.y_accuracy) # e.g., 2e-5 for your test noise = torch.randn_like(target_param) * strength target_param.add_(noise) @torch.no_grad() def _refresh_reset_jitter(self): """ Unified operation: refresh → reset momentum → jitter intercepts. Called inside _maybe_full_refresh(). """ refresh_fn = getattr(self.model, "refresh_intercepts_via_transform", None) if refresh_fn is None: return changed = refresh_fn() if changed > 0: logger.info( f"[FCOT-SEP] Full refresh updated {changed} intercepts at step {self._step_count}" ) self._reset_intercept_momentum() self._jitter_intercepts() def _maybe_update_temperature(self): """Cosine warmup from temp_max → temp_min over temp_warmup_iters steps (if set).""" if self._temp_min is None or self._temp_max is None: return if self._temp_warmup_iters is None or self._temp_warmup_iters <= 0: self.model.temp = self._temp_max return ratio = min(self._temp_step / self._temp_warmup_iters, 1.0) cos_scale = 0.5 * (1 - math.cos(math.pi * ratio)) self.model.temp = self._temp_min + (self._temp_max - self._temp_min) * cos_scale self._temp_step += 1 def save(self, address, iters_done): """ Save model + optimizer checkpoint to disk for caching/resume. """ torch.save( { "model_state_dict": self.model.state_dict(), "optimizer_state_dict": self.optimizer.state_dict(), "iters_done": iters_done, "outer_lr": self.outer_lr, "betas": self.betas, }, address, ) def load(self, address): """ Load saved checkpoint and return iterations already completed. """ checkpoint = torch.load(address, map_location=self.device) self.model.load_state_dict(checkpoint["model_state_dict"]) if "optimizer_state_dict" in checkpoint: self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) return checkpoint.get("iters_done", 0) def transport_X_to_Y( self, X, *, selection_mode: str = "hard", snap_to_grid: bool = True, ): """ Transport samples `X` to `Y` via the c-gradient (Monge map). The method differentiates through the softmax-infimal transforms to obtain ∇u(x) and then applies the provided `inverse_kx`. """ X = X.to(self.device) X.requires_grad_(True) # Compute u(x) with gradients # Use "hard" mode for envelope theorem gradients (more efficient) # Use "soft" mode for differentiable selection (smoother but slower) _, u_X = self.model.forward( X, selection_mode=selection_mode, snap_to_grid=snap_to_grid ) # Compute ∇u(x) via backprop grad_u = torch.autograd.grad( u_X.sum(), X, create_graph=False )[0] # Apply inverse c-gradient Y_pred = self.inverse_kx(X.detach(), grad_u.detach()) return Y_pred def _fit( self, X, Y, iters_done: int = 0, iters: int = 100, inner_steps: int | None = None, # Not used for separable print_every: int | None = 1, callback=None, convergence_tol: float | None = 1e-4, convergence_patience: int | None = 10, batch_size: int | None = None, eval_every: int = 100, warmup_steps: int = 0, warmup_until_converged: bool = False, log_every: int | None = None, log_level: str | None = None, ): """ Fit procedure delegating logging/optimization to the OT base class. inner_steps is always forced to 0 because separable transforms are exact. """ if inner_steps is not None and inner_steps != 0: logger.warning( f"[FCOT-SEP] inner_steps={inner_steps} ignored for separable model " "(transforms computed exactly)" ) # Call parent's _fit which handles all the logging and optimization loop effective_log_level = "info" if log_level is None else log_level return super()._fit( X, Y, iters_done=iters_done, iters=iters, inner_steps=0, # Always 0 for separable print_every=print_every, callback=callback, convergence_tol=convergence_tol, convergence_patience=convergence_patience, batch_size=batch_size, eval_every=eval_every, warmup_steps=warmup_steps, warmup_until_converged=warmup_until_converged, log_every=log_every, log_level=effective_log_level, ) from typing import Optional from torch import nn import torch from tools.utils import moded_max, moded_min from models.inf_convolution import InfConvolution from tools.feedback import logger # ===================================================================== # FiniteSeparableModel: finitely-convex/concave on product kernels # ===================================================================== class FiniteSeparableModel(nn.Module): """ Separable finite model for product kernels K(x,y) = sum_d k(x_d, y_d). The model discretizes the 1-D Y-domain once (Y = Y_0^d) and keeps one intercept per (y_i, dimension) pair. Forward evaluations split across dimensions. Sup/inf transforms can optionally run a coarse-to-fine search: - coarse_x_factor subsamples the X grid by this stride for the coarse pass. - coarse_top_k keeps the top-k coarse candidates per transform call. - coarse_window refines within ±(coarse_window * stride) around each seed. If coarse_x_factor is None or <= 1, transforms reduce over the full X grid. """ def __init__( self, kernel, num_dims: int, radius: float, y_accuracy: float = 1e-2, x_accuracy: float = 1e-2, mode: str = "convex", temp: float = 50.0, epsilon: float = 1e-4, cache_gradients: bool = False, coarse_x_factor: Optional[int] = None, coarse_top_k: int = 1, coarse_window: int = 0, ): super().__init__() assert mode in ("convex", "concave") self.kernel_fn = kernel self.num_dims = num_dims self.radius = radius self.y_accuracy = y_accuracy self.x_accuracy = x_accuracy self.mode = mode self.temp = temp self.epsilon = epsilon self.cache_gradients = cache_gradients # Discrete grids shared across dimensions ny = int(2 * radius / y_accuracy) + 1 nx = int(2 * radius / x_accuracy) + 1 Y_grid = torch.linspace(-radius, radius, ny) X_grid = torch.linspace(-radius, radius, nx) kernel_tensor = kernel(X_grid.reshape(-1, 1), Y_grid.reshape(1, -1)) # (nx, ny) self.register_buffer("Y_grid", Y_grid) self.register_buffer("X_grid", X_grid) self.register_buffer("kernel_tensor", kernel_tensor) if cache_gradients: dx_tensor, dy_tensor = self._precompute_kernel_derivatives() else: dx_tensor = torch.empty(0) dy_tensor = torch.empty(0) self.register_buffer("kernel_dx", dx_tensor) self.register_buffer("kernel_dy", dy_tensor) # Intercept parameterization with fixed first row (gauge: b[0, :] = 0) self.intercepts_param = FixedFirstIntercept(ny=ny, dim=num_dims) # Track last fully refreshed intercepts to make refresh idempotent self._last_refreshed_intercepts: Optional[torch.Tensor] = None # Optional coarse-to-fine transform search settings stride = coarse_x_factor if (coarse_x_factor is not None and coarse_x_factor > 1) else None if stride is not None: coarse_idx = torch.arange(0, nx, stride, dtype=torch.long) if coarse_idx[-1] != nx - 1: coarse_idx = torch.cat([coarse_idx, coarse_idx.new_tensor([nx - 1])]) coarse_top_k = max(1, coarse_top_k) coarse_window = max(0, coarse_window) window_radius = coarse_window * stride offsets = torch.arange(-window_radius, window_radius + 1, dtype=torch.long) else: coarse_idx = torch.empty(0, dtype=torch.long) offsets = torch.empty(0, dtype=torch.long) coarse_top_k = 0 coarse_window = 0 self.register_buffer("coarse_idx", coarse_idx) self.register_buffer("coarse_offsets", offsets) self.coarse_stride = stride if stride is not None else 1 self.coarse_top_k = coarse_top_k self.coarse_window = coarse_window self.use_coarse_search = coarse_idx.numel() > 0 @property def intercepts(self) -> torch.Tensor: """ Full intercept matrix b of shape (ny, num_dims) with gauge b[0, :] = 0. """ return self.intercepts_param.value def refresh_intercepts_via_transform(self) -> int: """ Recompute every intercept column via discrete biconjugation. For a separable model with scores K - b: - In concave mode, we compute: u(x) = min_y [K(x,y) - b(y)] b_new(y) = min_x [K(x,y) - u(x)] - In convex mode, we compute: u(x) = max_y [K(x,y) - b(y)] b_new(y) = max_x [K(x,y) - u(x)] This is the inf–inf / sup–sup biconjugation used in the main algorithm; intercepts are still defined up to a per-column gauge. """ K = self.kernel_tensor.to(device=self.intercepts_param.theta.device, dtype=self.intercepts_param.theta.dtype) # If intercepts already match the last refreshed state, do nothing (idempotent). current = self.intercepts.detach() if self._last_refreshed_intercepts is not None and torch.equal( current, self._last_refreshed_intercepts ): return 0 changed = 0 with torch.no_grad(): for dim in range(self.num_dims): column = self.intercepts[:, dim] scores = K - column.unsqueeze(0) if self.mode == "concave": u_grid = scores.min(dim=1).values refreshed = (K - u_grid.unsqueeze(1)).min(dim=0).values else: u_grid = scores.max(dim=1).values refreshed = (K - u_grid.unsqueeze(1)).max(dim=0).values changed += (refreshed - column).abs().gt(1e-12).sum().item() # Write refreshed column back through parameterization self.intercepts_param.set_column_from_raw_(dim, refreshed) # Cache refreshed intercepts snapshot for idempotence checks self._last_refreshed_intercepts = self.intercepts.detach().clone() return changed def project(self, T: torch.Tensor) -> torch.Tensor: """Clamp tensor entries to [-R+ε, R-ε].""" return torch.clamp(T, -self.radius + self.epsilon, self.radius - self.epsilon) def _kernel_grad_wrt_x(self, x_scalar: torch.Tensor, y_scalar: torch.Tensor) -> torch.Tensor: """Compute ∂k/∂x at (x_scalar, y_scalar) using the continuous kernel.""" with torch.enable_grad(): x_temp = x_scalar.detach().clone().requires_grad_(True) y_temp = y_scalar.detach().clone().reshape(1, 1) x_temp_view = x_temp.reshape(1, 1) k_val = self.kernel_fn(x_temp_view, y_temp).squeeze() dk_dx = torch.autograd.grad(k_val, x_temp, retain_graph=False, create_graph=False)[0] return dk_dx.detach() def _kernel_grad_wrt_y(self, x_scalar: torch.Tensor, y_scalar: torch.Tensor) -> torch.Tensor: """Compute ∂k/∂y at (x_scalar, y_scalar) using the continuous kernel.""" with torch.enable_grad(): y_temp = y_scalar.detach().clone().requires_grad_(True) x_temp = x_scalar.detach().clone().reshape(1, 1) y_temp_view = y_temp.reshape(1, 1) k_val = self.kernel_fn(x_temp, y_temp_view).squeeze() dk_dy = torch.autograd.grad(k_val, y_temp, retain_graph=False, create_graph=False)[0] return dk_dy.detach() def _precompute_kernel_derivatives(self): """ Precompute ∂k/∂x and ∂k/∂y on the (X_grid, Y_grid) lattice for use in cached straight-through gradients. """ device = self.X_grid.device dtype = self.X_grid.dtype with torch.enable_grad(): dx_list = [] for y in self.Y_grid: x_vals = self.X_grid.clone().detach().requires_grad_(True) y_tensor = y.clone().detach().to(device=device, dtype=dtype).unsqueeze(0).unsqueeze(0) vals = self.kernel_fn(x_vals.unsqueeze(-1), y_tensor).squeeze() grad_x = torch.autograd.grad( vals, x_vals, grad_outputs=torch.ones_like(vals), retain_graph=False, create_graph=False )[0] dx_list.append(grad_x.detach()) kernel_dx = torch.stack(dx_list, dim=1) dy_list = [] for x in self.X_grid: y_vals = self.Y_grid.clone().detach().requires_grad_(True) x_tensor = x.clone().detach().to(device=device, dtype=dtype).unsqueeze(0).unsqueeze(0) vals = self.kernel_fn(x_tensor, y_vals.unsqueeze(0)).squeeze() grad_y = torch.autograd.grad( vals, y_vals, grad_outputs=torch.ones_like(vals), retain_graph=False, create_graph=False )[0] dy_list.append(grad_y.detach()) kernel_dy = torch.stack(dy_list, dim=0) return kernel_dx, kernel_dy def _cached_dx(self, x_vals: torch.Tensor, y_vals: torch.Tensor) -> torch.Tensor: """Interpolate cached ∂k/∂x for each (x, y) pair.""" if self.kernel_dx.numel() == 0: raise RuntimeError("Cached gradients requested but kernel_dx not precomputed.") nx = self.X_grid.numel() ny = self.Y_grid.numel() x_pos = (x_vals + self.radius) / self.x_accuracy x_low = torch.floor(x_pos).long().clamp(0, nx - 1) x_high = torch.clamp(x_low + 1, 0, nx - 1) weight = (x_pos - x_low.float()).clamp(0.0, 1.0) y_idx = ((y_vals + self.radius) / self.y_accuracy).round().long().clamp(0, ny - 1) grad_low = self.kernel_dx[x_low, y_idx] grad_high = self.kernel_dx[x_high, y_idx] return grad_low + (grad_high - grad_low) * weight def _cached_dy(self, x_idx: torch.Tensor, y_vals: torch.Tensor) -> torch.Tensor: """Interpolate cached ∂k/∂y for fixed x-grid indices.""" if self.kernel_dy.numel() == 0: raise RuntimeError("Cached gradients requested but kernel_dy not precomputed.") ny = self.Y_grid.numel() y_pos = (y_vals + self.radius) / self.y_accuracy y_low = torch.floor(y_pos).long().clamp(0, ny - 1) y_high = torch.clamp(y_low + 1, 0, ny - 1) weight = (y_pos - y_low.float()).clamp(0.0, 1.0) grad_low = self.kernel_dy[x_idx, y_low] grad_high = self.kernel_dy[x_idx, y_high] return grad_low + (grad_high - grad_low) * weight def _per_dim_forward( self, X: torch.Tensor, along: int = 0, selection_mode: str = "soft", snap_to_grid: bool = True, ): """ Compute per-dimension contribution: convex → max_i k(x_d, y_i) - b_i^d concave → min_i k(x_d, y_i) - b_i^d Notes ----- - snap_to_grid=True (hard mode) rounds x_d to the nearest X_grid point before scoring; False evaluates the continuous kernel directly. - In hard mode with gradients, a straight-through update is used: if cache_gradients=True and precomputed dx exist, it interpolates them; otherwise it backpropagates through per-sample autograd calls. Either way, the forward choice stays discrete and gradients are approximate. """ x_vals = X[:, along] x_vals_clamped = self.project(x_vals) if selection_mode == "hard": if snap_to_grid: x_indices = self.get_indices_for_x(x_vals_clamped.detach()).long() kernel_scores = self.kernel_tensor[x_indices, :] else: kernel_scores = self.kernel_fn( x_vals_clamped.unsqueeze(-1), self.Y_grid.unsqueeze(0), ) else: kernel_scores = self.kernel_fn( x_vals_clamped.unsqueeze(-1), self.Y_grid.unsqueeze(0), ) b = self.intercepts[:, along].unsqueeze(0) scores = kernel_scores - b Y_candidates = self.Y_grid.view(1, -1, 1) if self.mode == "convex": choice, f_x, _ = moded_max(scores, Y_candidates, dim=1, temp=self.temp, mode=selection_mode) else: choice, f_x, _ = moded_min(scores, Y_candidates, dim=1, temp=self.temp, mode=selection_mode) choice = choice.squeeze(-1) if selection_mode == "hard" and x_vals.requires_grad: # Straight-through gradient: reuse cached derivatives if available, otherwise compute on the fly. if self.cache_gradients and self.kernel_dx.numel() > 0: approx_grads = self._cached_dx(x_vals_clamped, choice.detach()) else: approx_grads = [] for xs, ys in zip(x_vals_clamped, choice.detach()): approx_grads.append(self._kernel_grad_wrt_x(xs, ys)) approx_grads = torch.stack(approx_grads, dim=0) f_x = f_x + approx_grads * (x_vals_clamped - x_vals_clamped.detach()) return choice, f_x def forward( self, X: torch.Tensor, selection_mode: str = "soft", snap_to_grid: bool = True, ): """ Evaluate f(X) and argmax/argmin selections on the separable grid. Parameters ---------- X : torch.Tensor Input batch of shape (num_samples, num_dims). selection_mode : {"soft", "hard", "ste"} Soft/Hard/stochastic selection; "hard" uses straight-through gradients. snap_to_grid : bool If True, x-values are rounded to the X grid before scoring (only relevant in "hard"). Returns ------- tuple[torch.Tensor, torch.Tensor] (choice, f_x) with per-dimension choices and summed scores. """ args = [] values = [] for dim in range(self.num_dims): arg, val = self._per_dim_forward( X, along=dim, selection_mode=selection_mode, snap_to_grid=snap_to_grid ) args.append(arg) values.append(val) f_x = torch.stack(values, dim=1).sum(dim=1) choice = torch.stack(args, dim=1) return choice, f_x def sup_transform(self, Z): return self._transform_core(Z, maximize=True) def inf_transform(self, Z): return self._transform_core(Z, maximize=False) def _transform_core(self, Z: torch.Tensor, maximize: bool = True): """ Batch transform across dimensions with optional coarse-to-fine search. Parameters ---------- Z : torch.Tensor Targets of shape (num_samples, num_dims). maximize : bool True for sup-transform, False for inf-transform. Notes ----- If coarse_x_factor/coarse_top_k/coarse_window were provided at init, each per-dimension transform first scores a subsampled X grid, seeds the best few candidates, then refines within the configured window; otherwise it searches the full grid. Setting coarse_x_factor<=1 disables the coarse path. """ num_samples, num_dims = Z.shape device = Z.device dtype = Z.dtype X_points = [] total_values = torch.zeros(num_samples, device=device, dtype=dtype) for d in range(num_dims): x_opt_d, value_d = self._per_dimension_transform_batch(Z[:, d], along=d, maximize=maximize) X_points.append(x_opt_d) total_values = total_values + value_d X_opt = torch.stack(X_points, dim=1) return X_opt.detach(), total_values, True def _per_dimension_transform_batch(self, z_vals: torch.Tensor, along: int, maximize: bool): """ Vectorized per-dimension transform for an entire batch of z values. Gradients: If z_vals.requires_grad, the method adds a straight-through term: - with cache_gradients=True and precomputed kernel_dy, it interpolates cached ∂k/∂y at the chosen (x*, z) pairs; - otherwise it falls back to per-sample autograd on the continuous kernel. Forward choices remain discrete; gradients are approximate. """ device = z_vals.device dtype = z_vals.dtype z_clamped = self.project(z_vals) z_idx = self.get_indices_for_y(z_clamped.detach()).long() kernel_tensor = self.kernel_tensor kernel_vals = kernel_tensor[:, z_idx] # (nx, batch) rows = kernel_tensor b = self.intercepts[:, along] scores = rows - b.unsqueeze(0) if self.mode == "convex": f_subset = scores.max(dim=1).values else: f_subset = scores.min(dim=1).values objective_all = kernel_vals - f_subset.unsqueeze(1) nx, batch = objective_all.shape if not self.use_coarse_search or self.coarse_idx.numel() == 0: if maximize: best_pos = objective_all.argmax(dim=0) else: best_pos = objective_all.argmin(dim=0) best_values = objective_all[best_pos, torch.arange(batch, device=device)] best_idx = best_pos else: # Fully vectorized coarse-to-fine search over the batch coarse_idx = self.coarse_idx coarse_obj = objective_all.index_select(0, coarse_idx) # (n_coarse, batch) k = min(self.coarse_top_k, coarse_idx.numel()) if maximize: _, top_pos = torch.topk(coarse_obj, k=k, dim=0, largest=True) else: _, top_pos = torch.topk(-coarse_obj, k=k, dim=0, largest=True) # Seed indices on the full grid, shape (k, batch) seed_idx = coarse_idx[top_pos] if self.coarse_offsets.numel() == 0: refine_idx = seed_idx.view(-1, batch) else: offsets = self.coarse_offsets.view(1, 1, -1) # (1,1,n_offsets) refine_idx = seed_idx.unsqueeze(-1) + offsets # (k,batch,n_offsets) refine_idx = refine_idx.view(-1, batch) # (k*n_offsets, batch) refine_idx = refine_idx.clamp(0, nx - 1) # Gather candidate objective values for all refined indices candidate_obj = objective_all.gather(0, refine_idx) # (n_candidates, batch) if maximize: best_rows = candidate_obj.argmax(dim=0) else: best_rows = candidate_obj.argmin(dim=0) best_values = candidate_obj[best_rows, torch.arange(batch, device=device)] best_idx = refine_idx[best_rows, torch.arange(batch, device=device)] x_opt = self.X_grid[best_idx].to(device=device, dtype=dtype) values = best_values if z_vals.requires_grad: if self.cache_gradients and self.kernel_dy.numel() > 0: approx_grad = self._cached_dy(best_idx, z_clamped) else: grads = [] for x_opt_val, z_val in zip(x_opt, z_clamped): grads.append(self._kernel_grad_wrt_y(x_opt_val, z_val)) approx_grad = torch.stack(grads, dim=0) values = values + approx_grad * (z_clamped - z_clamped.detach()) return x_opt, values.to(device=device, dtype=dtype) def _objective_on_indices(self, indices, kernel_vals, along): """ Evaluate kernel(x_j, z) - f^d(x_j) for a subset of X-grid indices. Args: indices: 1-D LongTensor with rows of the X grid to evaluate. kernel_vals: precomputed k(x_j, z) for all j. along: dimension index. """ rows = self.kernel_tensor.index_select(0, indices) # (subset, ny) b = self.intercepts[:, along] # (ny,) scores = rows - b.unsqueeze(0) if self.mode == "convex": f_subset = scores.max(dim=1).values else: f_subset = scores.min(dim=1).values return kernel_vals[indices] - f_subset def _build_refine_indices(self, seeds): """ Expand coarse seed indices by the configured window. """ if self.coarse_offsets.numel() == 0: return seeds.unique(sorted=True) neighbors = seeds.unsqueeze(1) + self.coarse_offsets.unsqueeze(0) neighbors = neighbors.reshape(-1) neighbors = neighbors.clamp(0, self.X_grid.numel() - 1) return neighbors.unique(sorted=True) def _coarse_transform(self, kernel_vals, along, maximize): """ Run coarse-to-fine search for transforms: 1. Evaluate objective on subsampled coarse grid. 2. Select top-k coarse seeds. 3. Refine by evaluating all fine-grid points in the window. """ coarse_idx = self.coarse_idx coarse_obj = self._objective_on_indices(coarse_idx, kernel_vals, along) k = min(self.coarse_top_k, coarse_idx.numel()) if maximize: _, top_pos = torch.topk(coarse_obj, k=k, largest=True) else: _, top_pos = torch.topk(-coarse_obj, k=k, largest=True) seed_idx = coarse_idx[top_pos] refine_idx = self._build_refine_indices(seed_idx) obj_subset = self._objective_on_indices(refine_idx, kernel_vals, along) if maximize: best_pos = obj_subset.argmax() else: best_pos = obj_subset.argmin() best_idx = refine_idx[best_pos] value = obj_subset[best_pos] return best_idx, value def get_indices_for_x(self, X): idx = ((X + self.radius) / self.x_accuracy).round().long() nx = self.X_grid.numel() return idx.clamp(0, nx - 1) def get_indices_for_y(self, Y): idx = ((Y + self.radius) / self.y_accuracy).round().long() ny = self.Y_grid.numel() return idx.clamp(0, ny - 1) # ===================================================================== # FiniteModel: unified finitely-convex / finitely-concave representation # ===================================================================== class FiniteModel(nn.Module): """ Unified finite model representing either finitely convex or finitely concave functions. **convex mode**: f(x) = max_j [ k(x, y_j) - b_j ] **concave mode**: f(x) = min_j [ k(x, y_j) - b_j ] Supports: - soft / hard / ste selection (via moded_max and moded_min) - numerical sup/inf transforms - full batching support with *global warm-start* storage indexed by sample_idx """ # ------------------------------------------------------------------ def __init__( self, num_candidates: int, num_dims: int, kernel, # k(x,y) mode: str = "convex", temp: float = 50.0, is_y_parameter: bool = True, y_min: Optional[float] = None, y_max: Optional[float] = None, original_dist_to_bounds: float = 1e-1, is_there_default: bool = False, default_intercept: float = 0.0, sorted_model: bool = False, ): """ Initialize a finite convex/concave model over a discrete candidate set. Parameters ---------- num_candidates : int Number of non-default candidate points (columns in Y_rest). num_dims : int Dimensionality of each candidate/vector. kernel : callable Pairwise kernel k(x, y) used to score candidates. mode : {"convex", "concave"} Whether to take max or min over candidates. temp : float Softmax temperature used in soft/ste selection. is_y_parameter : bool If True, Y is trainable; otherwise treated as a buffer. y_min, y_max : float | None Optional bounds; when provided, Y is rescaled/clamped into [y_min, y_max]. original_dist_to_bounds : float Slack added when only one bound is set to avoid collapsing candidates. is_there_default : bool Include an extra default option with its own intercept. default_intercept : float Intercept value for the default option. sorted_model : bool If True, enforce nondecreasing coordinates per candidate via softplus increments. """ super().__init__() assert mode in ("convex", "concave") self.num_dims = num_dims self.num_candidates = num_candidates self.temp = temp self.mode = mode self.is_y_parameter = is_y_parameter self.is_there_default = is_there_default self.y_min = y_min self.y_max = y_max self.original_dist_to_bounds = original_dist_to_bounds self.sorted_model = sorted_model self.default_intercept = default_intercept self.kernel_fn = kernel # --------------------------------------------------------- # Initialize Y candidates # --------------------------------------------------------- Y_init = torch.randn(1, num_candidates, num_dims) min_val, max_val = Y_init.min(), Y_init.max() if sorted_model: Y_init = self._prepare_sorted_Y_init(Y_init) else: Y_init = self._initialize_Y_init(Y_init, min_val, max_val) if is_y_parameter: self._Y_rest_param = nn.Parameter(Y_init) else: self.register_buffer("_Y_rest_param", Y_init) self.intercept_rest = nn.Parameter(torch.zeros(1, num_candidates)) # Optional default option self.register_buffer("Y0", torch.zeros(1,1,num_dims)) intercept0_value = Y_init.new_full((1, 1), default_intercept) self.register_buffer("intercept0", intercept0_value) # --------------------------------------------------------- # GLOBAL per-datapoint warm start buffers # Will be lazily allocated when sample_idx is first seen. # --------------------------------------------------------- self._warm_X_global = None self._num_global_points = None def _initialize_Y_init(self, Y_init: torch.Tensor, min_val: torch.Tensor, max_val: torch.Tensor) -> torch.Tensor: if (self.y_min is not None) and (self.y_max is not None): centered = Y_init - Y_init.mean() max_abs = centered.abs().max() + 1e-4 mid = 0.5 * (self.y_min + self.y_max) half = 0.5 * (self.y_max - self.y_min) return mid + half * centered / max_abs elif self.y_min is not None: return Y_init - min_val + (self.y_min + self.original_dist_to_bounds) elif self.y_max is not None: return max_val - Y_init + (self.y_max - self.original_dist_to_bounds) return Y_init def _prepare_sorted_Y_init(self, Y_init: torch.Tensor) -> torch.Tensor: """Apply one-time rescale for sorted models so later updates are not clamped.""" # sorted path builds Y via softplus increments + cumsum; rescale once here increments = F.softplus(Y_init) sorted_Y = torch.cumsum(increments, dim=-1) if (self.y_min is not None) or (self.y_max is not None): sorted_Y = self._rescale_to_range(sorted_Y) diffs = torch.cat( [sorted_Y[..., :1], sorted_Y[..., 1:] - sorted_Y[..., :-1]], dim=-1, ) safe_diffs = diffs.clamp_min(1e-8) Y_init = torch.log(torch.expm1(safe_diffs)) return Y_init # ============================================================ # Helpers: Y accessors / full intercepts # ============================================================ @property def Y_rest(self): base = self._Y_rest_param if not self.sorted_model: # Unsorted model: just return the Y return base increments = F.softplus(base) sorted_Y = torch.cumsum(increments, dim=-1) return sorted_Y def _rescale_to_range(self, Y: torch.Tensor) -> torch.Tensor: if self.y_min is None and self.y_max is None: return Y if self.y_min is not None and self.y_max is not None: raw_min = Y.amin(dim=1, keepdim=True) raw_max = Y.amax(dim=1, keepdim=True) span = raw_max - raw_min if (span <= 1e-6).all(): return self._rescale_bounds_clamp(Y) target_span = self.y_max - self.y_min return (Y - raw_min) / span * target_span + self.y_min if self.y_min is not None: raw_min = Y.amin(dim=1, keepdim=True) return (Y + (self.y_min - raw_min)).clamp(min=self.y_min) raw_max = Y.amax(dim=1, keepdim=True) return (Y + (self.y_max - raw_max)).clamp(max=self.y_max) def _rescale_bounds_clamp(self, Y: torch.Tensor) -> torch.Tensor: lower = self.y_min if self.y_min is not None else float("-inf") upper = self.y_max if self.y_max is not None else float("inf") return torch.clamp(Y, min=lower, max=upper) @property def Y_rest_raw(self): return self._Y_rest_param def full_Y(self): if self.is_there_default: return torch.cat([self.Y0, self.Y_rest], dim=1) return self.Y_rest def full_intercept(self): if self.is_there_default: return torch.cat([self.intercept0, self.intercept_rest], dim=1) return self.intercept_rest # ============================================================ # Forward: MAX or MIN over candidates # ============================================================ def forward(self, X: torch.Tensor, selection_mode: str = "soft", already_sorted: bool = False): """ Compute f(x) for x ∈ R^{num_samples × num_dims} using either max_j or min_j of the kernel scores. """ if self.sorted_model and not already_sorted: X = torch.sort(X, dim=-1).values Y = self.full_Y() # (1, num_candidates, num_dims) b = self.full_intercept() # (1, num_candidates) batch_size = getattr(self, "kernel_batch_size", None) def _compute_chunk(X_chunk): scores_chunk = self.kernel_fn(X_chunk[:, None, :], Y) - b if self.mode == "convex": return moded_max(scores_chunk, Y, dim=1, temp=self.temp, mode=selection_mode) return moded_min(scores_chunk, Y, dim=1, temp=self.temp, mode=selection_mode) if batch_size is None or batch_size <= 0 or X.shape[0] <= batch_size: choice, f_x, aux = _compute_chunk(X) else: choices = [] fxs = [] last_aux = {} for start in range(0, X.shape[0], batch_size): end = start + batch_size X_chunk = X[start:end] choice_chunk, f_x_chunk, aux_chunk = _compute_chunk(X_chunk) choices.append(choice_chunk) fxs.append(f_x_chunk) last_aux = aux_chunk choice = torch.cat(choices, dim=0) f_x = torch.cat(fxs, dim=0) aux = last_aux # Diagnostics self._last_weights = aux.get("weights") self._last_idx = aux.get("idx") self._last_mean_max_weight = aux.get("mean_max") self._last_eff_temp = aux.get("eff_temp") return choice, f_x # ============================================================ # PUBLIC sup/inf transforms (batch-aware sample_idx added) # ============================================================ def sup_transform(self, Z, sample_idx=None, **kw): """ Compute sup_x [k(x,z) - f(x)] for each z. Returns: tuple: (X_opt, values, converged) where - X_opt: Optimized positions (detached) - values: Transform values (with gradients) - converged: Boolean indicating if optimization converged """ return self._transform_core(Z, sample_idx, maximize=True, **kw) def inf_transform(self, Z, sample_idx=None, **kw): """ Compute inf_x [k(x,z) - f(x)] for each z. Returns: tuple: (X_opt, values, converged) where - X_opt: Optimized positions (detached) - values: Transform values (with gradients) - converged: Boolean indicating if optimization converged """ return self._transform_core(Z, sample_idx, maximize=False, **kw) # ============================================================ # CORE routine with per-sample warm start via sample_idx # ============================================================ def _transform_core( self, Z: torch.Tensor, sample_idx: Optional[torch.Tensor], maximize: bool, steps: int = 50, lr: float = 1e-3, optimizer: str = "lbfgs", lam: float = 1e-3, tol: float = 1e-3, patience: int = 5, ): """ Main routine for sup_x or inf_x: maximize=True → sup_x (negated inf) maximize=False → inf_x Supports random mini-batching through sample_idx, enabling per-datapoint warm-start reuse. Uses InfConvolution for implicit differentiation. Returns: tuple: (X_opt, values, converged) where - X_opt: Optimized positions (detached) - values: Transform values (with gradients) - converged: Boolean indicating if optimization converged Convergence Detection: - LBFGS: Always considered converged (uses internal line search) - Adam/GD: Converged if |loss[i] - loss[i-1]| < tol for any step i - If max steps reached without meeting tolerance, converged=False Warm starts: If sample_idx is provided, per-sample warm-starts are stored and reused across calls (and grown dynamically as new indices appear); otherwise Z is used as the initialization. Optimizer controls: `optimizer` selects between "lbfgs" and first-order updates; `lam` is the implicit-diff regularizer; `patience` gates early stopping for first-order modes. The returned `converged` flag reflects these criteria only. """ num_samples, num_dims = Z.shape device = Z.device # ------------------------------------------------------------- # Allocate GLOBAL warm-start storage if first call # ------------------------------------------------------------- if sample_idx is not None: max_idx = int(sample_idx.max()) if self._warm_X_global is None: # First time: allocate storage for all possible Y indices self._num_global_points = max_idx + 1 self._warm_X_global = torch.zeros( self._num_global_points, num_dims, device=device ) elif max_idx + 1 > self._num_global_points: # Expand storage if new larger index appears new_size = max_idx + 1 new_tensor = torch.zeros(new_size, num_dims, device=device) new_tensor[:self._num_global_points] = self._warm_X_global self._warm_X_global = new_tensor self._num_global_points = new_size # Extract warm starts for this batch X_init = self._warm_X_global[sample_idx].clone() else: # Fallback: no sample_idx → use Z as initialization X_init = Z.clone() # ------------------------------------------------------------- # Create a wrapper module for f(x) to work with InfConvolution # InfConvolution expects f_net(x) to return scalar value # and x will be a 1D tensor of shape (num_dims,) # ------------------------------------------------------------- class FWrapper(nn.Module): def __init__(self, finite_model, negate=False): super().__init__() self.finite_model = finite_model self.negate = negate def forward(self, x): # x is (num_dims,) - need to add batch dimension # FiniteModel.forward expects (num_samples, num_dims) if x.dim() == 1: x = x.unsqueeze(0) # self.finite_model.forward returns (choice, f_x) _, f_x = self.finite_model.forward(x, selection_mode="soft") # f_x is (num_samples,) or scalar - return scalar result = f_x.squeeze() return -result if self.negate else result # ------------------------------------------------------------- # Define kernel function that handles maximize flag and dimensions # InfConvolution expects K(x, y) where x and y are 1D tensors # # For maximize=True (sup): # sup_x [k(x,z) - f(x)] = -inf_x [-(k(x,z) - f(x))] # = -inf_x [f(x) - k(x,z)] # InfConvolution computes: inf_x [K(x,z) - F(x)] # So we need: K(x,z) = -k(x,z) and F(x) = -f(x) # Then: inf_x [-k(x,z) - (-f(x))] = inf_x [f(x) - k(x,z)] # And negate result: -inf_x [f(x) - k(x,z)] = sup_x [k(x,z) - f(x)] # ------------------------------------------------------------- if maximize: f_wrapper = FWrapper(self, negate=True) def K_wrapper(x, z): # Ensure x and z are 2D for kernel_fn if x.dim() == 1: x = x.unsqueeze(0) if z.dim() == 1: z = z.unsqueeze(0) result = -self.kernel_fn(x, z) return result.squeeze() else: f_wrapper = FWrapper(self, negate=False) def K_wrapper(x, z): # Ensure x and z are 2D for kernel_fn if x.dim() == 1: x = x.unsqueeze(0) if z.dim() == 1: z = z.unsqueeze(0) result = self.kernel_fn(x, z) return result.squeeze() # ------------------------------------------------------------- # Apply InfConvolution for each sample in Z # InfConvolution now returns (g_value, converged, x_star) so we # don't need to re-solve the optimization # ------------------------------------------------------------- values_list = [] converged_list = [] X_opt_list = [] for i in range(num_samples): z_i = Z[i] # 1D tensor of shape (num_dims,) x_init_i = X_init[i] # 1D tensor of shape (num_dims,) # InfConvolution.apply returns (g_value, converged, x_star) # This computes the value with proper gradients via implicit differentiation # and returns the optimal x for warm-starting g_i, conv_i, x_star_i = InfConvolution.apply( z_i, f_wrapper, K_wrapper, x_init_i, steps, lr, optimizer, lam, tol, patience, *list(self.parameters()) ) values_list.append(g_i) converged_list.append(conv_i) X_opt_list.append(x_star_i) values = torch.stack(values_list) converged = all(converged_list) X_opt = torch.stack(X_opt_list) # Negate values if maximize if maximize: values = -values # ------------------------------------------------------------- # Save warm starts # ------------------------------------------------------------- with torch.no_grad(): if sample_idx is not None: self._warm_X_global[sample_idx] = X_opt.detach() return X_opt.detach(), values, converged