RF weight dropout and variational noise

[albertz]

Currently we don't have weight dropout in the RF. We should add it.

(I thought there was an issue already about it but I don't find it.)

Related:

• Weight dropout in RETURNN common: returnn_common.nn.utils.weight_dropout.weight_dropout. Related issue: https://github.com/rwth-i6/returnn_common/issues/59, https://github.com/rwth-i6/returnn_common/issues/250
• Weight dropout for TF net-dict. See: https://github.com/rwth-i6/returnn/issues/735
• Related discussion (on weight norm): https://github.com/rwth-i6/returnn/issues/1264#issuecomment-1446276707

In general, there is a whole class of similar features on parameter reparameterizations, which would require a similar mechanism, like weight norm.

Regarding implementation:

I think we could follow the PyTorch implementation of similar logic (e.g. weight norm) by using a forward-pre-hook. We already have support for hooks via rf.Module.register_forward_hook/rf.hooks.setup_post_hook_on_method.

Regarding PyTorch:

• torch.nn.utils.weight_norm uses register_forward_pre_hook. However, the doc says:

  This function is deprecated. Use torch.nn.utils.parametrizations.weight_norm() which uses the modern parametrization API.

• torch.nn.utils.parametrizations.weight_norm uses the modern parametrization API, i.e. torch.nn.utils.parametrize.register_parametrization. On register_parametrization (it's somewhat ugly (although our own hooks mechanism is also a bit complicated, but this here looks even worse):
  • It uses _inject_new_class to replace the original module class by a new dummy one, which is going to be extended by adding a property for the parameter.
  • It deletes the parameter (delattr) and stores it in a ParametrizationList.
  • It injects a property (_inject_property) for the parameter. This will call the parametrization to get it.
  • Note that this implementation of parametrization is not working with scripting. (I'm not exactly sure why though?)
  • Even tracing is not working together with caching? (Again not sure why?)
  • There is a caching mechanism which is disabled by default. The user need to explicitly enable it temporarily, like so:

    import torch.nn.utils.parametrize as P
    ...
    with P.cached():
    output = model(inputs)

    When leaving the cached() context, it will clear the (global) cache.

So, comparing the modern parametrization API register_parametrization to the way it was done in torch.nn.utils.weight_norm via register_forward_pre_hook:

• register_parametrization also works for any other module methods, while register_forward_pre_hook only works for forward.
• register_forward_pre_hook will create a temporary buffer to store the calculated parameter. This is buffer is not freed after forward. It is just updated in the next pre-hook. Via register_parametrization, the calculated parameter is freed after it's usage (and with caching enabled, after you leave the cached() context scope). So register_parametrization requires less memory.
• register_parametrization will recompute the parameter several times when it is accessed multiple times (whenever the property is accessed), unless the user uses the parameterization cached mechanism explicitly. But you usually would want to avoid the redundant computation, i.e. you want to use the cached context. But that means you also need to modify the way you call your model by installing the cached context somewhere. This is not automatic. On the other side, register_forward_pre_hook is always automatic. (Although, if you call forward multiple times, it would also cause redundant recomputations of the parameter.)
• register_forward_pre_hook is conceptually much simpler and less ugly than register_parametrization.
• PyTorch parametrization API introduction and discussions: PR https://github.com/pytorch/pytorch/pull/33344, issue https://github.com/pytorch/pytorch/issues/28937, issue https://github.com/pytorch/pytorch/issues/7313.

Concluding from that, I'm a bit unsure what way to go for RF... Using register_forward_pre_hook looks too error-prone (only covers the hooked function, nothing else)... but the other approach looks too complicated? But maybe still better. The caching mechanism is maybe also not so important for now? For all use cases, I think it would not matter (e.g. rf.Linear, rf.SelfAttention, rf.Conv, etc.).

Further, we should also support this with gradient checkpointing such that the weights are not stored twice in memory. In our existing TF implementation of variational noise, we already use gradient checkpointing, where only the random number generator state is stored and not the dropout mask nor the weight. Thus there is almost no memory overhead. See gradient_checkpoint_scope and co. For PyTorch, it is currently unclear how to do this. I moved this over to a separate issue: #1552

[albertz]

Btw, regarding gradient checkpointing, see this current code as an example for variational noise in our TF code:

if param_variational_noise and param.dtype.is_floating and isinstance(param, tf.Variable):
    with default_control_flow_ctx():  # make independent from loop/cond
        with reuse_name_scope_of_tensor(param, postfix="_variational_noise", add_tensor_name=True):

            def _apply_var_noise():
                rnd_state = tf_util.StatelessRandomSeed.create(shape=tf_util.get_shape(param))
                with gradient_checkpoint_scope():
                    noise = rnd_state.normal(stddev=param_variational_noise, dtype=param.dtype.base_dtype)
                    return param + noise

            param = self.network.cond_on_train(
                fn_train=_apply_var_noise,
                fn_eval=lambda: param,
            )

Specifically, check the code of gradient_checkpoint_scope and prepare_gradient_checkpointing.

I know that people also do gradient checkpointing in PyTorch, but I don't know exactly how that works.

[NeoLegends]

There is a gradient checkpointing API in PT: https://pytorch.org/docs/stable/checkpoint.html

It even saves/restores the RNG state so we could do Dropout in there. I'm not sure the RNG state there can be made explicit, but it seems suitable in all the other ways.

I saw you asking in the PT issue about JAX: the RNG there is by definition stateless, and follows a design where the RNG seed has to be threaded through the code and explicitly "split" to make new seeds.

Copying from https://jax.readthedocs.io/en/latest/jax.random.html:

seed = 1701
num_steps = 100
key = jax.random.key(seed)
for i in range(num_steps):
  key, subkey = jax.random.split(key)
  params = compiled_update(subkey, params, next(batches))

[NeoLegends]

It seems to me the API PT exposes for gradient checkpointing could be used as the RF frontend API and for the associated TF-backed implementation as well?

[albertz]

There is a gradient checkpointing API in PT: https://pytorch.org/docs/stable/checkpoint.html

Yea that is what I referred to when we talked about it. But I need to check it more how it is done there. Specifically, I'm still not exactly sure how I get what I want: that the dropout outputs are not stored but recomputed.

I saw you asking in the PT issue about JAX

No, I did not ask about JAX in there?

[NeoLegends]

Yea that is what I referred to when we talked about it. But I need to check it more how it is done there. Specifically, I'm still not exactly sure how I get what I want: that the dropout outputs are not stored but recomputed.

Yeah it would seem to me like applying only the dropout operation within the gradient checkpointed context might not be enough, but one would have to move more of the layer functionality into the checkpointed/recomputed area? Is this what you're referring to?

[albertz]

Yeah it would seem to me like applying only the dropout operation within the gradient checkpointed context might not be enough, but one would have to move more of the layer functionality into the checkpointed/recomputed area? Is this what you're referring to?

I don't know how this works. I don't want to recompute whatever comes after the dropout. I just don't want that it stores the dropout output in memory for the backprop, i.e. that it recomputes the dropout.

[albertz]

(Note, I made a separate issue just for the gradient checkpointing aspect in PyTorch: #1552. So this issue here can just focus on the RF specific question on how to implement weight dropout (or also weight noise / variational noise).)