I'm currently fine-tuning a T5-base model for a complex summarization task and I’m debating between a standard linear learning rate decay and a cosine annealing scheduler. My initial tests show that while linear decay is stable, the convergence seems to plateau early. I've read that cosine scheduling allows for higher learning rates for longer periods before a sharp final drop. What has been your experience with T5 specifically? Does the "warm restart" variant help in escaping local minima during text-to-text generation training?
3 answers
In my experience training T5 and other encoder-decoder architectures, cosine learning rate scheduling significantly outperforms linear decay in the final 20% of the training budget. The primary impact on convergence speed comes from the "gentle" initial slope of the cosine curve, which prevents the model from aggressive updates that can destabilize the early attention layers. Unlike linear decay, which reduces the rate at a constant pace, the cosine function keeps the learning rate relatively high during the middle of the training, allowing the T5 model to cover more ground in the loss landscape. By the time it reaches the final stages, the rapid decay helps "settle" the weights into a sharper local optimum, often resulting in a lower cross-entropy loss and better ROUGE scores compared to a linear schedule.
That is a very detailed breakdown, Kimberly. However, have you noticed if the cosine scheduler requires a longer warmup period for T5? I’ve found that if I don’t extend the linear warmup to at least 10% of the total steps, the cosine decay starts too early and the model fails to learn the basic sequence-to-sequence mappings.
I’ve had great results using CosineAnnealingWarmRestarts. The periodic resets help the T5 model discover different semantic interpretations, which is particularly useful for creative summarization tasks where the loss landscape is quite bumpy.
I agree with Charles. Warm restarts act like a regularizer for T5. One tip: make sure to increase the cycle length ($T_{mult}$) after each restart so the model has enough time to actually converge in the later stages of training.
You hit on a crucial point, Kimberly. For T5 specifically, since it’s pre-trained with a constant-then-inverse-square-root schedule, switching to cosine during fine-tuning can be tricky. I usually find that a 10% warmup is the minimum safe zone. If you go lower, the gradient spikes in the T5 decoder can cause early divergence. I'd recommend matching your peak learning rate to about $1 \times 10^{-4}$ when using cosine annealing to ensure the model doesn't "jump" out of a good region before the decay starts.