Why would you implement the position-wise feed-forward network of the transformer with convolution layers?

Question

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?

Update

Adding benchmarks in response to the answer by @mshlis:

import os
import typing as t
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import numpy as np

from keras import layers, models
from keras import backend as K
from tensorflow import Tensor


# Generate random data

n = 128000  # n samples
seq_l = 32  # sequence length
emb_dim = 512  # embedding size

x = np.random.normal(0, 1, size=(n, seq_l, emb_dim)).astype(np.float32)
y = np.random.binomial(1, 0.5, size=n).astype(np.int32)

---

# Define constructors

def ffn_dense(hid_dim: int, input_: Tensor) -> Tensor:
    output_dim = K.int_shape(input_)[-1]
    hidden = layers.Dense(hid_dim, activation='relu')(input_)
    return layers.Dense(output_dim, activation=None)(hidden)


def ffn_cnn(hid_dim: int, input_: Tensor) -> Tensor:
    output_dim = K.int_shape(input_)[-1]
    hidden = layers.Conv1D(hid_dim, 1, activation='relu')(input_)
    return layers.Conv1D(output_dim, 1, activation=None)(hidden)


def build_model(ffn_implementation: t.Callable[[int, Tensor], Tensor], 
                ffn_hid_dim: int, 
                input_shape: t.Tuple[int, int]) -> models.Model:
    input_ = layers.Input(shape=(seq_l, emb_dim))
    ffn = ffn_implementation(ffn_hid_dim, input_)
    flattened = layers.Flatten()(ffn)
    output = layers.Dense(1, activation='sigmoid')(flattened)
    model = models.Model(inputs=input_, outputs=output)
    model.compile(optimizer='Adam', loss='binary_crossentropy')
    return model

---

# Build the models

ffn_hid_dim = emb_dim * 4  # this rule is taken from the original paper
bath_size = 512  # the batchsize was selected to maximise GPU load, i.e. reduce PCI IO overhead

model_dense = build_model(ffn_dense, ffn_hid_dim, (seq_l, emb_dim))
model_cnn = build_model(ffn_cnn, ffn_hid_dim, (seq_l, emb_dim))

---

# Pre-heat the GPU and let TF apply memory stream optimisations

model_dense.fit(x=x, y=y[:, None], batch_size=bath_size, epochs=1)
%timeit model_dense.fit(x=x, y=y[:, None], batch_size=bath_size, epochs=1)

model_cnn.fit(x=x, y=y[:, None], batch_size=bath_size, epochs=1)
%timeit model_cnn.fit(x=x, y=y[:, None], batch_size=bath_size, epochs=1)

I am getting 14.8 seconds per epoch with the Dense implementation:

Epoch 1/1
128000/128000 [==============================] - 15s 116us/step - loss: 0.6332
Epoch 1/1
128000/128000 [==============================] - 15s 115us/step - loss: 0.5327
Epoch 1/1
128000/128000 [==============================] - 15s 117us/step - loss: 0.3828
Epoch 1/1
128000/128000 [==============================] - 14s 113us/step - loss: 0.2543
Epoch 1/1
128000/128000 [==============================] - 15s 116us/step - loss: 0.1908
Epoch 1/1
128000/128000 [==============================] - 15s 116us/step - loss: 0.1533
Epoch 1/1
128000/128000 [==============================] - 15s 117us/step - loss: 0.1475
Epoch 1/1
128000/128000 [==============================] - 15s 117us/step - loss: 0.1406

14.8 s ± 170 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

and 18.2 seconds for the CNN implementation. I am running this test on a standard Nvidia RTX 2080. So, from a performance perspective there seems to be no point in actually implementing an FFN block as a CNN in Keras. Considering that the maths are the same, the choice boils down to pure aesthetics.

Answer 1

I'm going to post another guess to this question - it won't be a complete answer, but hopefully it'll provide some direction towards finding a more legitimate answer.

The feed-forward networks as suggested by Vaswani are very reminiscent of the sparse autoencoders. Where the input / output dimensions are much greater than the hidden input dimension.

If you aren't familiar with sparse autoencoders, this is a little counter intuitive - WTF would you have a larger hidden dimension?

The intuition borrows from infinitely wide neural networks. If you have an infinitely wide neural network, you have basically have a Gaussian process and sample any function you'd like. So the wider the network you have, the more approximation power that you have. In the case of inputs, this is a matter of learning a dictionary. If you have only discrete inputs, this hidden layer will be capped at $O(2^N)$ width, where $N$ is the maximum number of bits it takes to represent the input (which would boil down to approximating a lookup table).

Of course, these aren't trivial to implement in practice. These layers are bound to be bloated with identifiability issues. Common approaches include $L_1$ regularization. I'm guessing that the convolutional layers + dropout are just another attempt to deal with these sorts of identifiability issues. Furthermore, the FFN is an attempt to learn an arbitrary mapping for individual words (you can think of mapping words to synonyms for instance).

These are all guesses though - more intuition is welcome.

Answer 2

1) The math is the exact same, so from an optimization or mathematical perspective there is no difference

2) Here are my guesses to a possible answer.

• Habit: People may just call one over the other out of habit
• Generality: Across frameworks a 1d convolution op would work, while Dense of FC may need adjustments to work on the temporal axis
• Parallel Workers: Convolution and Dense call different subroutines in the backend, and the one used by convolution may have better gains on sequential input for this purpose

Edit
Regarding bench-marking the 2, your experiment was shallow. I didn't have time to wait to do a full gird search, so i held 3 paramaters constant and fluctuated one. Here are the results (note the model was just a simple feed forward relu residual model)

Note that in a couple yeah dense out performs conv but it isn't consistent and there are scenarios where it is not true. This is only for a small grid that I chose but you can extend this yourself to check. So it is not as straightforward to say one is sheerly better than the other.