[Replication] "Transformer Feed-Forward Layers Are Key-Value Memories"

[quinn-dougherty]

Background

Feed-forward layers constitute two-thirds of a transformer model's parameters, yet their role in the network remains under-explored. We show that feed-forward layers in transformer-based language models operate as key-value memories, where each key correlates with textual patterns in the training examples, and each value induces a distribution over the output vocabulary. Our experiments show that the learned patterns are human-interpretable, and that lower layers tend to capture shallow patterns, while upper layers learn more semantic ones. The values complement the keys' input patterns by inducing output distributions that concentrate probability mass on tokens likely to appear immediately after each pattern, particularly in the upper layers. Finally, we demonstrate that the output of a feed-forward layer is a composition of its memories, which is subsequently refined throughout the model's layers via residual connections to produce the final output distribution.

What to Replicate?

The paper

Modifications

Do it with large autoregressive models

Related Papers/Frameworks

Paper: https://arxiv.org/abs/2012.14913

EleutherAI's language modeling framework: https://github.com/EleutherAI/gpt-neox

Code from the paper: https://github.com/mega002/ff-layers/

[StellaAthena]

@quinn-dougherty Who is working on this project? Where is the code and what are the goals?

[evhub]

I definitely agree that Geva et al. is in need of a replication--also Dai et al., which builds on Geva et al. but seems even more sketchy. See my conversation with Paul Christiano here for some thoughts on things that seem sketchy about Dai et al. (and also sort of by extension Geva et al.) that would be good to try and address in a replication.

[MicPie]

Hi, I’ll have a look into the publication and the accompanying repo (https://github.com/mega002/ff-layers) to setup GPT-NeoX for it. :-)

[StellaAthena]

@MicPie Did anything ever happen with this?

[MicPie]

Sorry, I was running out of time and I'm not sure if this project is already taken over by somebody else (there was a discussion some weeks ago).

[StellaAthena]

Sorry, I was running out of time and I'm not sure if this project is already taken over by somebody else (there was a discussion some weeks ago).

Not as far as I know. Someone took a look at it but nothing came of it.

[rokosbasilisk]

I am starting on this.