如何停止张量流中张量的某些条目的梯度

Question

我正在尝试实现一个嵌入层。嵌入将使用预先训练的手套嵌入来初始化。对于可以在手套中找到的单词，它将被修复。对于那些没有出现在手套中的单词，它会被随机初始化，并且是可训练的。我如何在张量流中做到这一点？我知道整个张量有一个 tf.stop_gradient ，有没有适合这种场景的 stop_gradient api？或者，有什么解决方法吗？任何建议表示赞赏

Answer 1

所以我们的想法是使用

mask

和

tf.stop_gradient

来解决这个问题：

res_matrix = tf.stop_gradient(mask_h*E) + mask*E

，

在矩阵

mask

中，1表示我想要应用梯度的条目，0表示我不想应用梯度的条目（将梯度设置为0），

mask_h

是

mask

的倒数（1翻转为0，0翻转为1）。然后我们可以从

res_matrix

中获取。这是测试代码：

import tensorflow as tf
import numpy as np

def entry_stop_gradients(target, mask):
    mask_h = tf.abs(mask-1)
    return tf.stop_gradient(mask_h * target) + mask * target

mask = np.array([1., 0, 1, 1, 0, 0, 1, 1, 0, 1])
mask_h = np.abs(mask-1)

emb = tf.constant(np.ones([10, 5]))

matrix = entry_stop_gradients(emb, tf.expand_dims(mask,1))

parm = np.random.randn(5, 1)
t_parm = tf.constant(parm)

loss = tf.reduce_sum(tf.matmul(matrix, t_parm))
grad1 = tf.gradients(loss, emb)
grad2 = tf.gradients(loss, matrix)
print matrix
with tf.Session() as sess:
    print sess.run(loss)
    print sess.run([grad1, grad2])

Answer 2

我建议您使用两个不同的张量来保存数据：一个 tf.constant 用于预训练数据，一个 tf.Variable 用于要训练的新数据，然后您可以将两者与串联和同样的张量混合加盟运营。

由于tf.constant无法训练，所以你不必担心停止梯度。

Answer 3

我对词嵌入了解不多，但我对你的问题的理解是你有一个变量

v

并且你只想训练它的某些（已知）条目。您可以通过使用“掩模”操纵梯度来实现这一点，即与

v

形状相同的常量张量，可训练条目的值为 1，否则为 0。

v = your_variable()
loss = your_loss() #some loss that uses v
mask = your_mask_as_explained_above()
opt = tf.train.GradientDescentOptimizer(learning_rate=0.1)

# Get list (length 1 in our example) of (gradient, variable)-pairs from the optimizer and extract the gradient w.r.t. v
grads_and_vars = opt.compute_gradients(loss, [v])
v_grad = grads_and_vars[0][0]

# Multiply the gradient with the mask before feeding it back to the optimizer
sgd_step = opt.apply_gradients([(v, v_grad*mask)])

根据问题的复杂性，这可能不是一个有效的解决方案，因为完整的梯度 w.r.t.

v

在每个步骤中计算。它根本没有在优化器更新中应用。

如果您不熟悉

opt.compute_gradients

和

opt.apply_gradients

，请在 docs 中找到解释。

Answer 4

这是自定义图层，我将其称为

GradientControlLayer

：

from keras.api.layers import Layer
from keras import ops

class GradientControlLayer(Layer):
  def __init__(self, **kwargs):
    super(GradientControlLayer, self).__init__(**kwargs)
  
  def call(self, inputs):
    target, mask = inputs
    
    # Cast the mask to the same type as the target
    mask = ops.cast(mask, dtype=target.dtype)
    mask_h = ops.logical_not(mask)
    mask_h = ops.cast(mask_h, dtype=target.dtype)
    
    # Stop the gradients where the mask is False
    stopped_gradient_part = ops.stop_gradient(target)
    
    # Multiply stopped gradient part with mask_h
    stopped_gradient_masked = stopped_gradient_part * mask_h
    
    # Multiply non-stopped gradient part with mask
    non_stopped_gradient_part = target * mask
    
    # Add the stopped and non-stopped parts
    return stopped_gradient_masked + non_stopped_gradient_part

使用示例：

from keras.api.layers import Input
from keras.api.models import Model

my_input = Input(...) # Replace with your use case
gradient_control_mask = Input(1,) # Must be (1,) because it's scalar.

# Assume prev layer is defined
stop_grad_layer = GradientControlLayer()([prev_layer, gradient_control_mask])
next_layer = Layer()(stop_grad_layer)

model = Model(inputs=[my_input, gradient_control_mask], outputs=[...])

其中值为 0 的 mask 表示阻止梯度，而值为 1 的 mask 表示允许梯度流。

这是工作证明：https://colab.research.google.com/drive/1DxMFoEAqxwPPC20QfJGs3-EDZrqd_Ecn?usp=sharing

这是绘图模型示例，其中输出 B 受到控制，而输出 A 不受控制：