Feature request: weight normalization

[appierys]

Is it possible to incorporate Weight Normalization (https://arxiv.org/abs/1602.07868) into tensorflow itself?

https://github.com/openai/weightnorm/tree/master/tensorflow

[jart]

@rmlarsen Any thoughts on this feature request?

[yaroslavvb]

@appierys that's just a link to a file with helper utilities, there no actual implementation of the algorithm there

[appierys]

@yaroslavvb I do think the file contains the implementation of WN. For example, here is the implementation of weight norm in a fully connected layer: https://github.com/openai/weightnorm/blob/master/tensorflow/nn.py#L166

[yaroslavvb]

@appierys that indeed looks like one of the formulas, but some other ones are missing. (ie, equation for updating the step length parameter)

[discoveredcheck]

Hi
Seconding @appierys 's request, we have also been experimenting with weight-normalisation in RNNs with great results. It seems to be significantly helpful in stabilising training and reducing its dependence on hyperparameters. Seeing as this has become a default RNN configuration for us, we are quite keen on integrating it with the tensorflow rnn machinery.

Our implementation is based on
https://github.com/openai/generating-reviews-discovering-sentiment/blob/master/encoder.py

which requires a few lines of additional code in tensorflow core. This is because in the current setup (dynamic_rnn etc.) RNN weights are created via the cell's call() function which is called from within tf.while(). Based on the above link, weight-normalised cells require weight creation and normalisation before entering the while loop.

This may not be the best way to implement it, but for the moment we have sent in a pull request:
#11573

We'd be keen to hear your thoughts and/or to engage in further testing/discussing a different implementation structure.

[tensorflowbutler]

It has been 14 days with no activity and the awaiting tensorflower label was assigned. Please update the label and/or status accordingly.

[jart]

We'd be interested in considering a contribution adding this support.

[llan-ml]

@discoveredcheck Thanks for your contribution, and I have read your PR about "WeightNormLSTMCell".

I notice that you define variables in the functions self._linear and self._normalize, and do not set tf.AUTO_REUSE flag. So, I wonder how to ensure the variables are reused when self.call are invoked multiple times? I test your code, and it works well. I just want to know how it works.

[discoveredcheck]

@llan-ml Apologies for the late reply. WeightNormCell follows convention from other RNNCellsand relies on the caller to use call() only once. From the doc - "Note that while_loop calls cond and body exactly once (inside the call to while_loop, and not at all during Session.run())"
https://www.tensorflow.org/versions/master/api_docs/python/tf/while_loop

Hope this helps.

Ashwini

[llan-ml]

Thanks for your explanations @discoveredcheck.

On the other hand, to enable weight normalization for dense and convolutional layers, can we just add the weight normalization in the self.build() method of tf.Layer.
For example, in the scenario of Dense layer, we can add weight normalization by modifying https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/layers/core.py#L132 as follows:

kernel = self.add_variable("kernel", ...)
if self._weight_norm:
    g = self.add_variable("wn/g", ...)
    self.kernel = nn_impl.l2_normalize(kernel, axis=0) * g
else:
    self.kernel = kernel

Similar modifications could be also made for tensorflow.python.layers.convolutional._Conv.
Do you have any suggestions?

[discoveredcheck]

Yes, adding the weight-norm update in self.build() seems to be the correct pattern. It is also correct to omit out with tf.control_dependencies(None) compared to the RNN weight-norm code.

Your link, however, revers to L152 which is inside self.call(). I assume you mean L132?

[llan-ml]

Yeah, it should be L132, and I have corrected the link.

[llan-ml]

I write some simple wrappers for dense and conv2d layers with weight normalization.
https://github.com/llan-ml/weightnorm

[seanpmorgan]

Duplicate of #14070

This has been implemented in tensorflow/addons. Issue can be closed.

Code
Example

[facaiy]

Thanks, Sean :-)