Dynamic Temperature Scheduling for Stochastic Neural Networks

In neural network training and inference, temperature is a hyperparameter that controls the sharpness of probability distributions. When applied to a softmax function, the temperature τ modifies the output distribution:

$$\text{softmax}_\tau(z_i) = \frac{e^{z_i / \tau}}{\sum_j e^{z_j / \tau}}$$

• High temperature (τ > 1): Produces a more uniform distribution, encouraging exploration and diverse outputs
• Low temperature (τ < 1): Produces a peaked distribution, favoring high-probability outcomes and more deterministic behavior
• Temperature = 1: Standard softmax behavior with no modification

Why Dynamic Temperature Scheduling Matters:

Rather than using a fixed temperature throughout training or generation, many state-of-the-art techniques dynamically adjust temperature over time. This approach is critical in:

• Text Generation: Starting with higher temperature for creativity, then reducing for coherence
• Gumbel-Softmax Training: Annealing temperature to sharpen discrete approximations during training
• Reinforcement Learning: Balancing exploration vs exploitation over training epochs
• Knowledge Distillation: Gradually transitioning from soft to hard label matching

Your Task:

Implement a temperature scheduler that computes the current temperature value at any training step. Support four common scheduling strategies:

1. Linear Decay: $$T(t) = T_0 - (T_0 - T_{min}) \cdot \frac{t}{T_{total}}$$

The temperature decreases at a constant rate from the initial value to the final value.

2. Exponential Decay: $$T(t) = T_0 \cdot \left(\frac{T_{min}}{T_0}\right)^{\frac{t}{T_{total}}}$$

The temperature decreases multiplicatively, with faster initial decay that gradually slows.

3. Cosine Annealing: $$T(t) = T_{min} + \frac{1}{2}(T_0 - T_{min}) \cdot \left(1 + \cos\left(\frac{\pi \cdot t}{T_{total}}\right)\right)$$

A smooth, wave-like decay inspired by learning rate schedules, popular in transformer training.

4. Constant: $$T(t) = T_0$$

The temperature remains unchanged throughout training (baseline for comparison).

Return the computed temperature value rounded to 2 decimal places.