"Making Keras Cooler!": Layers within Layers, and Masks

“让Keras更酷一些！”：层中层与mask 苏剑林 · 2019-07-16 · Machine Learning Keras

This installment of "Making Keras Cooler!" covers two topics. The first is "layers within layers" — as the name suggests, this is about reusing existing layers when defining a custom layer in Keras, which can massively cut down on the amount of code you need to write. The second topic, requested by readers, is an introduction to the principles and methods of masking in sequence models.

Layers within layers

In “Making Keras Cooler!”: Elegant Layers and Fancy Callbacks, we already introduced the basic method for defining custom layers in Keras, whose core steps involve defining the two functions build and call, where build is responsible for creating trainable weights and call defines the actual computation.

Refusing to repeat yourself

Anyone who often writes custom layers has probably noticed that we tend to repeat ourselves a lot. For instance, if we want to add a linear transformation, we have to add a kernel and a bias variable inside build (also handling variable initialization, regularization, etc.), and then apply them with K.dot inside call — sometimes we also need to worry about dimension alignment. It's all rather tedious. But in fact, a linear transformation is nothing more than a Dense layer without an activation function. If we could reuse existing layers when writing custom layers, we would obviously save a huge amount of code. more

In fact, if you're reasonably familiar with Python object-oriented programming and study the source code of Keras's Layer carefully, it isn't hard to figure out how to reuse existing layers. Below I've organized this into a fairly standardized workflow for readers' reference.

(Note: starting from Keras 2.3.0, the "layer within a layer" functionality is already built in, so you no longer need the custom OurLayer below — you can just use Layer directly.)

OurLayer

First, we define a new OurLayer class:

class OurLayer(Layer):
    """定义新的Layer，增加reuse方法，允许在定义Layer时调用现成的层
    """
    def reuse(self, layer, *args, **kwargs):
        if not layer.built:
            if len(args) > 0:
                inputs = args[0]
            else:
                inputs = kwargs['inputs']
            if isinstance(inputs, list):
                input_shape = [K.int_shape(x) for x in inputs]
            else:
                input_shape = K.int_shape(inputs)
            layer.build(input_shape)
        outputs = layer.call(*args, **kwargs)
        for w in layer.trainable_weights:
            if w not in self._trainable_weights:
                self._trainable_weights.append(w)
        for w in layer.non_trainable_weights:
            if w not in self._non_trainable_weights:
                self._non_trainable_weights.append(w)
        for u in layer.updates:
            if not hasattr(self, '_updates'):
                self._updates = []
            if u not in self._updates:
                self._updates.append(u)
        return outputs

This OurLayer class inherits from the original Layer class, adding to it a reuse method, which is exactly what lets us reuse existing layers.

Here is a simple example: defining a layer that performs the following computation:

$$y = g(f(xW_1 + b_1)W_2 + b_2)$$

Here $f,g$ is an activation function, which is actually just a composition of two Dense layers. If we followed the standard approach, we would need to define several weights inside build, define their shapes based on the input, define their initializers, and so on — many steps. But in fact, aren't all of these already written for us inside the Dense layer? We can just call it directly, as in the following reference code:

class OurDense(OurLayer):
    """原来是继承Layer类，现在继承OurLayer类
    """
    def __init__(self, hidden_dim, output_dim,
                 hidden_activation='linear',
                 output_activation='linear', **kwargs):
        super(OurDense, self).__init__(**kwargs)
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.hidden_activation = hidden_activation
        self.output_activation = output_activation
    def build(self, input_shape):
        """在build方法里边添加需要重用的层，
        当然也可以像标准写法一样条件可训练的权重。
        """
        super(OurDense, self).build(input_shape)
        self.h_dense = Dense(self.hidden_dim,
                             activation=self.hidden_activation)
        self.o_dense = Dense(self.output_dim,
                             activation=self.output_activation)
    def call(self, inputs):
        """直接reuse一下层，等价于o_dense(h_dense(inputs))
        """
        h = self.reuse(self.h_dense, inputs)
        o = self.reuse(self.o_dense, h)
        return o
    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.output_dim,)

Isn't that so much cleaner?

Mask

In this section we'll discuss padding and masking when dealing with variable-length sequences.

Prove that you've thought about it

In several of the models I've recently open-sourced, masks are used extensively, and quite a few readers who apparently have never encountered this before have flooded me with all sorts of questions. There's nothing wrong with having questions about something new, but the problem is that asking questions without first thinking things through is simply irresponsible. I've always believed that when asking someone a question, you should simultaneously "prove" that you've actually thought about it. For example, if you want me to explain something about masks, I'll first ask you to answer this:

Roughly what does the sequence look like before masking? Which positions in the sequence change after masking? What do they become?

These three questions have nothing to do with the underlying principle of masking — they're just meant to check whether you understand what computation the mask is actually performing. Only once you understand that can we go on to discuss why the computation is done this way. If you can't even understand the computation itself, then you have only two options: give up trying to understand this topic, or go study Keras properly for a few months and then come back and discuss it with me.

Below, I'll assume the reader already understands the mask computation, and we'll briefly discuss the basic principle behind masking.

Excluding padding

Masking goes hand in hand with padding. Since neural network inputs need to be regular tensors, while text is usually of variable length, we need some way of trimming or padding sequences to make them a fixed length. By convention, we usually use 0 as the padding symbol.

Let's describe the principle of padding using a simple vector. Suppose we have a vector of length 5:

$$x = [1, 0, 3, 4, 5]$$

After padding it becomes a vector of length 8:

$$x = [1, 0, 3, 4, 5, 0, 0, 0]$$

When you feed this length-8 vector into a model, the model has no way of knowing whether this is truly "a vector of length 8" or "a vector of length 5, padded with 3 meaningless zeros." In order to indicate which parts are meaningful and which are padding, we need an additional mask vector (or matrix):

$$m = [1, 1, 1, 1, 1, 0, 0, 0]$$

This is a 0/1 vector (or matrix), where 1 indicates the meaningful part and 0 indicates the meaningless padded part.

What we call "masking" is precisely the operation between $x$ and $m$ that excludes the effect introduced by padding. For example, suppose we want the mean of $x$. The result we'd expect is:

$$\text{avg}(x) = \frac{1 + 0 + 3 + 4 + 5}{5} = 2.6$$

But since the vector has already been padded, computing it directly would give:

$$\frac{1 + 0 + 3 + 4 + 5 + 0 + 0 + 0}{8} = 1.625$$

which introduces a bias. Even worse, for the same input, the number of zeros padded on might not be fixed each time, so the same sample could end up with a different mean every time — which clearly doesn't make sense. Once we have the mask vector $m$, we can rewrite the mean computation as:

$$\text{avg}(x) = \frac{\text{sum}(x\otimes m)}{\text{sum}(m)}$$

Here $\otimes$ denotes elementwise multiplication. This way, the numerator only sums over the non-padded part, and the denominator only counts the non-padded part — no matter how many zeros you pad with, the final result will always be the same.

What if we want the maximum of $x$? We have $\max([1, 0, 3, 4, 5]) = \max([1, 0, 3, 4, 5, 0, 0, 0]) = 5$, which might seem to already exclude the effect of padding. That's true in this particular example, but consider another case:

$$x = [-1, -2, -3, -4, -5]$$

After padding this becomes

$$x = [-1, -2, -3, -4, -5, 0, 0, 0]$$

If we directly take $\max$ over the padded $x$, we get 0, and 0 is not within the original range of values. The way to fix this is to make the padded part small enough that $\max$ can (almost) never land on the padded part, e.g.

$$\max(x) = \max\left(x - (1 - m) \times 10^{10}\right)$$

Normally, the magnitude of neural network inputs and outputs won't be very large, so after $x - (1 - m) \times 10^{10}$, the padded part sits at the $-10^{10}$ order of magnitude, which guarantees that taking $\max$ will never select the padded part.

The same idea applies to handling padding in softmax. In attention or pointer networks, we may need to compute softmax over variable-length vectors. If we compute softmax directly over a padded vector, the padded part will also end up absorbing some of the probability mass, causing the probabilities over the actually meaningful part to no longer sum to 1. The solution is the same as for $\max$: make the padded part small enough that $e^x$ is close enough to 0 to be negligible:

$$\text{sofmax}(x) = \text{softmax}\left(x - (1 - m) \times 10^{10}\right)$$

The masking treatment for the operators above is somewhat special. For most other operations (aside from bidirectional RNNs), it basically suffices to output

$$x\otimes m$$

that is, to keep the padded part equal to 0.

Key points for implementing in Keras

Keras comes with built-in mask support, but I don't recommend using it, because the built-in masking isn't clear or flexible enough, and it doesn't support every layer. I strongly recommend implementing masks yourself.

Several of the models I've open-sourced recently already contain plenty of masking examples, and I believe that if readers carefully read through the source code, they'll find it quite easy to understand how masking is implemented. Let me just mention a few key points here. Generally speaking, the input to an NLP model is a word ID matrix of shape $\text{[batch_size, seq_len]}$, where I use 0 as the padding ID and 1 as the UNK ID (the rest is up to you). I then use a Lambda layer to generate the mask matrix:

# x是词ID矩阵
mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x, 2), 0), 'float32'))(x)

The resulting mask matrix has shape $\text{[batch_size, seq_len, 1]}$, and after the word ID matrix passes through the Embedding layer it has shape $\text{[batch_size, seq_len, word_size]}$, so we can then use the mask matrix to process the output. This is just my personal convention, not the one and only standard.

Combination: bidirectional RNN

Our discussion so far has excluded bidirectional RNNs, because RNNs are recursive models and can't simply be masked (this is mainly an issue for the backward-direction RNN part). A so-called bidirectional RNN runs the RNN forward and backward and then concatenates or sums the results. Suppose we run a backward RNN over $[1, 0, 3, 4, 5, 0, 0, 0]$: the final output will inevitably include the padded zeros (since the padded part is already involved in the computation from the very start). So this can't be fixed after the fact — it has to be excluded beforehand.

The solution is: before running the backward RNN, first reverse $[1, 0, 3, 4, 5, 0, 0, 0]$ into $[5, 4, 3, 0, 1, 0, 0, 0]$, then run a forward RNN, and finally reverse the result back. Note that when reversing, you should only reverse the non-padded part (this is the only way to guarantee that the padded part never participates in the recursive computation, and that the result stays aligned with the forward RNN's output). TensorFlow conveniently provides a ready-made function tf.reverse_sequence() for this.

Unfortunately, Keras's built-in Bidirectional doesn't offer this functionality, so I rewrote it myself, for readers' reference:

class OurBidirectional(OurLayer):
    """自己封装双向RNN，允许传入mask，保证对齐
    """
    def __init__(self, layer, **args):
        super(OurBidirectional, self).__init__(**args)
        self.forward_layer = layer.__class__.from_config(layer.get_config())
        self.backward_layer = layer.__class__.from_config(layer.get_config())
        self.forward_layer.name = 'forward_' + self.forward_layer.name
        self.backward_layer.name = 'backward_' + self.backward_layer.name
    def reverse_sequence(self, x, mask):
        """这里的mask.shape是[batch_size, seq_len, 1]
        """
        seq_len = K.round(K.sum(mask, 1)[:, 0])
        seq_len = K.cast(seq_len, 'int32')
        return tf.reverse_sequence(x, seq_len, seq_dim=1)
    def call(self, inputs):
        x, mask = inputs
        x_forward = self.reuse(self.forward_layer, x)
        x_backward = self.reverse_sequence(x, mask)
        x_backward = self.reuse(self.backward_layer, x_backward)
        x_backward = self.reverse_sequence(x_backward, mask)
        x = K.concatenate([x_forward, x_backward], -1)
        if K.ndim(x) == 3:
            return x * mask
        else:
            return x
    def compute_output_shape(self, input_shape):
        return input_shape[0][:-1] + (self.forward_layer.units * 2,)

It's used essentially the same way as the built-in Bidirectional, except you need to additionally pass in the mask matrix, e.g.:

x = OurBidirectional(LSTM(128))([x, x_mask])

Summary

Keras is an extremely friendly, extremely flexible high-level deep learning API. Don't believe the rumor circulating online that "Keras is friendly to beginners but lacks flexibility" — Keras is friendly to beginners, even friendlier to experienced users, and friendlier still to users who need to frequently write custom modules.