"Making Keras Cooler!": Per-Layer Learning Rates and Free Gradient Manipulation

“让Keras更酷一些！”：分层的学习率和自由的梯度 苏剑林 · 2019-03-10 · Machine Learning model Keras gradient learning rate

Flying the flag of "Making Keras Cooler!" once again, to unlock the infinite possibilities of Keras~

Today we'll accomplish two important things with Keras: setting per-layer learning rates, and flexibly manipulating gradients.

First, per-layer learning rates. The use case here is obvious: for instance, when fine-tuning an existing model, sometimes we want to freeze certain layers, but other times we don't want to freeze them entirely — instead we want them to update at a lower learning rate than the other layers. That is exactly the per-layer learning rate requirement. There has been some discussion online about setting per-layer learning rates in Keras, and the conclusion is always that you need to rewrite the optimizer to achieve it. Clearly, this approach is unfriendly both in implementation and in use.

Next is manipulating gradients. The most direct example of gradient manipulation is gradient clipping, i.e., constraining the gradient to a certain range — Keras has this built in. But Keras's built-in clipping is global. What if I want to apply a different clipping scheme to each gradient individually? Or what if I have some other idea for manipulating gradients — how would I implement that? Surely not by rewriting the optimizer again?

This post aims to give the simplest possible solutions to the above problems. more

Per-Layer Learning Rates

For the matter of setting per-layer learning rates, rewriting the optimizer is of course feasible, but it's too cumbersome. If we want a simpler approach, we need a bit of mathematical insight to guide us.

Optimization Under a Parameter Transformation

First, let's consider the update formula for gradient descent:

\begin{equation}\boldsymbol{\theta}_{n+1}=\boldsymbol{\theta}_{n}-\alpha \frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}\label{eq:sgd-1}\end{equation}

where $L$ is the loss function with parameter $\boldsymbol{\theta}$, $\alpha$ is the learning rate, and $\frac{\partial L(\boldsymbol{\theta}_{n})}{\partial \boldsymbol{\theta}_n}$ is the gradient, which we sometimes also write as $\nabla_{\boldsymbol{\theta}} L(\boldsymbol{\theta}_{n})$. The notation is fairly arbitrary — what matters is understanding what it means~

Now consider the transformation $\boldsymbol{\theta}=\lambda \boldsymbol{\phi}$, where $\lambda$ is a fixed scalar and $\boldsymbol{\phi}$ is also a parameter. We now optimize $\boldsymbol{\phi}$, with the corresponding update formula:

\begin{equation}\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}\end{equation}

where the second equality is simply the chain rule. Now multiply both sides by $\lambda$, giving us

\begin{equation}\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}\end{equation}

Comparing $\eqref{eq:sgd-1}$ and $\eqref{eq:sgd-2}$, I think you can already see what I'm getting at:

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

However, for adaptive learning rate optimizers (such as RMSprop, Adam, etc.), things are a bit different, because adaptive learning rates use the gradient (as a denominator) to rescale the learning rate, which cancels out a factor of $\lambda$, so that (I leave the derivation to interested readers)

In adaptive learning rate optimizers such as RMSprop and Adam, if we apply the parameter transformation $\boldsymbol{\theta}=\lambda \boldsymbol{\phi}$, the equivalent effect is that the learning rate changes from $\alpha$ to $\lambda\alpha$.

Sleight of Hand: Adjusting the Learning Rate

With these two conclusions in hand, all we need is a way to implement the parameter transformation, without having to rewrite the optimizer ourselves, in order to set per-layer learning rates.

Implementing a parameter transformation isn't hard either — we already covered the method previously in the post "Making Keras Cooler!": Flexible Outputs and Free Normalization when discussing weight normalization. Because when Keras builds a layer, it actually splits the process into two steps, build and call, we can insert some operations right after build and then call call.

Here is a packaged implementation:

import keras.backend as K

class SetLearningRate:
    """层的一个包装，用来设置当前层的学习率
    """

    def __init__(self, layer, lamb, is_ada=False):
        self.layer = layer
        self.lamb = lamb # 学习率比例
        self.is_ada = is_ada # 是否自适应学习率优化器

    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 # 自适应学习率优化器直接保持lamb比例
                else:
                    lamb = self.lamb**0.5 # SGD（包括动量加速），lamb要开平方
                K.set_value(weight, K.eval(weight) / lamb) # 更改初始化
                setattr(self.layer, key, weight * lamb) # 按比例替换
        return self.layer(inputs)

Usage example:

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

# 默认情况下是x = Embedding(100, 1000, weights=[word_vecs])(x)
# 下面这一句表示：后面将会用自适应学习率优化器，并且Embedding层以总体的十分之一的学习率更新。
# word_vecs是预训练好的词向量
x = SetLearningRate(Embedding(100, 1000, weights=[word_vecs]), 0.1, True)(x)

# 后面部分自己想象了～
x = LSTM(100)(x)

model = Model(x_in, x)
model.compile(loss='mse', optimizer='adam') # 用自适应学习率优化器优化

A few notes:

1. Currently this approach can only be inserted when you're manually writing the code to build a model — it can't be applied to an already-constructed model.
2. If you have pretrained weights, there are two ways to load them. The first is, as in the usage example above, to pass them in via the weights argument when defining the layer. The second is to build the model (with SetLearningRate already inserted in the appropriate places) and then use model.set_weights(weights) to assign the values, where weights is "the pretrained weights of the original model, already divided by $\lambda$ or $\sqrt{\lambda}$ at the positions where SetLearningRate is applied."
3. The second method for loading pretrained weights might sound a bit cryptic, but if you've already understood the principle behind this section, you should be able to see what I mean. Since setting the learning rate is implemented via weight * lamb, the initialization of weight needs to 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 overall learning rate, it is very hard to later change it to 1/5 or some other ratio on top of that. (Of course, if you've truly and thoroughly understood the principle in this section, and also figured out the second method of loading pretrained weights, then there is a way — and at that point I trust you'll be able to work it out yourself).
5. This approach has the above limitations precisely because we don't want to modify or rewrite the optimizer to implement this functionality. If you decide you do want to modify the optimizer yourself, please refer to "Making Keras Cooler!": Niche Custom Optimizers.

Free Gradient Manipulation

In this part, we'll learn how to exert more free control over gradients. This involves modifying the optimizer, but doesn't require rewriting it entirely.

The Structure of Keras Optimizers

To modify an optimizer, we first need to understand the structure of Keras optimizers. We already took a preliminary look at this in "Making Keras Cooler!": Niche Custom Optimizers; let's go over it again now.

The Keras optimizer code is at

https://github.com/keras-team/keras/blob/master/keras/optimizers.py

Take a casual look at any optimizer, and you'll find that to define a custom optimizer, all you need to do is subclass Optimizer and define the get_updates method. But in this post we don't want to create a new optimizer — we just want control over the gradients. As it turns out, the gradients are actually fetched in the parent class Optimizer's get_gradients method:

    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 in this method is where the raw gradients are obtained, and what follows provides two gradient-clipping methods. It's not hard to see that simply overriding the optimizer's get_gradients method lets us perform arbitrary manipulations on the gradients, without affecting the optimizer's update step (i.e., without affecting the get_updates method).

Everything Is an Object: Just Override It

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 get_gradients is a method of the optimizer, but we could equally say it's an attribute (object) of get_gradients — and since it's an attribute, we can simply override it by direct assignment.

Let's give the crudest possible example (a bit of 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)

This example is actually pretty silly — it just zeroes out all the gradients (and then no matter how you try to optimize it, nothing moves...) — but this prank example is already representative enough: you could zero out all gradients, or you could perform any operation you like on the gradients. For instance, 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 some layers? That's also simple — when defining the layer, give it a distinguishable name, and then perform different operations based on the name of params. At this point, I believe it's basically a case of "master one, master all."

The Elegance of Keras

Perhaps in many people's eyes, Keras is just a convenient but rigidly "locked-down" high-level framework, but in my eyes, all I see is its boundless flexibility~

That is an impeccable piece of encapsulation.