Проблема дифференциации в прогнозирующем выравнивании для реализации внимания

Question

Я пытаюсь реализовать внимание local-p на основе этой статьи: https://arxiv.org/pdf/1508.04025.pdf В частности, уравнение (9) получает позицию выравнивания, основанную на взятии сигмоиды некоторых нелинейных функций, а затем умножении полученного значения с количеством временных шагов. Поскольку сигмоид возвращает значения от 0 до 1, это умножение дает действительный индекс от 0 до количества временных шагов. Я могу мягко обойти это, чтобы вывести прогнозируемую позицию, однако я не смог найти способ преобразовать это в целое число для использования в операциях среза / индексации, поскольку tf.cast () не дифференцируем. Другая проблема заключается в том, что производные позиции имеют форму (B, 1) и, следовательно, одну выровненную позицию для каждого примера в партии. См. Ниже, чтобы понять эти операции:

"""B = batch size, S = sequence length (num. timesteps), V = vocabulary size, H = number of hidden dimensions"""
class LocalAttention(Layer):
    def __init__(self, size, window_width=None, **kwargs):
        super(LocalAttention, self).__init__(**kwargs)
        self.size = size
        self.window_width = window_width # 2*D

    def build(self, input_shape): 
        self.W_p = Dense(units=input_shape[2], use_bias=False)
        self.W_p.build(input_shape=(None, None, input_shape[2])) # (B, 1, H)
        self._trainable_weights += self.W_p.trainable_weights

        self.v_p = Dense(units=1, use_bias=False)
        self.v_p.build(input_shape=(None, None, input_shape[2])) # (B, 1, H)
        self._trainable_weights += self.v_p.trainable_weights

        super(Attention, self).build(input_shape)

    def call(self, inputs):
        sequence_length = inputs.shape[1]
        ## Get h_t, the current (target) hidden state ##
        target_hidden_state = Lambda(function=lambda x: x[:, -1, :])(inputs) # (B, H)
        ## Get h_s, source hidden states ##
        aligned_position = self.W_p(target_hidden_state) # (B, H)
        aligned_position = Activation('tanh')(aligned_position) # (B, H)
        aligned_position = self.v_p(aligned_position) # (B, 1)
        aligned_position = Activation('sigmoid')(aligned_position) # (B, 1)
        aligned_position = aligned_position * sequence_length # (B, 1)

Допустим, тензор aligned_position имеет элементы [24.2, 15.1, 12.3] для размера партии = B = 3 для упрощения. Затем исходные скрытые состояния выводятся из входных скрытых состояний (B = 3, S, H), так что для первого примера мы берем временные шаги, начиная с 24, следовательно, что-то вроде first_batch_states = Lambda(function=lambda x: x[:, 24:, :])(inputs) и так далее. Обратите внимание, что реализация local-p более сложная, чем эта, но я упростил ее здесь. Следовательно, основной проблемой является преобразование 24,2 в 24 без потери дифференцируемости или использование какой-либо маскирующей операции для получения индексов через скалярное произведение. Операция маски предпочтительна, так как нам придется делать это для каждого примера в пакетном режиме, и наличие цикла внутри пользовательского слоя Keras не является аккуратным. У вас есть идеи, как решить эту задачу? Буду благодарен за любые ответы и комментарии!

Answer 1

Я нашел два пути решения этой проблемы.

• Применение гауссовского распределения на основе выровненной позиции, показанной в исходном вопросе, к весам внимания, что делает процесс дифференцируемым, как предложил @Siddhant:

gaussian_estimation = lambda s: tf.exp(-tf.square(s - aligned_position) /
                                                   (2 * tf.square(self.window_width / 2)))
gaussian_factor = gaussian_estimation(0)
for i in range(1, sequence_length):
    gaussian_factor = Concatenate()([gaussian_factor, gaussian_estimation(i)])
# Adjust weights via gaussian_factor: (B, S*) to allow differentiability
attention_weights = attention_weights * gaussian_factor # (B, S*)

Следует отметить, что здесь не используется сложная операция нарезки, только простая настройка в соответствии с расстоянием.

• Сохранение верхних значений n и обнуление остальных, как предложено здесь @Vlad, Как реализовать пользовательский слой keras, который содержит только верхние значения n и нули всех остальных? :

aligned_position = self.W_p(inputs) # (B, S, H)
aligned_position = Activation('tanh')(aligned_position) # (B, S, H)
aligned_position = self.v_p(aligned_position) # (B, S, 1)
aligned_position = Activation('sigmoid')(aligned_position) # (B, S, 1)
## Only keep top D values out of the sigmoid activation, and zero-out the rest ##
aligned_position = tf.squeeze(aligned_position, axis=-1) # (B, S)
top_probabilities = tf.nn.top_k(input=aligned_position,
                                k=self.window_width,
                                sorted=False) # (values:(B, D), indices:(B, D))
onehot_vector = tf.one_hot(indices=top_probabilities.indices,
                           depth=sequence_length) # (B, D, S)
onehot_vector = tf.reduce_sum(onehot_vector, axis=1) # (B, S)
aligned_position = Multiply()([aligned_position, onehot_vector]) # (B, S)
aligned_position = tf.expand_dims(aligned_position, axis=-1) # (B, S, 1)
source_hidden_states = Multiply()([inputs, aligned_position]) # (B, S*=S(D), H)
## Scale back-to approximately original hidden state values ##
aligned_position += 1 # (B, S, 1)
source_hidden_states /= aligned_position # (B, S*=S(D), H)

Следует отметить, что здесь мы вместо этого применяем плотные слои ко всем скрытым исходным состояниям, чтобы получить форму (B,S,1) вместо (B,1) для aligned_position. Я считаю, что это настолько близко, насколько мы можем добраться до того, что предлагает газета.

Любой, кто пытается внедрить механизмы внимания, может проверить мое репо https://github.com/ongunuzaymacar/attention-mechanisms. Слои здесь предназначены для последовательных задач «многие-к-одному», но могут быть адаптированы к другим формам с незначительными изменениями.