Обновление только части матрицы встраивания слов в Tensorflow

Question

Предполагая, что я хочу обновить предварительно обученную матрицу встраивания слов во время обучения, есть ли способ обновить только подмножество матрицы встраивания слов?

Я посмотрел на страницу API Tensorflow и нашел это:

# Create an optimizer.
opt = GradientDescentOptimizer(learning_rate=0.1)

# Compute the gradients for a list of variables.
grads_and_vars = opt.compute_gradients(loss, <list of variables>)

# grads_and_vars is a list of tuples (gradient, variable).  Do whatever you
# need to the 'gradient' part, for example cap them, etc.
capped_grads_and_vars = [(MyCapper(gv[0]), gv[1])) for gv in grads_and_vars]

# Ask the optimizer to apply the capped gradients.
opt.apply_gradients(capped_grads_and_vars)

Однако, как я могу применить это к матрице вложения слов. Предположим, я делаю:

word_emb = tf.Variable(0.2 * tf.random_uniform([syn0.shape[0],s['es']], minval=-1.0, maxval=1.0, dtype=tf.float32),name='word_emb',trainable=False)

gather_emb = tf.gather(word_emb,indices) #assuming that I pass some indices as placeholder through feed_dict

opt = tf.train.AdamOptimizer(1e-4)
grad = opt.compute_gradients(loss,gather_emb)

Как мне тогда использовать opt.apply_gradients а также tf.scatter_update обновить оригинальную матрицу Embeddign? (Кроме того, тензор потока выдает ошибку, если второй аргумент compute_gradient это не tf.Variable)

Answer 1

TL;DR: реализация по умолчанию opt.minimize(loss), TensorFlow сгенерирует редкое обновление для word_emb который изменяет только строки word_emb что участвовал в форварде.

Градиент tf.gather(word_emb, indices) оп относительно word_emb это tf.IndexedSlices объект ( см. реализацию для более подробной информации). Этот объект представляет разреженный тензор, который везде равен нулю, за исключением строк, выбранных indices, Вызов opt.minimize(loss) звонки AdamOptimizer._apply_sparse(word_emb_grad, word_emb), который звонит tf.scatter_sub(word_emb, ...) * обновляет только строки word_emb которые были выбраны indices,

Если, с другой стороны, вы хотите изменить tf.IndexedSlices что возвращается opt.compute_gradients(loss, word_emb), вы можете выполнять произвольные операции TensorFlow на его indices а также values свойства, и создать новый tf.IndexedSlices которые могут быть переданы opt.apply_gradients([(word_emb, ...)]), Например, вы можете ограничить градиенты, используя MyCapper() (как в примере), используя следующие вызовы:

grad, = opt.compute_gradients(loss, word_emb)
train_op = opt.apply_gradients(
    [tf.IndexedSlices(MyCapper(grad.values), grad.indices)])

Точно так же вы можете изменить набор индексов, которые будут изменены, создав новый tf.IndexedSlices с другими показателями.

---

* В общем, если вы хотите обновить только часть переменной в TensorFlow, вы можете использовать tf.scatter_update() , tf.scatter_add() , или же tf.scatter_sub() операторы, которые соответственно устанавливают, добавляют к (+=) или вычесть из (-=) значение, ранее сохраненное в переменной.

Answer 2

Поскольку вы просто хотите выбрать элементы для обновления (а не для изменения градиентов), вы можете сделать следующее.

Позволять indices_to_update быть булевым тензором, указывающим индексы, которые вы хотите обновить, и entry_stop_gradients определяется по ссылке, тогда:

gather_emb = entry_stop_gradients(gather_emb, indices_to_update)

( Источник)

Answer 3

Собственно, я тоже боролся с такой проблемой. В моем случае мне нужно было обучить модель с вложениями w2v, но не все токены существовали в матрице встраивания. Таким образом, для тех токенов, которых не было в матрице, я произвел случайную инициализацию. Конечно, токены, для которых уже были обучены вложения, не должны обновляться, поэтому я придумал такое решение:

      class PartialEmbeddingsUpdate(tf.keras.layers.Layer):
def __init__(self, len_vocab, 
             weights,
            indices_to_update):
    super(PartialEmbeddingsUpdate, self).__init__()
    self.embeddings = tf.Variable(weights, name='embedding', dtype=tf.float32)
    self.bool_mask = tf.equal(tf.expand_dims(tf.range(0,len_vocab),1), tf.expand_dims(indices_to_update,0))
    self.bool_mask = tf.reduce_any(self.bool_mask,1)
    self.bool_mask_not = tf.logical_not(self.bool_mask)
    self.bool_mask_not = tf.expand_dims(tf.cast(self.bool_mask_not, dtype=self.embeddings.dtype),1)
    self.bool_mask = tf.expand_dims(tf.cast(self.bool_mask, dtype=self.embeddings.dtype),1)
    
def call(self, input):
    input = tf.cast(input, dtype=tf.int32)
    embeddings = tf.stop_gradient(self.bool_mask_not * self.embeddings) + self.bool_mask * self.embeddings
    return tf.gather(embeddings,input)

Где len_vocab - длина вашего словаря, weights - матрица весов (некоторые из которых не должны обновляться) и indices_to_update - индексы тех токенов, которые нужно обновить. После этого я применил этот слой вместо tf.keras.layers.Embeddings. Надеюсь, это поможет всем, кто столкнулся с такой же проблемой.