Layer-wise Learning Rate Decay

The problem, from first principles

A pretrained network's early layers learn general features (edges, textures, word pieces) that transfer to almost any task, while its last layers are specific to the original task. When fine-tuning on new data it is therefore wise to update the early layers gently and the later layers boldly. Layer-wise learning-rate decay formalizes this: the top layer gets the base learning rate, and each layer below it gets the rate of the layer above multiplied by a decay factor d < 1. The result is a different learning rate per parameter group, which optimizers support directly. Parameter names carry the layer index, so the group can be recovered from the name.

From theory to code

Implement layer_index and layerwise_lr.

Constraints

• Parameters are named embed.<...>, layers.<i>.<...> or head.<...>. With n_layers blocks, assign depth indices: embed is 0, layers.i is i + 1, head is n_layers + 1. layer_index(name, n_layers) returns this index. Any other name raises ValueError.
• layerwise_lr(name, n_layers, base_lr, decay) returns base_lr * decay ** (n_layers + 1 - layer_index(name, n_layers)), so the head gets base_lr and every step down multiplies by decay.

Hints