File size: 2,003 Bytes
97bec8a | 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 | # src/optimizer/models.py
import torch
import torch.nn as nn
import geoopt
from .utils import to_device # Optional: for easy device placement in scripts
class SphereRosenbrockModel(nn.Module):
"""
A simple batched model representing points on the 3-sphere (S^3) parameterized as unit quaternions.
Used primarily for testing optimizers on the stereographically compactified 3D Rosenbrock function,
which creates a challenging landscape with narrow valleys and pole singularities.
Attributes:
q: ManifoldParameter on the Sphere manifold (shape: [num_instances, 4])
"""
def __init__(self, num_instances: int = 32, device: torch.device | None = None):
super().__init__()
# Initialize near the north pole (challenging starting region)
init = torch.randn(num_instances, 4, dtype=torch.float64)
init[..., 0] = 0.95 + 0.05 * torch.randn(num_instances) # w component biased high
init = init / init.norm(dim=-1, keepdim=True) # Project to unit sphere
self.manifold = geoopt.manifolds.Sphere()
self.q = geoopt.ManifoldParameter(init, manifold=self.manifold)
# Optional: move to device immediately if specified
if device is not None:
self.to(device)
def forward(self) -> torch.Tensor:
"""
Forward pass: simply return the quaternion parameters on the sphere.
Returns:
q: Tensor of shape [num_instances, 4]
"""
return self.q
# Future-proof placeholders for additional benchmark models
# ------------------------------------------------------------------
# class StiefelOrthogonalModel(nn.Module):
# """Example: Model with parameters on the Stiefel manifold (orthogonal frames)."""
# ...
#
# class PoincareBallModel(nn.Module):
# """Example: Hyperbolic embedding model on the Poincaré ball."""
# ...
# ------------------------------------------------------------------
|