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(
    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))

      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) +
              logits=policy_logits, labels=target_policy))

      loss += tf.scale_gradient(l, gradient_scale)

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


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

  • 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
    Commented Jan 2, 2020 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
    Commented Jan 6, 2020 at 16:55

2 Answers 2


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)
  • $\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
    Commented Feb 14, 2020 at 23:14

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')

  • $\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
    Commented Jan 6, 2020 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
    Commented Jan 6, 2020 at 16:56

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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