Medium End-to-End

Learning Rate Scheduler

A learning-rate schedule supplies a step size for each training update. Cosine annealing lowers that rate along a smooth cosine curve. Here, build the schedule as an array; do not create an optimizer or train a model.

Your task

Implement cosine_schedule(lr_max, lr_min, T):

lr_max and lr_min are finite floats with 0 <= lr_min <= lr_max.
T is a positive integer giving the number of values to return.

For every integer step t from 0 through T - 1, compute:

lr[t] = lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(pi * t / T))

Use radians for cosine and ordinary division for t / T. Return a one-dimensional float32 tensor of shape (T,), in increasing step order. There is no warmup or restart in this schedule.

The endpoint matters

At t = 0, the formula gives lr_max. It reaches lr_min at t = T, but that step is not included in the returned array. Use T as the denominator, not T - 1, and do not append an extra minimum value.

For T = 1, return just [lr_max]. If lr_min = lr_max, every returned value is that same rate. With distinct bounds, the mathematical schedule decreases toward the minimum; floating-point rounding can make nearby values equal.

Example

For lr_max = 0.1, lr_min = 0 and T = 3, the cosine arguments are 0, pi/3 and 2*pi/3. Their cosines are 1, 0.5 and -0.5, so return approximately [0.1, 0.075, 0.025]. There are exactly three values, and the last is above zero.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Set the learning rate and control gradient size

A learning-rate schedule changes the step size during training. Check its exact indexing: a warmup may start at zero or at a first positive fraction. Test the initial value, the warmup boundary, the decay endpoint, and any later steps. The schedule's name alone does not specify those choices.

Global norm clipping scales all gradients by the same factor:

$$\min\left(1,\frac{c}{\|g\|}\right),$$

where g represents all parameter gradients together and c is the norm limit. For scalar gradients 3 and 4, the global norm is 5. Clipping to 2.5 multiplies both by 0.5, giving [1.5, 2]. Clipping each tensor separately can change their relative direction. Handle zero norm without dividing by zero.

Keep a smoothed copy of the parameters

An exponential moving average (EMA) updates a separate parameter copy:

$$\bar\theta_t=\beta\bar\theta_{t-1}+(1-\beta)\theta_t.$$

For old average 2, new parameter 4, and beta = 0.9, the new average is 0.9*2 + 0.1*4 = 2.2. A larger beta retains more of the past. Initialization and whether EMA runs before or after the optimizer step affect the result.

Use validation to decide when to stop

Evaluate using the intended inference behavior without changing model parameters. For mean per-example loss, weight each batch's contribution by its example count. Giving unequal batches equal weight changes the overall mean.

Early stopping monitors a validation metric and stops after a specified number of observations without improvement. Define ties, the minimum improvement, and patience precisely. Save an independent copy of the best parameters: the final step may not be the best one. Keep the test set out of the stopping decision so it remains an independent evaluation.

Detecting runtime…