I’m trying to train a similar model and I am getting the same problem. It does work for me however with a relu activation on the last classification layer instead of softmax and a smaller learning rate optimizer = keras.optimizers.Adam(learning_rate=3e-5).
I’m not sure why it isn’t working with softmax. FYI here’s the model I’m using
ids_input = keras.Input(shape=(max_tokens,), dtype=np.int32)
attention_mask_input = keras.Input(shape=(max_tokens,), dtype=np.int32)
bert_model = TFDistilBertModel.from_pretrained('distilbert-base-uncased')
dropout = layers.Dropout(0.1)
token_classifier_layer = layers.Dense(num_labels, activation="relu")

bert_output = bert_model({'input_ids': ids_input, 'attention_mask': attention_mask_input}, return_dict=True)
x = dropout(bert_output['last_hidden_state'])
x = token_classifier_layer(x)
word_classifier_model = keras.Model(inputs=[ids_input, attention_mask_input], outputs=x)

optimizer = keras.optimizers.Adam(learning_rate=3e-5)
word_classifier_model.compile(optimizer=optimizer, loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])

word_classifier_model.fit(x=[input_ids, attention_mask], y=model_labels, epochs=config.epochs, batch_size=config.batch_size, validation_split=0.2)

word_classifier_model.save_weights(config.model_file_path)