I'm trying to solve the word ordering task: given a syntactically unordered sentence, recover the right order of the words. The adopted approach is to transform each sentence in a dependency tree and then order each subtree to obtain the local order. Once you have the local order of each subtree, you can restore the global order.

So, I built two neural network to achieve this. The output of NN is a real number for each word of the subtree: each output of the subtree is then ranked in order to obtain integer number representing word position.

I trained a simple MLP, which yields good result. This is the architecture:

mlp = keras.models.Sequential()

keras.layers.Dense(
units=training_dataset.shape[1],
input_shape = (training_dataset.shape[1],),
kernel_initializer=keras.initializers.RandomUniform(minval=-0.05, maxval=0.05, seed=None),
activation='relu')
)

keras.layers.Dense(
units=training_dataset.shape[1] + 10,
input_shape = (training_dataset.shape[1] + 10,),
kernel_initializer=keras.initializers.RandomUniform(minval=-0.05, maxval=0.05, seed=None),
bias_initializer='zeros',
activation='relu')
)

keras.layers.Dense(
units=1,
input_shape = (1, ),
kernel_initializer=keras.initializers.RandomUniform(minval=-0.05, maxval=0.05, seed=None),
bias_initializer='zeros',
activation='linear')
)


Number of neurons in the first layer are 837 and in the hidden are 847. In this architecture I used train_on_batch to train the model, when each batch is a subtree. The loss of this model starts from 0.92 and end around 0.88.

Now, I'm trying to build an LSTM to achieve better results, because I think that the past information are very important in this task. Since LSTM wants a 3d input (sequence) I padded each subtree to the same length (maximum subtree length found in the dataset). This is the architecture of LSTM:

model = tf.keras.Sequential()

Where features is equal to 837 and timesteps is 20. The loss of this model starts from 0.33 and end around to 0.31. Also the validation loss is around those values.
However, the results of the LSTM are rather worse than the MLP. In the prediction phase, from the 20 real numbers predicted for each subtree, I take only a number of values corresponding to the number of words of the subtree, excluding the predictions of the pad vectors.
I also tried a Bidirectional LSTM, but loss and results are the same.
Do you have any advice or explanation for LSTM's behavior? Do you see something wrong?