Converting a Standard LSTM RNN over to a Transformer Model

Question

I am looking for some advice on converting my existing CNN/LSTM RNN over to a Transformer type model. This regression model takes a sliding window size of 240 rows with 33 features. It aims to predict the percent difference of the next row of data we are working with. I have fully trained this model below with ~85% accuracy which is fantastic for its use case.

Here is the existing model:

model = Sequential()

model.add(Conv1D(filters=128, kernel_size=3, activation='relu', input_shape=(inp_history_size, len(features))))
model.add(BatchNormalization())

model.add(Conv1D(filters=256, kernel_size=3, activation='relu'))
model.add(BatchNormalization())


model.add(Bidirectional(LSTM(256, return_sequences=True)))

model.add(Bidirectional(LSTM(128, return_sequences=True)))

model.add(Bidirectional(LSTM(64)))
model.add(Dropout(0.3))

model.add(Dense(64, activation='elu', kernel_regularizer=l1_l2(l1=l1_reg_strength, l2=l2_reg_strength)))
model.add(BatchNormalization())

model.add(Dense(1))

I have tried to craft a basic transformer model in a bid to try a different approach and possibly gain better accuracy. However, no matter what I seem to do, it stalls at ~55% accuracy and losses stop falling.

Here is the basic transformer code I have:

def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
    x = layers.LayerNormalization(epsilon=1e-6)(inputs)
    x = layers.MultiHeadAttention(
        key_dim=head_size, num_heads=num_heads, dropout=dropout
    )(x, x)
    x = layers.Dropout(dropout)(x)
    res = x + inputs

    x = layers.LayerNormalization(epsilon=1e-6)(res)
    x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="elu")(x)
    x = layers.Dropout(dropout)(x)
    x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
    return x + res
           
    
    
def build_model(
    input_shape,
    head_size,
    num_heads,
    ff_dim,
    num_transformer_blocks,
    mlp_units,
    dropout=0,
    mlp_dropout=0,
):
    inputs = keras.Input(shape=input_shape)
    x = inputs
    for _ in range(num_transformer_blocks):
        x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)

    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
    for dim in mlp_units:
        x = layers.Dense(dim, activation="elu")(x)
        x = layers.Dropout(mlp_dropout)(x)
    outputs = layers.Dense(1)(x)
    return keras.Model(inputs, outputs)




input_shape=(inp_history_size, len(features))

  
    
model = build_model(
    input_shape,
    head_size=512,
    num_heads=6,
    ff_dim=8,
    num_transformer_blocks=8,
    mlp_units=[256],
    mlp_dropout=0.2,
    dropout=0.1,
)

I am curious as to why the transformer model stalls in learning. I have run several tests tweaking the hyperparameters but all end up with the same result.

I run a learning rate finder before every test. I hope I haven't missed anything but will update this question with further information if needed.

Can anyone lend any advice as to getting this up and running and learning well? Thank you for your help, tips and code examples in advance.

Answer 1

With a bit of elastic net, dynamic gradient clipping and adjustments to the transformer model build training is progressing nicely now.

Here is the build that fixed it:

l1_reg_strength = 0.0010
l2_reg_strength = 0.0010


class AutoClipper:
    def __init__(self, clip_percentile, history_size=10000):
        self.clip_percentile = clip_percentile
        self.grad_history = tf.Variable(tf.zeros(history_size), trainable=False)
        self.i = tf.Variable(0, trainable=False)
        self.history_size = history_size

    def __call__(self, grads_and_vars):
        grad_norms = [self._get_grad_norm(g) for g, _ in grads_and_vars]
        total_norm = tf.norm(grad_norms)
        assign_idx = tf.math.mod(self.i, self.history_size)
        self.grad_history = self.grad_history[assign_idx].assign(total_norm)
        self.i = self.i.assign_add(1)
        clip_value = tfp.stats.percentile(self.grad_history[: self.i], q=self.clip_percentile)
        return [(tf.clip_by_norm(g, clip_value), v) for g, v in grads_and_vars]

    def _get_grad_norm(self, t, axes=None, name=None):
        values = tf.convert_to_tensor(t.values if isinstance(t, tf.IndexedSlices) else t, name="t")

        # Calculate L2-norm, clip elements by ratio of clip_norm to L2-norm
        l2sum = tf.math.reduce_sum(values * values, axes, keepdims=True)
        pred = l2sum > 0
        # Two-tap tf.where trick to bypass NaN gradients
        l2sum_safe = tf.where(pred, l2sum, tf.ones_like(l2sum))
        return tf.squeeze(tf.where(pred, tf.math.sqrt(l2sum_safe), l2sum))



def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
    # Normalization and Attention
    x = layers.LayerNormalization(epsilon=1e-6)(inputs)
    x = layers.MultiHeadAttention(
        key_dim=head_size, num_heads=num_heads, dropout=dropout
    )(x, x)
    x = layers.Dropout(dropout)(x)
    res = x + inputs

    # Feed Forward Part
    x = layers.LayerNormalization(epsilon=1e-6)(res)
    x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="gelu")(x)
    x = layers.Dropout(dropout)(x)
    x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
    return x + res
    
    
    
    
    
def build_model(
    input_shape,
    head_size,
    num_heads,
    ff_dim,
    num_transformer_blocks,
    mlp_units,
    dropout=0,
    mlp_dropout=0,
):
    inputs = keras.Input(shape=input_shape)
    x = inputs
    for _ in range(num_transformer_blocks):
        x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)

    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
    for dim in mlp_units:
        x = layers.Dense(dim, activation="gelu", kernel_regularizer=l1_l2(l1=l1_reg_strength, l2=l2_reg_strength))(x)
        x = layers.Dropout(mlp_dropout)(x)
    outputs = layers.Dense(1)(x)
    return keras.Model(inputs, outputs)




input_shape=(inp_history_size, len(features))

  
    
model = build_model(
    input_shape,
    head_size=512,
    num_heads=8,
    ff_dim=8,
    num_transformer_blocks=8,
    mlp_units=[512],
    mlp_dropout=0.2,
    dropout=0.1,
)