MultiHeadAttention layer

[google-deepmind]

MultiHeadAttention layer

[trevorcai]

You can see an example implementation of MultiHeadAttention here:
https://github.com/deepmind/dm-haiku/blob/5457f4b0bf199c336be9021cb3010a1efd2d95da/examples/transformer/model.py#L26

I'd recommend forking that & modifying/simplifying it for your use case.

To explain our thought process a little more:

Since we're on the hook for maintaining modules we put into Haiku, we try to be careful about making sure that the modules inside are canonical and generally agreed upon (at least within DeepMind).

Admittedly, MultiHeadAttention as used in transformers is pretty canonical (except initialization). But any architecture changes people make to attention (e.g. Transformer-XL, sparse attention, all-attention) require changes to the module.

Because most people using transformers at DeepMind want to experiment with changes like these to attention, we think consumption of the canonical MHA will be low (unlike e.g. an LSTM, which very few people are modifying these days).

We've similarly discussed whether we should promote a transformer implementation to hk.nets, like we have for ResNet-50.
The annoying thing about transformers is that people's idea of what "canonical" is varies based on their use case - where does the layer norm go (residual or main path?); how should we initialize, etc.

Overall, we think there's sufficient movement in the field re: ideal transformer architecture and attention mechanisms and sufficient diversity re: expected canonical implementation that we've decided not to make a canonical one. Instead, we recommend folks to fork from our examples, or other existing codebases.

Does that make sense? WDYT?

[cgarciae]

I think this makes sense to not to make it official until its stable, yet given PyTorch and Tensorflow Addons (I contributed this one) do have at least placeholder implementations it would be good to provide something. Maybe consider an hk.experimental module or a haiku-addons repo?

[trevorcai]

Thanks for the feedback - I appreciate it!

I'm personally onboard with the value proposition of tensorflow-addons - but I don't think either Tom or I have bandwidth to support a Haiku equivalent right now.
(My understanding is that maintaining tf.contrib was quite a bit of work for the TF team.)

It seems like forking (or directly importing from) the examples hits 90% efficacy for our core use case (supporting deep learning research in JAX), so I'm inclined to keep the status quo here.

Please reopen if you think this is misguided!

[trevorcai]

I've been convinced by helpful colleagues that my thinking on this is misguided! We're standing by the "no canonical transformer" thinking, but are open to having a MHA implementation in core Haiku.

[cgarciae]

@trevorcai Nice! Can I start a PR with this?

[trevorcai]

We have some well-proven implementations internally. When I get the chance, I'll move one into Haiku core.

[trevorcai]

Sorry for the wait, we upstreamed one here:
https://github.com/deepmind/dm-haiku/blob/610a574c54c78ca71a9a1213bdaff9de189555d4/haiku/_src/attention.py