"Making Keras Cooler!": Layers within Layers and Masking

Avoiding Redundant Work

Readers who frequently use custom layers might feel that they are often repeating work. For example, if we want to add a linear transformation, we have to add kernel and bias variables in build (and customize variable initialization, regularization, etc.), and then use K.dot in call to execute it. Sometimes, dimension alignment also needs to be considered, making the steps quite tedious. However, a linear transformation is essentially just a Dense layer without an activation function. If we can reuse existing layers when customizing a layer, it would clearly save a lot of code.

In fact, if you are familiar with Python’s object-oriented programming and carefully study the source code of Keras’s Layer class, it is not difficult to find a way to reuse existing layers. Below, I have organized this into a standardized process for readers to reference.

(Note: Starting from Keras 2.3.0, the "layers within layers" functionality is built-in. You no longer need the custom OurLayer defined below; you can simply use the standard Layer class.)

OurLayer

First, we define a new OurLayer class:

class OurLayer(Layer):
    """Defines a new Layer and adds a reuse method, 
    allowing existing layers to be called during Layer definition.
    """
    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 and adds a reuse method, which allows us to reuse existing layers.

Below is a simple example defining a layer with the following operation: y = g(f(xW_1 + b_1)W_2 + b_2) Here f and g are activation functions. This is essentially a composition of two Dense layers. If written in the standard way, we would need to define several weights in build, define shapes based on the input, and handle initialization. However, since these are already implemented in the Dense layer, we can call them directly. The reference code is as follows:

class OurDense(OurLayer):
    """Originally inherited from the Layer class, now inherits from 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):
        """Add layers to be reused within the build method.
        Of course, you can still add trainable weights like the standard way.
        """
        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):
        """Simply reuse the layers; equivalent to 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 much cleaner?

Prove You Have Thought About It

Recently, I have used masking extensively in several open-source models, and many readers seem to have never encountered it before, leading to many questions. While it is natural to have questions about something new, asking without thinking can be irresponsible. I have always believed that when asking others a question, one should simultaneously "prove" that they have thought about it. For instance, if you want an explanation of masking, I would first ask you to answer:

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

These three questions are not about the principle of masking; they are simply about whether you understand the computation being performed. Based on that, we can discuss why such a computation is necessary. If you don’t even understand the computation itself, there are only two paths: give up on understanding the problem, or study Keras for a few months before we discuss it further.

Assuming the reader has understood the computation of masking, let’s briefly discuss its basic principles.

Excluding Padding

Masking appears alongside padding because neural network inputs require a regular tensor, while text is usually of variable length. Consequently, cropping or padding is needed to make them a fixed length. By convention, we use 0 as the padding symbol.

Let’s use a simple vector to describe the principle of padding. Suppose there is a vector of length 5: x = [1, 0, 3, 4, 5] After padding to length 8: x = [1, 0, 3, 4, 5, 0, 0, 0] When you input this length-8 vector into a model, the model does not know whether it is a "vector of length 8" or a "vector of length 5 padded with three meaningless zeros." To indicate which parts are meaningful and which are padding, we also need a mask vector (matrix): m = [1, 1, 1, 1, 1, 0, 0, 0] This is a 0/1 vector (matrix), where 1 represents meaningful parts and 0 represents meaningless padding.

Masking is the operation between x and m to exclude the effects of padding. For example, if we want to find the mean of x, the expected result is: \text{avg}(x) = \frac{1 + 0 + 3 + 4 + 5}{5} = 2.6 However, because the vector has been padded, a direct calculation would yield: \frac{1 + 0 + 3 + 4 + 5 + 0 + 0 + 0}{8} = 1.625 This introduces a bias. More seriously, for the same input, the number of zeros padded each time might not be fixed, so the same sample could yield different means, which is unreasonable. With the mask vector m, we can rewrite the mean operation: \text{avg}(x) = \frac{\text{sum}(x \otimes m)}{\text{sum}(m)} Here \otimes denotes element-wise multiplication. In this way, the numerator only sums the non-padding parts, and the denominator counts the non-padding parts. Regardless of how many zeros are padded, the final result remains the same.

What if we want the maximum value of x? We have \max([1, 0, 3, 4, 5]) = \max([1, 0, 3, 4, 5, 0, 0, 0]) = 5. Does it seem like we don’t need to exclude padding effects? In this example, yes, but consider: x = [-1, -2, -3, -4, -5] After padding, it becomes: x = [-1, -2, -3, -4, -5, 0, 0, 0] If we directly take the \max of the padded x, we get 0, which is not within the original range. The solution here is to make the padding part small enough so that the \max (almost) never picks it. For example: \max(x) = \max\left(x - (1 - m) \times 10^{10}\right) Normally, the magnitude of neural network inputs and outputs is not very large, so after x - (1 - m) \times 10^{10}, the padding parts are in the range of -10^{10}, ensuring that the \max operation will not select them.

Handling padding for softmax is similar. In Attention or pointer networks, we may encounter softmax over variable-length vectors. If we directly apply softmax to the padded vector, the padding parts will also take up some probability, causing the sum of probabilities of the meaningful parts to be less than 1. The solution is the same as for \max: make the padding parts small enough so that e^x is close to 0, effectively ignoring them: \text{softmax}(x) = \text{softmax}\left(x - (1 - m) \times 10^{10}\right)

The mask processing for these operators is somewhat special. For most other operations (except bidirectional RNNs), mask processing basically only requires outputting: x \otimes m which keeps the padding parts as 0.

Keras Implementation Points

Keras has built-in masking functionality, but it is not recommended because it is not clear or flexible enough and does not support all layers. It is strongly suggested that readers implement masking themselves.

Several recently open-sourced models have provided many masking examples. I believe that by reading the source code carefully, readers will easily understand how to implement masking. Here are a few key points. Generally, the input to an NLP model is a word ID matrix of shape [batch_size, seq_len]. I use 0 as the padding ID and 1 as the UNK ID. Then, I use a Lambda layer to generate the mask matrix:

# x is the word ID matrix
mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x, 2), 0), 'float32'))(x)

The generated mask matrix has a shape of [batch_size, seq_len, 1]. After the word ID matrix passes through the Embedding layer, the output shape is [batch_size, seq_len, word_size]. Thus, the mask matrix can be used to process the output. This approach is just my personal habit, not the only standard.