Multihead attention

[OneAdder]

Hello, Milan! I hope, I'm not bothering you too much with my pull requests, but this is a good one. At this stage it is a draft of MultiHead Attention. It cannot be merged until work on input2d_layer and linear2d_layer is completed.
Implementation of dropout would also help improve MHA, but it can be added later.

MultiHead Attention

MultiHead Attention is the main component of Transformer architecture, which is the most advanced modern approach in the area of Natural Language Processing, as well as some other areas.
Here I propose an implementation based on the Transformer article. It works and its output conforms with SOTA implementation in PyTorch.

Python Reference

https://github.com/OneAdder/neural-fortran-references/blob/main/self_attention.py

[ricor07]

Hello Michael, I saw your pull requests and I think what you do is very interesting. Could you take a look at mine? milancurcic#2

What you have to look here is not the locally connected 1d layer but rather the reshape architecture I am trying to make. Your help would be very appreciated. Thanks

[milancurcic]

Amazing, thank you! Yes, let's wrap up Input2d in #198, then Linear2d in #197, to avoid the 3d crutch.

[ricor07]

@OneAdder i see you are talking about a 2d handling. Would you like to make this together? I have to make this as well since I'm implementing a conv 1d layer

[milancurcic]

Hi guys, thanks for pushing this forward. Today I'm finishing a lot of busywork with some proposals so next week I'm getting back more actively with neural-fortran work, and will be able to contribute to the reviews more actively

[OneAdder]

@ricor07 Great idea! But we have a problem with generics again. The issues is that a predict_2d and predict_1d_batch have the same rank of input. I think we should simply make a separate generic for predict_batch. @milancurcic your thoughts?

[ricor07]

Yes, I think we can make a generic predict. But I suggest you to create a new branch

[milancurcic]

I think it's fine to make predict_batch its own generic name because it's getting in the way. 👍

[OneAdder]

@milancurcic Done, here: #198
@ricor07 There is still one piece of the puzzle missing: a general flatten layer with interfaces for both 3D and 2D. Do you want to implement it? Or I can do it together with linear2d layer

[ricor07]

You can make it. I'll work on maxpool

[OneAdder]

@milancurcic I think it's ready for review! I'll add the complicated example later. At this stage I added a simple example that converges nicely and doesn't require datasets and extra deps

[OneAdder]

That is intentional that there is no plumbing for this one yet. I suggest that we add it at later stage when we have more components for seq2seq models. At this stage it can be added like this: without any public access

[OneAdder]

Suggestion for future: create nf_utils.f90 (or similar) and put this procedure there

[OneAdder]

attention_mask is not accessible to the users by design at this point. It will be used by transformer decoder later and I'll add corresponding logic later

[jvdp1]

@OneAdder Nice set of PRs. I added a few suggestions. Feel free to ignore them.
Furthermore, I usually try to avoid reshape and using pointers instead, mainly for performance reasons. I think this approach could be used at some places, but it could be for a next PR.

[milancurcic]

@OneAdder:

self_attention_layer and cross_attention_layer both extend multihead_attention_layer, but only self_attention is exposed in the user-facing module nf and in constructors.

1. Should cross_attention be made accessible to the user like self_attention is?
2. Is multihead_attention intended to be used only internally as a building block, or is there value to making it part of the public API?

The answers to these questions will also guide which layers to include in the layers table in the README.

[OneAdder]

@jvdp1 Great idea! I think we should extend it even further and make smth like nf_utils with a subroutine that does those inplace reshapes. So that we could use it whenever possible. But I think this PR is already pretty long. So, I suggest we do it later

[OneAdder]

@milancurcic

1. Yes, but at this point we don't have any layers to make it useful. Therefore I made it ready to be included but not actually included as of yet. I suggest it makes it to public API later.
2. I don't know any scenarios where Attention is neither self nor cross. If there are, I think they should also be separate classes extending multihead_attention_layer, as it itself doesn't implement forward and backward methods

[milancurcic]

Thank you!!