GRU implementation details

[mkschleg]

There are multiple variations of the gated recurrent unit tied to different versions of the arxiv paper. The variation only changes where the reset vector (r) is applied (either before or after the matrix operation). Currently, we are implementing the version found in v1 of the paper, which follows from what PyTorch does . Tensorflow implements these different variations, with the v3 version being the default. Another driver of this decision is CuDNN seems to only support GRU from v1 of the paper.

Would there be interest in a PR to add the v3 variation?

If we were to add v3 (and maybe other variations), we should likely keep v1 as the default (so as not to break assumptions). But I think it is unclear in the literature exactly what people are using when using GRUs. This paper from Yoshua Bengio's group seems to use v3 to empirically study GRUs. But everyone using Tensorflow or Pytorch are likely using v1 of the GRUs. Tensorflow says the default is v3, but it looks like they are setting the argument reset_after=true. This to me means they are actually using v1 (I haven't dug into their code to figure out what is doing yet because I am never a happy camper when looking at tensorflow...).

[DhairyaLGandhi]

That sounds interesting! If you have an implementation of the different versions it would be great to have a PR.

[ToucheSir]

This came up recently elsewhere, thanks for digging into it! Unsurprisingly, it seems like CuDNN is the bottleneck.

[mkschleg]

I mentioned this in slack, so that might be what you are thinking about. I can work on a PR for this.

[CarloLucibello]

Not strictly related, but we also miss bridging v1 to cudnn right now. From what I remember from @jeremiedb's posts, we would need to move away from the vector of vectors representation of a series to a matrix representation

[ToucheSir]

Yeah, I don't see us trying to re-integrate with CuDNN unless the RNN interface itself undergoes some major changes. And even then, there may be a way to get comparable performance without bending over backwards to accommodate it like other frameworks have.

[mkschleg]

I do wonder how much performance we are leaving on the table by not using CuDNN. We might be able to get it back with some specialized kernels written in Julia? It might also be related to how memory is laid out. But I'm not sure what CuDNN offers in terms of optimizations over what we can do ourselves. The biggest hurdle would be how we assume data is presented to a network (i.e. array of matrices in Flux vs a 3 dimensional array for CuDNN).

[mkschleg]

In any case, there are a ton of recurrent architectures out there. And a lot of variants for LSTMs/GRUs (which have support in Pytorch/Tensorflow). It might be useful to figure out how to incorporate new variants of existing models without adding too much complexity. Flux has been a savior for developing novel recurrent architectures, mostly because you can very easily look at the update functions and know how the entire stack is working. So maintaining simplicity where possible would be a must.

The obvious answer is to just have a new type for each variant and have a more complex constructor or just many constructors. But maybe there is another way?

[ToucheSir]

#1367 (comment) has some performance numbers between the old CuDNN path and current implementation.

As I noted on that thread, Flux is not the only library that eschews CuDNN for its RNN implementation. Certainly they are much easier to implement efficiently than e.g. convs because the only ops involved are simple matmuls and pointwise broadcasting.

[DhairyaLGandhi]

There's a bunch of optimizations to be had with either approach. It's always a better idea to have multiple backends as well for a ground truth implementation.

@mkschleg what kinds of variations were you thinking about? It's likely that we can share a lot of code and API, so we should avoid many complicated constructors. One key is that we make the API transparent - and be explicit about what we will do with user inputs, and what we expect from them. Usually, we don't need sentinel/ magical values, which would be good to maintain. We also likely want a fairly flat type hierarchy, so there may be multiple types involved, but we can break down the problem so as to maximize how many of these can share interfaces/ behaviour. As you said, maintaining simplicity (and having generic but performant code) is a must.

[mkschleg]

This paper lists a bunch. But I'm not sold we actually need to support all the variants, just some. Looking at what Pytorch and Tensorflow support might be a place to start. Another concrete place in the literature to look could be the competitors in this paper.

My research focus doesn't tend to go outside of GRUs and LSTMs as I tend to use them as competitors when developing RL specific recurrent architectures. So I'm less familiar with what is currently becoming popular, but it might be good to think through how we would want to structure this if there was interest in expanding the included architectures in Flux.

I will just start with v1 vs v3 for the GRU and we can go from there.

[mkschleg]

@ToucheSir It seems that comparison was made without a fused 3d array, no? I'm sure the optimizations to be had might be more apparent when using the curnn kernel as intended (likely in how memory is handled).

I'm not sure how we would back that with the current paradigm Flux has (i.e. use broadcasting to iterate over time steps). Maybe we could have a special function that does this? But that seems like it might be hard to design to work generally across the library.

[ToucheSir]

Correct on all points. The old CuDNN path was making pretty poor use of the library because it fed in one timestep at a time. Unfortunately, there's not a consensus on how to (re)design the RNN interface to work more efficiently with 3D+ input data either. That said, if anyone is interested in drafting a proposal for it, I'd be happy to help review and give feedback :)

[jeremiedb]

I see 2 different components to the discussion.

1. Expanding the available recurrent cells provided by flux (currently RNN, LSTM, GRU)
2. Leveraging CUDNN for performance.

Regarding 1, I think that current form of the recurrent structure in Flux make it pretty much as straightforward as it gets to define a new/custom recurrent structure. That said, expanding recurrent cell zoo with new ones would be nice. When it comes to variations on LSTM or GRU, I guess having a constructor that takes some extra input configuration options to generate the appropriate structure such as {GRUv3} or {LSTMpeephole} may be the simplest way to go.

As for 2, in the context of an "single time step" paradigm as currently implemented in Flux, I'd disagree with the reliance upon CuDNN as it doesn't improve performance and as for the comfortmity to ground truth, I'd personally prefer the safety from having a single code definition of the cell which dispatches on cpu or gpu thanks to CUDA.jl rather than having an independent implementation from CuDNN taking over when on gpu. Moreover, given the lightweight implementation of the recurrent cells, I feel such approach is more transparent and easy to assess by the user. It's also my perception that such transparent cpu/gpu execution model is part of the attractiveness of the Julia/Flux ecosystem, more than CuDNN bindings.

That said, having a support for a 3D fused sounds like a good addition to leverage optims that are then available. I'd see something of the form `FusedRNN/FusedLSTM...), which would be restricted to 3D inputs and dispatch to CUDNN. As for the cpu support, maybe the simplest to start with would be to define the fused function as a loop over the seq_length dimension, thus reusing the current non-fused implementation (rather than trying to build an optimised 3D cpu implementation)?

[ToucheSir]

I wonder if we could intercept the broadcasting of certain RNN layers on 3d inputs and use that to dispatch to backends like CuDNN. Layers/functions that map or fold (in either direction) an RNNCell over a sequence would also be an option.

[jeremiedb]

intercept the broadcasting of certain RNN layers on 3d inputs and use that to dispatch to backends like CuDNN

Could you expand a bit on what you would have in mind? My first thought was that if CuDNN is involved, then we are dealing with gpu execution and in such case, CuDNN's RNN/LSTM/GRU is already set to consume the entire 3D array as a single input and apply the full stack of layers onto it. So if CuDNN is involved, I don't see a case for broadcasting as this would seem to brake its optimisation which is the motivation for its usage.

Roughly speaking, my take on a 3D support would go along the lines:

function (m::FusedLSTM(x::CuArray{T,3}))
    cudnn_LSTM(x)
end

function (m::FusedLSTM(x::Array{T,3}))
    # init out    
    for layer in m.layers # support stack of layers as in CuDNN
        out = [layer(x[:,seq,:]) for seq in 1:size(x,2)]
    end
    # need to reconcat into 3D output
    return out
end