| """ |
| Learning rate schedulers for gradient ascent optimization. |
| |
| Provides various LR scheduling strategies for reward-guided gradient ascent, |
| including cosine annealing, linear decay, and custom schedules. |
| """ |
|
|
| import math |
| from typing import Optional, Literal |
|
|
|
|
| class LRScheduler: |
| """Base class for learning rate schedulers.""" |
| |
| def __init__(self, initial_lr: float, num_steps: int): |
| """ |
| Initialize LR scheduler. |
| |
| Args: |
| initial_lr: Initial learning rate |
| num_steps: Total number of optimization steps |
| """ |
| self.initial_lr = initial_lr |
| self.num_steps = num_steps |
| self.current_step = 0 |
| |
| def get_lr(self) -> float: |
| """Get current learning rate.""" |
| raise NotImplementedError |
| |
| def step(self): |
| """Update scheduler state after a step.""" |
| self.current_step += 1 |
| |
| def reset(self): |
| """Reset scheduler state.""" |
| self.current_step = 0 |
|
|
|
|
| class ConstantLR(LRScheduler): |
| """Constant learning rate (no scheduling).""" |
| |
| def get_lr(self) -> float: |
| return self.initial_lr |
|
|
|
|
| class LinearLR(LRScheduler): |
| """Linear learning rate decay.""" |
| |
| def __init__( |
| self, |
| initial_lr: float, |
| num_steps: int, |
| end_lr: float = 0.0, |
| start_step: int = 0, |
| ): |
| """ |
| Initialize linear LR scheduler. |
| |
| Args: |
| initial_lr: Starting learning rate |
| num_steps: Total number of steps |
| end_lr: Ending learning rate (default: 0.0) |
| start_step: Step to begin decay (default: 0) |
| """ |
| super().__init__(initial_lr, num_steps) |
| self.end_lr = end_lr |
| self.start_step = start_step |
| |
| def get_lr(self) -> float: |
| if self.current_step < self.start_step: |
| return self.initial_lr |
| |
| progress = (self.current_step - self.start_step) / (self.num_steps - self.start_step) |
| progress = min(1.0, progress) |
| |
| return self.initial_lr + (self.end_lr - self.initial_lr) * progress |
|
|
|
|
| class CosineLR(LRScheduler): |
| """Cosine annealing learning rate schedule.""" |
| |
| def __init__( |
| self, |
| initial_lr: float, |
| num_steps: int, |
| min_lr: float = 0.0, |
| warmup_steps: int = 0, |
| ): |
| """ |
| Initialize cosine LR scheduler. |
| |
| Args: |
| initial_lr: Maximum learning rate |
| num_steps: Total number of steps |
| min_lr: Minimum learning rate (default: 0.0) |
| warmup_steps: Number of linear warmup steps (default: 0) |
| """ |
| super().__init__(initial_lr, num_steps) |
| self.min_lr = min_lr |
| self.warmup_steps = warmup_steps |
| |
| def get_lr(self) -> float: |
| if self.current_step < self.warmup_steps: |
| |
| return self.initial_lr * (self.current_step / self.warmup_steps) |
| |
| |
| progress = (self.current_step - self.warmup_steps) / (self.num_steps - self.warmup_steps) |
| progress = min(1.0, progress) |
| |
| cosine_decay = 0.5 * (1 + math.cos(math.pi * progress)) |
| return self.min_lr + (self.initial_lr - self.min_lr) * cosine_decay |
|
|
|
|
| class ExponentialLR(LRScheduler): |
| """Exponential learning rate decay.""" |
| |
| def __init__( |
| self, |
| initial_lr: float, |
| num_steps: int, |
| gamma: float = 0.95, |
| ): |
| """ |
| Initialize exponential LR scheduler. |
| |
| Args: |
| initial_lr: Starting learning rate |
| num_steps: Total number of steps |
| gamma: Multiplicative decay factor per step |
| """ |
| super().__init__(initial_lr, num_steps) |
| self.gamma = gamma |
| |
| def get_lr(self) -> float: |
| return self.initial_lr * (self.gamma ** self.current_step) |
|
|
|
|
| class StepLR(LRScheduler): |
| """Step-wise learning rate decay.""" |
| |
| def __init__( |
| self, |
| initial_lr: float, |
| num_steps: int, |
| step_size: int, |
| gamma: float = 0.1, |
| ): |
| """ |
| Initialize step LR scheduler. |
| |
| Args: |
| initial_lr: Starting learning rate |
| num_steps: Total number of steps |
| step_size: Number of steps between each decay |
| gamma: Multiplicative decay factor |
| """ |
| super().__init__(initial_lr, num_steps) |
| self.step_size = step_size |
| self.gamma = gamma |
| |
| def get_lr(self) -> float: |
| num_decays = self.current_step // self.step_size |
| return self.initial_lr * (self.gamma ** num_decays) |
|
|
|
|
| def create_lr_scheduler( |
| scheduler_type: Literal["constant", "linear", "cosine", "exponential", "step"], |
| initial_lr: float, |
| num_steps: int, |
| **kwargs |
| ) -> LRScheduler: |
| """ |
| Factory function to create learning rate schedulers. |
| |
| Args: |
| scheduler_type: Type of scheduler ("constant", "linear", "cosine", "exponential", "step") |
| initial_lr: Initial learning rate |
| num_steps: Total number of optimization steps |
| **kwargs: Additional scheduler-specific arguments |
| For linear: end_lr, start_step |
| For cosine: min_lr, warmup_steps |
| For exponential: gamma |
| For step: step_size, gamma |
| |
| Returns: |
| LRScheduler instance |
| |
| Examples: |
| # Constant LR |
| scheduler = create_lr_scheduler("constant", initial_lr=0.1, num_steps=100) |
| |
| # Linear decay |
| scheduler = create_lr_scheduler("linear", initial_lr=0.1, num_steps=100, end_lr=0.01) |
| |
| # Cosine annealing with warmup |
| scheduler = create_lr_scheduler("cosine", initial_lr=0.1, num_steps=100, |
| min_lr=0.001, warmup_steps=10) |
| """ |
| if scheduler_type == "constant": |
| return ConstantLR(initial_lr, num_steps) |
| |
| elif scheduler_type == "linear": |
| return LinearLR( |
| initial_lr, num_steps, |
| end_lr=kwargs.get("end_lr", 0.0), |
| start_step=kwargs.get("start_step", 0), |
| ) |
| |
| elif scheduler_type == "cosine": |
| return CosineLR( |
| initial_lr, num_steps, |
| min_lr=kwargs.get("min_lr", 0.0), |
| warmup_steps=kwargs.get("warmup_steps", 0), |
| ) |
| |
| elif scheduler_type == "exponential": |
| return ExponentialLR( |
| initial_lr, num_steps, |
| gamma=kwargs.get("gamma", 0.95), |
| ) |
| |
| elif scheduler_type == "step": |
| return StepLR( |
| initial_lr, num_steps, |
| step_size=kwargs.get("step_size", 10), |
| gamma=kwargs.get("gamma", 0.1), |
| ) |
| |
| else: |
| raise ValueError(f"Unknown scheduler type: {scheduler_type}. " |
| f"Choose from: constant, linear, cosine, exponential, step") |
|
|
