neelnanda-io
commented
1 year ago Some caveats: BERT is encoder-only, not encoder-decoder. It should be a HookedBERT or HookedEncoder class, in my opinion.
Note that post-LayerNorm means "LayerNorm applied after the layer". Confusingly, the LayerNormPre class means "LayerNorm with weights folded into the subsequent layer". Pre-LayerNorm when used in papers means LayerNorm at the start of a layer. shrug Naming conventions are hard man
On Wed, 26 Apr 2023 at 11:01, Joseph Bloom @.***> wrote:
[image: Screenshot 2023-04-26 at 8 00 40 pm] https://user-images.githubusercontent.com/69127271/234541975-b5f48a88-0354-4e2c-a060-0deb52d60de8.png Proposal
A glaring gap in models supported by TransformerLens is the lack of support for encoder/decoder models like BERT and T5. In order to rectify this, we need to build support for encoder-decoder models, cross attention and likely a few other details. Motivation
- I've been asked a number of times by people whether we have T5 or BERT support coming.
- I think this is a reasonable contribution for people skilling up. You can ask for help or collaborate with other people working on building it.
- I (Joseph) and likely others will be able to give feedback on your work and pitch in if needed!
Pitch
In order to complete this card we are going to build a class called HookedEncoderDecoder which can be instantiated and have the BERT or T5 weights from huggingface https://huggingface.co/bert-base-uncased loaded in. This is the same strategy the package uses for other models so there's nothing fundamentally new, only components we haven't yet built with hooks. Once we have a class that has corresponding weights as needed, we will want to ensure all components are "hooked" correctly, enabling methods like .run_with_cache to work. It's not as complicated as it sounds to get it working.
Possibly the "best" version of this is tricky but we're willing to tolerate some code duplication as long we're confident the implementation is accurate (ie: tests show the hooked activations are correct and inference is identical to what we get without using TransformerLens.
Tasks:
- Make a fork for BERT support (keep it synced if building this takes more than a moment)
- Write some tests to be used during development, these should cover any novel architecture components these will involve during development. I'd start with two placeholder tests. 1. for do you have a model architecture which you can load weights into. and 2. if you had that architecture, would it produce identical results to the hugging-face model.
- Load the BERT model, look at each component. Maybe read the BERT paper if needed to get context. Make a list of every nn.module that we don't currently support. Some items that will be on the list. 1. a linear block before the unembed, post-layerNorm (can't be unfolded RIP), sinusoidal positional embeddings. T5 might have other stuff.
- Keep in mind you might need a new config class to configure the entire architecture. If so, build it.
- If you can pass the test for the model doing inference correctly and show that you can use .run_activation_cache and all the other methods that work on hooked transformer (within reason), then we're likely ready to merge.
Some things to note:
- You should refine the to do list and add detail as you go.
- Writing tests should make your life easier and create feedback cycles. It will also make it easier if multiple people want to collaborate on this.
- A previous PR https://github.com/neelnanda-io/TransformerLens/pull/82 attempted to provide BERT support is here. Unfortunately with drift it was too complicated to Merge. To avoid this, I recommend either syncing regularly or getting someone else to finish the PR if you can't.
- Possibly quite useful, a link to the original ARENA materials where we did this https://github.com/callummcdougall/arena-v1/blob/main/w2d2/solutions_build_bert.py without the hooks during the bootcamp
Additional context
- BERT paper https://arxiv.org/abs/1810.04805
- BERT on huggingface https://huggingface.co/bert-base-uncased
- Related T5 Model issue https://github.com/neelnanda-io/TransformerLens/issues/94
Checklist
- I have checked that there is no similar issue https://github.com/Farama-Foundation/MiniGrid/issues in the repo ( required)
— Reply to this email directly, view it on GitHub https://github.com/neelnanda-io/TransformerLens/issues/258, or unsubscribe https://github.com/notifications/unsubscribe-auth/ASRPNKJUT3735X5MX6SOC7DXDDW7VANCNFSM6AAAAAAXMGIZMA . You are receiving this because you are subscribed to this thread.Message ID: @.***>
rusheb
alan-cooney
jbloomAus
fzyzcjy
Proposal
A glaring gap in models supported by TransformerLens is the lack of support for encoder/decoder models like BERT and T5. In order to rectify this, we need to build support for encoder-decoder models, cross attention and likely a few other details.
Motivation
Pitch
In order to complete this card we are going to build a class called HookedEncoderDecoder which can be instantiated and have the BERT or T5 weights from huggingface loaded in. This is the same strategy the package uses for other models so there's nothing fundamentally new, only components we haven't yet built with hooks. Once we have a class that has corresponding weights as needed, we will want to ensure all components are "hooked" correctly, enabling methods like .run_with_cache to work. It's not as complicated as it sounds to get it working.
Possibly the "best" version of this is tricky but we're willing to tolerate some code duplication as long we're confident the implementation is accurate (ie: tests show the hooked activations are correct and inference is identical to what we get without using TransformerLens.
Tasks:
Some things to note:
Additional context
Checklist