Adding kNN language modeling and Machine Translation

[yjernite]

🌟 Adding the kNN-LM and kNN-MT models

Model description

The kNN Language Model and Machine Translation leverage additional data at test time by looking up examples in the data store that are similar to the test example and using them to inform the prediction. This leads to decreased perplexity for language modeling and increased BLEU for MT.

Ideally, we'd have a general KNNDecoder class which can use any model in the library along with a 🤗datasets data store, similarly to what RAG does currently.

Open source status

• [x] the model implementation is available: for kNN-LM, an implementation using Fairseq and FAISS can be found here
• [x] the model weights are available: the methods use pre-trained models which are already in the library, GPT-2 and the Facebook WMT'19 models
• [x] who are the authors: first author @urvashik

[sshleifer]

are you interested @stas00 ? this will be a tough one but cool!

[stas00]

This definitely sounds interested, thank you for thinking of me, @sshleifer!

But I also want to make sure that since I will be shortly starting on a big project which could take time, if someone wants to work on this one by all means do that.

So I propose to assign me tentatively but if another developer/contributor wants to do it then they are welcome to assign themselves. Thank you!

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[stas00]

@LysandreJik, can we have some kind of way to tell the stale bot off for feature requests?

Surely if an issue is a bug and it has become old and everybody lost interest in it, it's probably ok to close it automatically, but setting such a short life span for feature-request type of issue just calls for a wasted time to get the bot off one's back and then one is too busy to ping the issue and the feature request disappears.

Thank you!

[stas00]

Looks like I never had the time to start on this one. I have unassigned myself for now as I honestly won't have time to work on it any time soon. And now it's open to others to work on.

[ncoop57]

I think using this more efficient version of the KNN-LM might make a future KNNDecoder class more practical: https://github.com/jxhe/efficient-knnlm. Maybe it can be a flag, but anyways it does dimensionality reduction on the embeddings that are generated to reduce the size of the datastore.

I'm seeing the following task to implementing this in the library:

• [ ] The construction of the faiss index using HF datasets built in version and running the model over the examples with some given context length and grabbing the final hidden state for the embedding
• [ ] Optionally reducing the embedding using PCA or UMAP
• [ ] Get the embeddings of the input and performing the nearest neighbor look up, also using the HF datasets built in version
• [ ] Interpolating the softmax predictions as done in the paper and using the new softmax to sample from with whatever other strategies are enabled.

Is there any other steps anyone can think for implementing this?

[urialon]

Hi @yjernite @sshleifer @stas00 @ncoop57 !

I just released an implementation of kNN-LM and our new RetoMaton model (ICML'2022) based on the transformers library at https://github.com/neulab/knn-transformers .

All previous implementations that I know about of kNN-based models are:

• Implemented in fairseq.
• Forking/duplicating the entire library (fairseq) to implement their modification. This makes it very hard to work with, fetch new updates to the library, and incorporate into other projects.

What I really like about my implementation is that I used hooks (using layer.register_forward_hook) to avoid forking the entire library, making the KNNWrapper standalone, and it can simply wrap any language model by:

knn_wrapper = KNNWrapper(...) # or: RetomatonWrapper(...)
knn_wrapper.break_into(model)

That's it! The model now internally uses kNN-LM or RetoMaton. In theory, this implementation should also make it very easy to reproduce the kNN-MT experiments, but I haven't tried yet.

Let me know if anyone has questions or wants to implement this inside transformers.

Best, Uri