Temperature Decay Scheduler

#231 · Deep Learning · Medium

Problem

Implement a temperature decay scheduler for training. The temperature starts at a high value and decays over training steps according to a specified schedule (e.g., linear, exponential, or cosine decay).

Solution

import math

def temperature_decay(
    initial_temp: float,
    final_temp: float,
    current_step: int,
    total_steps: int,
    schedule: str = "cosine",
) -> float:
    if total_steps <= 0 or current_step >= total_steps:
        return final_temp
    if current_step <= 0:
        return initial_temp

    progress = current_step / total_steps

    if schedule == "linear":
        temp = initial_temp + (final_temp - initial_temp) * progress
    elif schedule == "exponential":
        temp = initial_temp * (final_temp / initial_temp) ** progress
    elif schedule == "cosine":
        temp = final_temp + 0.5 * (initial_temp - final_temp) * (1 + math.cos(math.pi * progress))
    else:
        temp = initial_temp

    return round(temp, 6)

Explanation

Linear decay: Temperature decreases linearly from initial to final over total steps.
Exponential decay: Temperature decays geometrically, T = T_init * (T_final / T_init)^(step / total).
Cosine decay: Uses a cosine annealing schedule for smooth decay: T = T_final + 0.5 * (T_init - T_final) * (1 + cos(pi * progress)).
The scheduler clamps to the final temperature once all steps are complete.

Complexity

Time: O(1)
Space: O(1)

← #230 #232 →