Введение

Я пытаюсь реализовать архитектуру трансформатора набора в TensorFlow.

Одним из модулей, предложенных исследователями, является модуль объединения с многоголовым вниманием (PMA). Если я правильно понимаю, этот модуль использует обучаемый начальный вектор. Этот вектор используется в качестве запроса к многоголовому модулю внимания, ключи и значения которого являются векторами, заданными для объединения.

В авторской реализации этого модуля (с помощью PyTorch) мы видим, что этот вектор повторяется во время прямого прохода для соответствия размерности пакета:

 def forward(self, X):
        return self.mab(self.S.repeat(X.size(0), 1, 1), X)
 

Следовательно, моя реализация прямого прохода с тензорным потоком (как a tf.keras.layers.Layer ) является:

 def call(self, inputs, **kwargs):
    q = inputs
    # !!! Is the repeat operation allowing the back-propagation of gradient ?
    s = tf.expand_dims(self.seed_vector, axis=0)
    s = tf.repeat(s, tf.shape(q)[0], axis=0)
    return self.mab((s, q))
 

( self.mab будучи, в обеих реализациях, Многоголовым блоком внимания, определенным в документе)

В моей реализации self.seed_vector к слою добавляется обучаемый вес:

 def build(self, input_shape):
    # [...]
    self.seed_vector = self.add_weight(
        shape=(1, input_shape[2]),
        initializer="random_normal",
        dtype=tf.float32,
        trainable=True
    )
 

Я не нашел более простого способа выполнить повторение исходного вектора для соответствия ограничению размерности пакета, но код компилируется, и модель может быть обучена.

Проблема

Моя проблема в том, что я заметил большую нестабильность потери обучения/проверки модели во время обучения, когда я использую модуль PMA вместо простого усредняющего слоя объединения. Следовательно, мне интересно, правильна ли моя реализация PMA. Более конкретно, поскольку начальный вектор должен быть изучен с помощью весов модели, мне интересно узнать о влиянии tf.repeat операции на вычисления градиента ошибок в отношении этих обучаемых весов векторов. Может ли TensorFlow «отцепить» этот вектор при вычислении градиента ошибки и правильно обновить вектор в обратном проходе ?

Я не заметил ничего, связанного с этой частью в учебном коде реализации PyTorch авторов. Но вычисление градиентов PyTorch и обратное распространение могут отличаться от того, что реализовано в TensorFlow, и я не смог найти в документации TensorFlow ничего, связанного с этим потенциальным воздействием tf.repeat операции. Если бы у кого-то были идеи по этому поводу, я был бы очень благодарен !

Nota bene

I know that unexpected results can be caused by many factors and I am, of course, investigating them. But I am really suspicious on this specific part of the model since its introduction seems to significantly hinder the training behavior. If the implementation is OK, the reason might simply be that the use of this module is not a good idea for my task/data. I am not asking for help on the whole model implementation.