Building custom model over the final embedding layer

BERT supposedly generates 768 dimensional embeddings for tokens. I am trying to build a multi-class classification model on top of this. My assumption is that the output of layer Encoder-12-FeedForward-Norm of shape (None, [seq_length], 768) would give this embeddings. This is what I am trying :

model = load_trained_model_from_checkpoint(config_path, checkpoint_path, training=True, seq_len=seq_len)

new_out = Bidirectional(LSTM(50, return_sequences=True, 
                       dropout=0.1, 
                       recurrent_dropout=0.1))(model.layers[-9].output)
new_out = GlobalMaxPool1D()(new_out)
new_out = Dense(50, activation='relu')(new_out)
new_out = Dropout(0.1)(new_out)
new_out = Dense(6, activation='sigmoid')(new_out)

newModel = Model(model.inputs[:2], new_out)

I get the following error for new_out = GlobalMaxPool1D()(new_out) :

TypeError: Layer global_max_pooling1d_11 does not support masking, but was passed an input_mask: Tensor("Encoder-12-FeedForward-Add/All:0", shape=(?, 128), dtype=bool)

I am not sure how masking is involved if I am just using the output of the encoder.

The paper mentions that the output corresponding to just the first [CLS] token should be used for classification. On trying this :

new_out = Lambda(lambda x: x[:,0,:])(model.layers[-9].output)

the model trains (although with poor results).

How can the pre-loaded model be used for classification?

[CyberZHG]

I forgot to return a None mask in MaskedGlobalMaxPool1D. I’ve fixed it and made a release.

[CyberZHG]

#7 Sentence Embedding

GlobalMaxPool1D doesn’t support masking. Following is a modification that suits this case:

https://github.com/CyberZHG/keras-bert/blob/b7ecdc34637435a849695b1b1a4ebdc4da842832/keras_bert/layers/pooling.py#L5-L21

I’ve added a demo for sentence embedding with pooling:

https://github.com/CyberZHG/keras-bert/blob/02c7eb20c56b02ed74226d056354658c247ab52c/demo/load_model/load_and_pool.py#L20-L39