``Make Keras Cooler!'': Layer-wise Learning Rates and Free Gradients

Optimization under Parameter Transformation

First, let’s consider the update formula for Stochastic Gradient Descent (SGD): \boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n}-\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\label{eq:sgd-1} where L is the loss function with parameters \boldsymbol{\theta}, \alpha is the learning rate, and \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n} is the gradient, sometimes written as \nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n}). Notation is flexible; the key is to understand its meaning.

Now, consider the transformation \boldsymbol{\theta}=\lambda \boldsymbol{\phi}, where \lambda is a fixed scalar and \boldsymbol{\phi} is the parameter. Now we optimize \boldsymbol{\phi}, and the corresponding update formula is: \begin{aligned}\boldsymbol{\phi}_{n+1}=&\boldsymbol{\phi}_{n}-\alpha \frac{\partial L(\lambda\boldsymbol{\phi}_{n})}{\partial \boldsymbol{\phi}_n}\\ =&\boldsymbol{\phi}_{n}-\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\frac{\partial \boldsymbol{\theta}_{n}}{\partial \boldsymbol{\phi}_n}\\ =&\boldsymbol{\phi}_{n}-\lambda\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\end{aligned} The second equality is simply the chain rule. Now, multiplying both sides by \lambda, we get: \lambda\boldsymbol{\phi}_{n+1}=\lambda\boldsymbol{\phi}_{n}-\lambda^2\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\quad\Rightarrow\quad\boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n}-\lambda^2\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\label{eq:sgd-2} Comparing [eq:sgd-1] and [eq:sgd-2], you should see what I am getting at:

In the SGD optimizer, if we perform the parameter transformation \boldsymbol{\theta}=\lambda \boldsymbol{\phi}, the equivalent result is that the learning rate changes from \alpha to \lambda^2\alpha.

However, in adaptive learning rate optimizers (such as RMSprop, Adam, etc.), the situation is slightly different because the adaptive learning rate uses the gradient (as the denominator) to adjust the learning rate, which cancels out one \lambda. Thus (interested readers are encouraged to derive this themselves):

In adaptive learning rate optimizers like RMSprop and Adam, if we perform the parameter transformation \boldsymbol{\theta}=\lambda \boldsymbol{\phi}, the equivalent result is that the learning rate changes from \alpha to \lambda\alpha.

Adjusting Learning Rates via Grafting

With these two conclusions, we only need to find a way to implement parameter transformation without rewriting the optimizer to achieve layer-wise learning rates.

Implementing parameter transformation is not difficult. We previously discussed this in the article “Make Keras Cooler!”: Arbitrary Outputs and Flexible Normalization when talking about Weight Normalization. Since Keras separates the build and call steps when constructing a layer, we can insert some operations after build and before call.

Below is an encapsulated implementation:

import keras.backend as K

class SetLearningRate:
    """A wrapper for layers to set the learning rate of the current layer.
    """

    def __init__(self, layer, lamb, is_ada=False):
        self.layer = layer
        self.lamb = lamb # Learning rate ratio
        self.is_ada = is_ada # Whether it is an adaptive learning rate optimizer

    def __call__(self, inputs):
        with K.name_scope(self.layer.name):
            if not self.layer.built:
                input_shape = K.int_shape(inputs)
                self.layer.build(input_shape)
                self.layer.built = True
                if self.layer._initial_weights is not None:
                    self.layer.set_weights(self.layer._initial_weights)
        for key in ['kernel', 'bias', 'embeddings', 'depthwise_kernel', 'pointwise_kernel', 'recurrent_kernel', 'gamma', 'beta']:
            if hasattr(self.layer, key):
                weight = getattr(self.layer, key)
                if self.is_ada:
                    lamb = self.lamb # Adaptive optimizers maintain the lamb ratio directly
                else:
                    lamb = self.lamb**0.5 # SGD (including momentum), lamb needs a square root
                K.set_value(weight, K.eval(weight) / lamb) # Change initialization
                setattr(self.layer, key, weight * lamb) # Replace proportionally
        return self.layer(inputs)

Usage example:

x_in = Input(shape=(None,))
x = x_in

# Normally it would be x = Embedding(100, 1000, weights=[word_vecs])(x)
# The following line indicates: an adaptive optimizer will be used later, 
# and the Embedding layer will update at 1/10th of the global learning rate.
# word_vecs are pre-trained word vectors.
x = SetLearningRate(Embedding(100, 1000, weights=[word_vecs]), 0.1, True)(x)

# Imagine the rest of the model...
x = LSTM(100)(x)

model = Model(x_in, x)
model.compile(loss='mse', optimizer='adam') # Optimize with an adaptive optimizer

Several points to note:

1. Currently, this method can only be inserted when building the model manually via code; it cannot be applied to a model that has already been built.

2. If there are pre-trained weights, there are two ways to load them. The first is to pass them through the weights parameter when defining the layer, as shown in the example. The second method is to assign them after the model is built (with SetLearningRate already inserted) using model.set_weights(weights), where weights are the original pre-trained weights divided by \lambda or \sqrt{\lambda} at the SetLearningRate positions.

3. The second method for loading weights might seem confusing, but if you understand the principle of this section, you should know what I mean. Since the learning rate is set via weight * lamb, the initialization of the weight must become weight / lamb.

4. This operation is basically irreversible. For example, if you initially set the Embedding layer to update at 1/10 of the global learning rate, it is difficult to change it to 1/5 or another ratio later. (Of course, if you truly understand the principles of this section and the second weight loading method, there is still a way, and I believe you could figure it out).

5. These limitations exist because we want to avoid modifying or rewriting the optimizer. If you decide to modify the optimizer yourself, please refer to “Make Keras Cooler!”: Niche Custom Optimizers.

Structure of Keras Optimizers

To modify an optimizer, one must first understand its structure. We previously looked at this in “Make Keras Cooler!”: Niche Custom Optimizers; let’s review it again.

The Keras optimizer code is located at: https://github.com/keras-team/keras/blob/master/keras/optimizers.py

By observing any optimizer, you will find that to customize an optimizer, you only need to inherit the Optimizer class and define the get_updates method. But in this article, we don’t want to create a new optimizer; we just want to have control over the gradients. As we can see, the acquisition of gradients actually happens in the get_gradients method of the parent class Optimizer:

    def get_gradients(self, loss, params):
        grads = K.gradients(loss, params)
        if None in grads:
            raise ValueError('An operation has `None` for gradient. '
                             'Please make sure that all of your ops have a '
                             'gradient defined (i.e. are differentiable). '
                             'Common ops without gradient: '
                             'K.argmax, K.round, K.eval.')
        if hasattr(self, 'clipnorm') and self.clipnorm > 0:
            norm = K.sqrt(sum([K.sum(K.square(g)) for g in grads]))
            grads = [clip_norm(g, self.clipnorm, norm) for g in grads]
        if hasattr(self, 'clipvalue') and self.clipvalue > 0:
            grads = [K.clip(g, -self.clipvalue, self.clipvalue) for g in grads]
        return grads

The first line of the method is for obtaining the raw gradients, followed by two gradient clipping methods. It is not hard to imagine that by simply overriding the get_gradients method of the optimizer, one can implement any operation on the gradients, and this operation does not affect the update steps of the optimizer (i.e., it does not affect the get_updates method).

Everything is an Object: Just Override

How can we modify only the get_gradients method? This is thanks to Python’s philosophy—“Everything is an object.” Python is an object-oriented programming language, and almost every variable you encounter in Python is an object. We say that get_gradients is a method of the optimizer, but we can also say it is an attribute (object) of the optimizer. Since it is an attribute, we can simply overwrite it with a new assignment.

Let’s look at a blunt example (a prank):

def our_get_gradients(loss, params):
    return [K.zeros_like(p) for p in params]

adam_opt = Adam(1e-3)
adam_opt.get_gradients = our_get_gradients

model.compile(loss='categorical_crossentropy',
              optimizer=adam_opt)

Actually, this example is quite boring—it just sets all gradients to zero (so no matter how you optimize, the model won’t move...). But this prank example is representative enough: you can set all gradients to zero, or you can perform any operation you like on the gradients. For example, clipping gradients according to the l_1 norm instead of the l_2 norm, or making other adjustments.

What if I only want to manipulate the gradients of certain layers? That’s also simple. When defining a layer, give it a distinguishable name, and then perform different operations based on the names in params. By this stage, I believe the principle of “one method opens all doors” applies.