[Proposal] Documentation: Map the Act Names to the Transformer

[JuVogt]

Proposal

Create a figure that maps the act names to the transformer architecture.

Motivation

Names are just conventions. I find it hard to get the exact position within the transformer block just from the act name. I.e. the resid_pre might be before the split happens or before the merge happens. So I put it in context to the other act names and work by exclusion process or modify it to see what values will change.

Pitch

I suggest using the images from the Vasvani paper and adding labeled arrows pointing to the hook positions.

Alternatives

A list or table of (act name, description) pairs.

Checklist

• I have checked that there is no similar issue in the repo (required)

[bryce13950]

@JuVogt Do you have time to handle this issue?

[tjbai]

I could put together something this week as a first PR for this project

[JuVogt]

I am willing to contribute as well, but I am currently out of time, sorry for that. I can come back after I finish my thesis at the end of the year and design something, but a first sketch would definitely help. Maybe I could then add a list with the act names including some more information about i.e. the dimensions and calculations behind it if someone already contributed a sketch or vice versa.

Also, I could add some more documentation with minimal examples beside the colabs that I think would help me in the beginning.