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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348

import numpy as np
import torch
from torch.optim.lr_scheduler import ExponentialLR, LambdaLR
from typing import Callable, Optional
 
# =====================================================
# Implementation 1: PyTorch Native ExponentialLR
# =====================================================
def create_exponential_scheduler(
    optimizer,
    gamma: float = 0.99,
    last_epoch: int = -1
):
    """
    Standard PyTorch exponential decay.
    
    LR(t) = LR(0) * gamma^t
    
    Args:
        optimizer: Wrapped optimizer
        gamma: Multiplicative factor (0 < gamma < 1)
        last_epoch: For resuming training
    
    Example: gamma=0.99 means 1% reduction per epoch
    """
    return ExponentialLR(optimizer, gamma=gamma, last_epoch=last_epoch)
 
 
# =====================================================
# Implementation 2: Target-Based Exponential Decay
# =====================================================
class TargetExponentialDecay:
    """
    Specify initial and final LR; decay rate computed automatically.
    
    More intuitive than raw gamma: "I want LR to go from 0.1 to 0.0001
    over 100 epochs" rather than "I want gamma = 0.9313".
    """
    
    def __init__(
        self,
        optimizer,
        initial_lr: float,
        final_lr: float,
        total_epochs: int,
        last_epoch: int = -1
    ):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.final_lr = final_lr
        self.total_epochs = total_epochs
        
        # Compute gamma from endpoints
        # final_lr = initial_lr * gamma^total_epochs
        # gamma = (final_lr / initial_lr)^(1/total_epochs)
        self.gamma = (final_lr / initial_lr) ** (1.0 / total_epochs)
        
        self.current_epoch = last_epoch
        
        # Validate parameters
        if not (0 < self.gamma < 1):
            raise ValueError(
                f"Invalid decay: from {initial_lr} to {final_lr} in "
                f"{total_epochs} epochs gives gamma={self.gamma:.6f}"
            )
        
        print(f"TargetExponentialDecay: gamma={self.gamma:.6f}, "
              f"half-life={-np.log(2)/np.log(self.gamma):.1f} epochs")
    
    def get_lr(self) -> float:
        if self.current_epoch < 0:
            return self.initial_lr
        return self.initial_lr * (self.gamma ** self.current_epoch)
    
    def step(self, epoch: Optional[int] = None):
        if epoch is not None:
            self.current_epoch = epoch
        else:
            self.current_epoch += 1
        
        new_lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr
    
    def state_dict(self):
        return {
            'initial_lr': self.initial_lr,
            'final_lr': self.final_lr,
            'total_epochs': self.total_epochs,
            'gamma': self.gamma,
            'current_epoch': self.current_epoch
        }
    
    def load_state_dict(self, state_dict):
        self.initial_lr = state_dict['initial_lr']
        self.final_lr = state_dict['final_lr']
        self.total_epochs = state_dict['total_epochs']
        self.gamma = state_dict['gamma']
        self.current_epoch = state_dict['current_epoch']
 
 
# =====================================================
# Implementation 3: Per-Step Exponential Decay
# =====================================================
class PerStepExponentialDecay:
    """
    Decay learning rate every step (iteration) rather than epoch.
    
    Provides smoother decay curve at the cost of more frequent updates.
    Gamma should be much closer to 1 (e.g., 0.9999 vs 0.95 for per-epoch).
    """
    
    def __init__(
        self,
        optimizer,
        initial_lr: float,
        final_lr: float,
        total_steps: int
    ):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.final_lr = final_lr
        self.total_steps = total_steps
        
        # Per-step gamma for smooth transition
        self.gamma = (final_lr / initial_lr) ** (1.0 / total_steps)
        self.current_step = 0
        
        print(f"PerStepExponentialDecay: gamma={self.gamma:.8f} "
              f"(applied {total_steps} times)")
    
    def step(self):
        """Call after each optimizer.step(), not after epoch."""
        self.current_step += 1
        new_lr = self.initial_lr * (self.gamma ** self.current_step)
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr
    
    def get_lr(self) -> float:
        return self.initial_lr * (self.gamma ** self.current_step)
 
 
# =====================================================
# Implementation 4: Exponential Decay with Warmup
# =====================================================
class WarmupExponentialDecay:
    """
    Linear warmup followed by exponential decay.
    
    Warmup prevents unstable early training (especially with
    large batches or transformer architectures).
    """
    
    def __init__(
        self,
        optimizer,
        base_lr: float,
        warmup_epochs: int,
        decay_epochs: int,
        final_lr: float,
        warmup_start_lr: float = 1e-7
    ):
        self.optimizer = optimizer
        self.base_lr = base_lr
        self.warmup_epochs = warmup_epochs
        self.decay_epochs = decay_epochs
        self.final_lr = final_lr
        self.warmup_start_lr = warmup_start_lr
        
        # Compute decay gamma for post-warmup period
        self.gamma = (final_lr / base_lr) ** (1.0 / decay_epochs)
        self.current_epoch = 0
        
        self.total_epochs = warmup_epochs + decay_epochs
    
    def get_lr(self) -> float:
        if self.current_epoch < self.warmup_epochs:
            # Linear warmup
            alpha = self.current_epoch / self.warmup_epochs
            return self.warmup_start_lr + alpha * (self.base_lr - self.warmup_start_lr)
        else:
            # Exponential decay
            decay_step = self.current_epoch - self.warmup_epochs
            return self.base_lr * (self.gamma ** decay_step)
    
    def step(self, epoch: Optional[int] = None):
        if epoch is not None:
            self.current_epoch = epoch
        else:
            self.current_epoch += 1
        
        new_lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr
 
 
# =====================================================
# Implementation 5: Lambda-Based Flexible Schedule
# =====================================================
def create_lambda_exponential(
    optimizer,
    gamma: float = 0.99,
    warmup_epochs: int = 0,
    min_lr_ratio: float = 0.001,
    last_epoch: int = -1
):
    """
    Use LambdaLR for maximum flexibility.
    
    LambdaLR takes a function that maps epoch -> lr_multiplier.
    This enables combining exponential decay with warmup and floors.
    """
    
    def lr_lambda(epoch: int) -> float:
        if epoch < warmup_epochs:
            # Linear warmup
            return epoch / warmup_epochs
        else:
            # Exponential decay with floor
            decay_factor = gamma ** (epoch - warmup_epochs)
            return max(decay_factor, min_lr_ratio)
    
    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
 
 
# =====================================================
# Implementation 6: Staircase Exponential (Hybrid)
# =====================================================
class StaircaseExponential:
    """
    Apply exponential decay but update LR only at fixed intervals.
    
    Combines smoothness of exponential with interpretability of step.
    LR is computed exponentially but applied in discrete steps.
    
    Example: Compute LR exponentially, but update every 5 epochs.
    """
    
    def __init__(
        self,
        optimizer,
        initial_lr: float,
        gamma: float,
        staircase_interval: int = 10
    ):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.gamma = gamma
        self.staircase_interval = staircase_interval
        self.current_epoch = 0
        self.last_applied_lr = initial_lr
    
    def get_lr(self) -> float:
        # Compute exponential value at nearest staircase step
        staircase_epoch = (self.current_epoch // self.staircase_interval) * self.staircase_interval
        return self.initial_lr * (self.gamma ** staircase_epoch)
    
    def step(self, epoch: Optional[int] = None):
        if epoch is not None:
            self.current_epoch = epoch
        else:
            self.current_epoch += 1
        
        new_lr = self.get_lr()
        
        # Only log when LR actually changes
        if abs(new_lr - self.last_applied_lr) > 1e-10:
            print(f"Epoch {self.current_epoch}: LR -> {new_lr:.6f}")
            self.last_applied_lr = new_lr
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr
 
 
# =====================================================
# Visualization and Analysis Tools
# =====================================================
def compare_exponential_schedules(total_epochs: int = 100):
    """
    Visualize different exponential decay configurations.
    """
    import matplotlib.pyplot as plt
    
    epochs = np.arange(total_epochs)
    initial_lr = 0.1
    
    schedules = {
        'γ=0.99 (Slow)': initial_lr * (0.99 ** epochs),
        'γ=0.95 (Medium)': initial_lr * (0.95 ** epochs),
        'γ=0.9 (Fast)': initial_lr * (0.9 ** epochs),
        'Target: 0.1→0.001': initial_lr * ((0.001/0.1)**(epochs/total_epochs)),
    }
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # Linear scale
    ax1 = axes[0]
    for name, lr_curve in schedules.items():
        ax1.plot(epochs, lr_curve, label=name, linewidth=2)
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Learning Rate')
    ax1.set_title('Exponential Decay Schedules (Linear Scale)')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Log scale (shows exponential is linear)
    ax2 = axes[1]
    for name, lr_curve in schedules.items():
        ax2.semilogy(epochs, lr_curve, label=name, linewidth=2)
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Learning Rate (log scale)')
    ax2.set_title('Exponential Decay Schedules (Log Scale)')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig
 
 
def analyze_decay_rate(
    initial_lr: float,
    final_lr: float,
    total_epochs: int
) -> dict:
    """
    Analyze properties of exponential decay for given parameters.
    """
    gamma = (final_lr / initial_lr) ** (1.0 / total_epochs)
    half_life = -np.log(2) / np.log(gamma)
    
    # Compute LR at key points
    key_epochs = [0, total_epochs//4, total_epochs//2, 
                  3*total_epochs//4, total_epochs-1]
    
    lr_values = {
        epoch: initial_lr * (gamma ** epoch) 
        for epoch in key_epochs
    }
    
    return {
        'gamma': gamma,
        'lambda': -np.log(gamma),
        'half_life_epochs': half_life,
        'decay_per_epoch_pct': (1 - gamma) * 100,
        'lr_at_epochs': lr_values,
        'total_decay_factor': final_lr / initial_lr
    }

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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

import numpy as np
 
def calculate_decay_params(
    initial_lr: float,
    final_lr: float,
    total_epochs: int
) -> dict:
    """
    Calculate all decay parameters from initial/final LR goals.
    
    This is the recommended approach: specify endpoints, compute decay.
    """
    
    # Compute gamma
    gamma = (final_lr / initial_lr) ** (1.0 / total_epochs)
    
    # Compute lambda (continuous decay rate)
    lambda_rate = -np.log(gamma)
    
    # Compute half-life
    half_life = np.log(2) / lambda_rate
    
    # Decay per epoch in percentage
    decay_per_epoch = (1 - gamma) * 100
    
    # Time to reach 1% of initial (useful benchmark)
    time_to_1pct = np.log(0.01) / np.log(gamma)
    
    return {
        'gamma': gamma,
        'lambda': lambda_rate,
        'half_life_epochs': half_life,
        'decay_per_epoch_pct': decay_per_epoch,
        'epochs_to_1pct': time_to_1pct,
        'summary': f"LR halves every {half_life:.1f} epochs, "
                  f"decays {decay_per_epoch:.2f}% per epoch"
    }
 
 
def gamma_from_half_life(half_life_epochs: float) -> float:
    """
    Compute gamma that achieves specified half-life.
    """
    return np.exp(-np.log(2) / half_life_epochs)
 
 
def validate_schedule(
    initial_lr: float,
    gamma: float,
    total_epochs: int,
    min_acceptable_lr: float = 1e-8
) -> dict:
    """
    Validate an exponential schedule for common issues.
    """
    final_lr = initial_lr * (gamma ** total_epochs)
    half_life = -np.log(2) / np.log(gamma)
    
    issues = []
    
    if final_lr < min_acceptable_lr:
        issues.append(
            f"Final LR ({final_lr:.2e}) below minimum threshold "
            f"({min_acceptable_lr}). May cause numerical issues."
        )
    
    if half_life > total_epochs:
        issues.append(
            f"Half-life ({half_life:.1f}) exceeds total epochs ({total_epochs}). "
            "LR won't reduce by even 50%."
        )
    
    if half_life < total_epochs / 10:
        issues.append(
            f"Half-life ({half_life:.1f}) very short relative to training. "
            "LR may decay too quickly."
        )
    
    # Check for reasonable number of halvings
    num_halvings = total_epochs / half_life
    if num_halvings < 1:
        issues.append("Less than one full halving during training - consider faster decay.")
    elif num_halvings > 10:
        issues.append("More than 10 halvings - LR may become negligibly small.")
    
    return {
        'valid': len(issues) == 0,
        'issues': issues,
        'final_lr': final_lr,
        'half_life': half_life,
        'num_halvings': num_halvings
    }
 
 
# Interactive exploration
if __name__ == "__main__":
    # Example: Design schedule for 150-epoch training
    result = calculate_decay_params(
        initial_lr=0.1,
        final_lr=0.0001,  # 1/1000 of initial
        total_epochs=150
    )
    
    print("Schedule Design:")
    for key, value in result.items():
        print(f"  {key}: {value}")
    
    # Validation
    validation = validate_schedule(
        initial_lr=0.1,
        gamma=result['gamma'],
        total_epochs=150
    )
    
    print("
Validation:")
    print(f"  Valid: {validation['valid']}")
    for issue in validation['issues']:
        print(f"  ⚠️ {issue}")

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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

import torch
import torch.distributed as dist
from typing import Optional
import json
 
class ProductionExponentialDecay:
    """
    Production-ready exponential decay with:
    - Minimum LR floor
    - Warmup support
    - Checkpointing
    - Distributed training compatibility
    - Logging hooks
    """
    
    def __init__(
        self,
        optimizer,
        initial_lr: float,
        gamma: float,
        min_lr: float = 1e-7,
        warmup_epochs: int = 0,
        warmup_start_lr: float = 1e-7,
        log_lr_every: int = 1,
        distributed: bool = False
    ):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.gamma = gamma
        self.min_lr = min_lr
        self.warmup_epochs = warmup_epochs
        self.warmup_start_lr = warmup_start_lr
        self.log_lr_every = log_lr_every
        self.distributed = distributed
        
        self.current_epoch = 0
        self.current_lr = warmup_start_lr if warmup_epochs > 0 else initial_lr
        
        # Apply initial LR
        self._apply_lr(self.current_lr)
        
        # Compute schedule statistics
        self._log_schedule_info()
    
    def _log_schedule_info(self):
        """Log schedule configuration for debugging."""
        half_life = -np.log(2) / np.log(self.gamma)
        print(f"ExponentialDecay Schedule:")
        print(f"  Initial LR: {self.initial_lr}")
        print(f"  Gamma: {self.gamma}")
        print(f"  Half-life: {half_life:.1f} epochs")
        print(f"  Min LR: {self.min_lr}")
        print(f"  Warmup: {self.warmup_epochs} epochs")
    
    def _compute_lr(self, epoch: int) -> float:
        """Compute learning rate for given epoch."""
        if epoch < self.warmup_epochs:
            # Linear warmup
            alpha = epoch / self.warmup_epochs
            return self.warmup_start_lr + alpha * (self.initial_lr - self.warmup_start_lr)
        else:
            # Exponential decay with floor
            decay_epoch = epoch - self.warmup_epochs
            lr = self.initial_lr * (self.gamma ** decay_epoch)
            return max(lr, self.min_lr)
    
    def _apply_lr(self, lr: float):
        """Apply learning rate to all parameter groups."""
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
    
    def step(self, epoch: Optional[int] = None):
        """
        Advance scheduler by one epoch.
        
        In distributed mode, ensures all ranks use same LR.
        """
        if epoch is not None:
            self.current_epoch = epoch
        else:
            self.current_epoch += 1
        
        self.current_lr = self._compute_lr(self.current_epoch)
        
        # In distributed mode, broadcast LR from rank 0
        if self.distributed and dist.is_initialized():
            lr_tensor = torch.tensor([self.current_lr], dtype=torch.float32)
            dist.broadcast(lr_tensor, src=0)
            self.current_lr = lr_tensor.item()
        
        self._apply_lr(self.current_lr)
        
        # Logging
        if self.current_epoch % self.log_lr_every == 0:
            print(f"Epoch {self.current_epoch}: LR = {self.current_lr:.2e}")
    
    def get_lr(self) -> float:
        return self.current_lr
    
    def state_dict(self) -> dict:
        """Return scheduler state for checkpointing."""
        return {
            'current_epoch': self.current_epoch,
            'current_lr': self.current_lr,
            'initial_lr': self.initial_lr,
            'gamma': self.gamma,
            'min_lr': self.min_lr,
            'warmup_epochs': self.warmup_epochs,
            'warmup_start_lr': self.warmup_start_lr,
        }
    
    def load_state_dict(self, state_dict: dict):
        """Load scheduler state from checkpoint."""
        self.current_epoch = state_dict['current_epoch']
        self.current_lr = state_dict['current_lr']
        self.initial_lr = state_dict['initial_lr']
        self.gamma = state_dict['gamma']
        self.min_lr = state_dict['min_lr']
        self.warmup_epochs = state_dict['warmup_epochs']
        self.warmup_start_lr = state_dict['warmup_start_lr']
        
        # Apply loaded LR
        self._apply_lr(self.current_lr)
        print(f"Loaded scheduler: epoch={self.current_epoch}, lr={self.current_lr:.2e}")
 
 
def save_training_state(
    model, optimizer, scheduler, 
    epoch: int, loss: float, path: str
):
    """Proper checkpoint saving with scheduler state."""
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': loss,
    }, path)
    print(f"Checkpoint saved: {path}")
 
 
def load_training_state(
    model, optimizer, scheduler, path: str
) -> int:
    """Load checkpoint and return starting epoch."""
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    
    start_epoch = checkpoint['epoch'] + 1
    print(f"Resumed from epoch {checkpoint['epoch']}, loss={checkpoint['loss']:.4f}")
    return start_epoch