2
$\begingroup$

In the MuZero paper pseudocode, they have the following line of code:

hidden_state = tf.scale_gradient(hidden_state, 0.5)

What does this do? Why is it there?

I've searched for tf.scale_gradient and it doesn't exist in tensorflow. And, unlike scalar_loss, they don't seem to have defined it in their own code.

For context, here's the entire function:

def update_weights(optimizer: tf.train.Optimizer, network: Network, batch,
                   weight_decay: float):
  loss = 0
  for image, actions, targets in batch:
    # Initial step, from the real observation.
    value, reward, policy_logits, hidden_state = network.initial_inference(
        image)
    predictions = [(1.0, value, reward, policy_logits)]

    # Recurrent steps, from action and previous hidden state.
    for action in actions:
      value, reward, policy_logits, hidden_state = network.recurrent_inference(
          hidden_state, action)
      predictions.append((1.0 / len(actions), value, reward, policy_logits))

      # THIS LINE HERE
      hidden_state = tf.scale_gradient(hidden_state, 0.5)

    for prediction, target in zip(predictions, targets):
      gradient_scale, value, reward, policy_logits = prediction
      target_value, target_reward, target_policy = target

      l = (
          scalar_loss(value, target_value) +
          scalar_loss(reward, target_reward) +
          tf.nn.softmax_cross_entropy_with_logits(
              logits=policy_logits, labels=target_policy))

      # AND AGAIN HERE
      loss += tf.scale_gradient(l, gradient_scale)

  for weights in network.get_weights():
    loss += weight_decay * tf.nn.l2_loss(weights)

  optimizer.minimize(loss)

What does scaling the gradient do, and why are they doing it there?

$\endgroup$
2
  • 1
    $\begingroup$ I have searched for tf.scale_gradient() on the TensorFlow's website. As the results show, nothing comes out. It must be a function from old TF versions that now has been abandoned. For sure, it's no more available in TF 2.0. $\endgroup$ – Leevo Jan 2 '20 at 8:37
  • $\begingroup$ I don't believe it's ever been a function in tensorflow, given the lack of results from a Google search for it. $\endgroup$ – Pro Q Jan 6 '20 at 16:55
4
$\begingroup$

Author of the paper here - I missed that this is apparently not a TensorFlow function, it's equivalent to Sonnet's scale_gradient, or the following function:

def scale_gradient(tensor, scale):
  """Scales the gradient for the backward pass."""
  return tensor * scale + tf.stop_gradient(tensor) * (1 - scale)
$\endgroup$
1
  • $\begingroup$ Thank you very much for the reply! If you would be willing to look at stackoverflow.com/q/60234530 (another MuZero question), I would greatly appreciate it. $\endgroup$ – Pro Q Feb 14 '20 at 23:14
1
$\begingroup$

Given that its pseude code? (since its not in TF 2.0) I would go with gradient clipping or batch normalisation ('scaling of activation functions')

$\endgroup$
2
  • $\begingroup$ From the link you provided, it looks like this would likely be gradient norm scaling, which translates to setting a clipnorm parameter in the optimizer. However, in the code they use gradient scaling twice in the code with different values each time. The clipnorm parameter would not allow me to do this. Do you know how I could? $\endgroup$ – Pro Q Jan 6 '20 at 16:47
  • $\begingroup$ Also, the hidden state of a model doesn't seem like something that should be clipped. (I don't understand why clipping it would be helpful at all.) Explaining what gradient clipping would be doing there would be extremely helpful for me to be certain that your answer is correct. $\endgroup$ – Pro Q Jan 6 '20 at 16:56

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.