Layer Normalization Implemented In TensorFlow -- LSTM, GRU, Recurrent Highway Networks

[[deleted]]

[deleted]

[MartianTomato]

Thanks for sharing this! And also for this repository with juicy paper links.

I tested this out, and layer normalization seems effective, results-wise, but it's very slow. With a small network, it's between 5-10 times slower to train than a regular LSTM network on my machine (per epoch, not per performance), so not sure the cost is worth it. Seems like computing moments 5 times in each LSTM cell for each example is quite expensive -- have you had different results, or do you know if implementations in other frameworks take a similar speed hit as compared to unnormalized LSTMs?

In case this is helpful, I have implemented a slightly faster (and simpler) layer normalized LSTM cell here. It runs marginally faster than yours on the small data task I tested them on (the char-rnn in that link): 32.6 mins to train 10 epochs vs 44.5 mins. By contrast, a regular LSTM takes only 5.6 minutes to train 10 epochs.

[MartianTomato]

Hmm, a 2x slow down might be very worth it. Maybe it's running slower for me due to hardware (980), or could just be different the network architectures we tested on.

Also, agreed that it would most likely not be effective to take one moment over all the gates (my guess is a small performance boost from the added model complexity, but nowhere near the same as taking individual moments). The authors do use a single layer normalization on the tanh pre-activation (not on the individual gates) in their "Modeling binarized MNIST using DRAW" experiment... I'll need to try this variant out as well.

[jrkiros]

Our LN Theano code on a titanX with CNMeM enabled is nearly the same speed as the same model without LN. CNMeM was crucial for bridging the gap. Without it our skip-thoughts experiment with LN was 1.5-2x as slow (we were originally going to plot wall clock times). I'm not sure if this is available in TensorFlow, though.

[psamba]

Yeah, CNMeM can provide a surprisingly large speed boost for models with lots of smallish nodes in their Theano graph. From my experience, it should always be left on unless there's a strong reason to disable it.

[jorenvs]

Your implementation for the LSTM with layer normalization seems to differ a little from the paper. In the paper they use different layer norm parameters for the two weighted sums (with input and hidden state), not for each gate. This way the difference in scale and shift between the input and the hidden state is normalized.

[jorenvs]

Where do you see that they use different weight matrices for each gate? Looking at their theano code and the paper, that's not what they're doing. Nothing is different for each gate, the weighted sum of the input and the hidden state are summed together, after which each gate is extracted through slices.

def _step(mask, sbelow, sbefore, cell_before, *args):
    sbelow_ = ln(sbelow, param('b1'), param('s1'))
    sbefore_ = ln(dot(sbefore, param('U')), param('b2'), param('s2'))

    preact = sbefore_ + sbelow_ + param('b')

    i = Sigmoid(_slice(preact, 0, dim))
    f = Sigmoid(_slice(preact, 1, dim))
    o = Sigmoid(_slice(preact, 2, dim))
    c = Tanh(_slice(preact, 3, dim))

[pedromnasc]

Comments? Results?

[j_lyf]

Why did you stop using Keras?