Learning Rates and Optimizer State in RNN Training

While working on an RNN that generates Shakespearean-style text, two training problems surfaced: a loss function that refused to decrease, and training that stalled after initial progress. Both had straightforward causes rooted in how gradient descent is configured and run.

How Training Works in Brief

Deep learning training can be summarized as a continuous optimization problem:

  • The model combines training data and weights to produce a single number: the loss.
  • Weights are the parameters of every matrix in the model; a 64×64 matrix contributes 4096 weights.
  • For optimization, the loss is treated as a function of the weights, which change during training while the data does not.
  • Training runs gradient descent: compute the derivative of the loss with respect to all weights.
  • That derivative is obtained via the chain rule, applied through an algorithm called "backpropagation."
  • Weights are then adjusted by a multiple of the gradient: parameters -= learning_rate * gradient.
  • The multiplier, the learning rate, is notoriously hard to choose; heuristics like Adam exist to automate good choices.

Too High a Learning Rate Means No Learning

The first failure mode appeared as a flat loss curve. The model never learned, and it took a while to trace the cause back to the learning rate. A rate around 0.01 was simply too aggressive; dropping it to roughly 0.002 immediately allowed progress.

The model began producing text like this:

 erlon, w oller. is. d y ivell iver ave esiheres tligh? e ispeafeink
 teldenauke'envexes. h exinkes ror h. ser. sat ly. spon, exang oighis yn, y
 hire aning is's es itrt. for ineull ul'cl r er. s unt. y ch er e s out twiof
 uranter h measaker h exaw; speclare y towessithisil's  aches? s es, tith s aat

That was a clear improvement over its previous output:

kf ;o 'gen '9k ',nrhna 'v ;3; ;'rph 'g ;o kpr ;3;tavrnad 'ps ;]; ;];oraropr
;9vnotararaelpot ;9vr ;9

Even with the corrected rate, however, training eventually stalled again.

Resetting the Optimizer Breaks Training

The second stall came from a subtler bug. The training loop constructed and reinitialized the optimizer inside the loop, resetting its state periodically. That structure looked like this:

for i in range(something):
    optimizer = torch.optim.Adam(rnn.parameters())
    ... do training things

The mistake was treating the optimizer as stateless. Adam's internal state matters—it gradually adjusts the effective step size as training proceeds. Each reset wiped out that accumulated knowledge, and training could not continue effectively.

The fix was to create the optimizer once, before the training loop, and reuse it throughout. When saving the model, the optimizer's state had to be saved as well:

torch.save({'model_state_dict': rnn.state_dict(), 'optimizer_dict': optimizer.state_dict()}, MODEL_PATH)

Once the optimizer stopped being reset, the generated text improved considerably, even producing recognizable English fragments like "Woman steed and oppet!"

at soerin, I kanth as jow gill fimes, To metes think our wink we in fatching
and, Drose, How the wit? our arpear War, our in wioken alous, To thigh dies wit
stain! navinge a sput pie, thick done a my wiscian. Hark's king, and Evit night
and find. Woman steed and oppet, I diplifire, and evole witk ud

What's Next

The immediate follow-up is to understand "BPTT" (backpropagation through time), which should allow faster training and possibly a larger hidden state than the current 87 parameters. With that in hand, more ambitious models become feasible.