The Transformer model introduced in "Attention is all you need" by Vaswani et al. incorporates a so-called position-wise feed-forward network (FFN):

In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.

$$\text{FFN}(x) = \max(0, x \times {W}_{1} + {b}_{1}) \times {W}_{2} + {b}_{2}$$

While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way of describing this is as two convolutions with kernel size 1. The dimensionality of input and output is ${d}_{\text{model}} = 512$, and the inner-layer has dimensionality ${d}_{ff} = 2048$.

I have seen at least one implementation in Keras that directly follows the convolution analogy. Here is an excerpt from attention-is-all-you-need-keras.

class PositionwiseFeedForward():
    def __init__(self, d_hid, d_inner_hid, dropout=0.1):
        self.w_1 = Conv1D(d_inner_hid, 1, activation='relu')
        self.w_2 = Conv1D(d_hid, 1)
        self.layer_norm = LayerNormalization()
        self.dropout = Dropout(dropout)
    def __call__(self, x):
        output = self.w_1(x) 
        output = self.w_2(output)
        output = self.dropout(output)
        output = Add()([output, x])
        return self.layer_norm(output)

Yet, in Keras you can apply a single Dense layer across all time-steps using the TimeDistributed wrapper (moreover, a simple Dense layer applied to a 2D input implicitly behaves like a TimeDistributed layer). Therefore, in Keras a stack of two Dense layers (one with a ReLU and the other one without an activation) is exactly the same thing as the aforementioned position-wise FFN. So, why would you implement it using convolutions?


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.