Проблемы в реализации RNN с несколькими выходными слоями TimeDistributed

Question

Я пытаюсь перевести следующую прямую связь DNN

input_img = Input(shape=(1025, 70, 3 ) )
x = ( Flatten())(input_img)

for i in range(0,4):
   x = ( Dense(1024*3))(x)
   x = ( BatchNormalization() )(x)
   x = ( LeakyReLU())(x)
o0 = ( Dense(1025, activation='sigmoid'))(x)
o1 = ( Dense(1025, activation='sigmoid'))(x)
o2 = ( Dense(1025, activation='sigmoid'))(x)

в RNN с тремя слоями LSTM и тремя выходами o0, o1, o2, которые должны быть тремя плотными слоями.

Размер входного сигнала RNN равен (10,70,1025 * 3) (немного изменен для DNN с прямой связью, но с такими же точными данными), в то время как каждый выход имеет размер (10,70,1025)

Я попытался создать три слоя LSTM с субнейронной сетью (здесь называемой подсетью), которая использует функцию TimeDistributed.

Из-за сложности реализации функции TimeDistributed я подумал, что было бы неплохо получить три выходных слоя, разделив длинный плотный слой, как показано ниже:

def lambda_fun(y) :

    x = tf.split(y,[1025,1025,1025],-1)
    return x

input_img = Input(batch_shape=(10,70,3*1025 ) )
x = Bidirectional(LSTM(128,return_sequences=True, dropout=0.5, recurrent_dropout=0.2))(input_img)
x = Bidirectional(LSTM(128,return_sequences=True, dropout=0.5, recurrent_dropout=0.2))(x)
x = Bidirectional(LSTM(128,  return_sequences=True, dropout=0.5, recurrent_dropout=0.2))(x)

subnet = Sequential()
subnet.add(Dense(1025*3, activation='relu',input_shape=(x._keras_shape[-1],)))
subnet.add(BatchNormalization())
subnet.add(LeakyReLU())
y = TimeDistributed(subnet)(x)

o0, o1, o2 =  Lambda(lambda_fun )(y)

Хотя обучение RNN проходит гладко, функция стоимости снижается гораздо медленнее, чем в случае использования простой обратной связи, что заставляет меня думать, что моя попытка внедрения может быть неудачной. Есть предложения?