Skip to content

We write sample code for two tower models for retrieval and add RLHF/RLAIF style alignment with a ranking model to make the retrieval more aligned with the ranking model on top

License

Notifications You must be signed in to change notification settings

ms207/two_tower_models

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

30 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

two_tower_models

This is companion repository to two tower models, a commonly used approach to retrieval / candidate generation in recommender systems. The goal of this repository is to show how to increase the alignment of retrieval with ranking.

two_tower_schematic

This shows a sample implementation of two tower models in PyTorch.

This is text-book implementation of self attention and positional encodings to summarise user history for the user tower.

This uses the UserHistoryEncoder above to add to the user tower.

Since we are using net user value to estimate the relevance, this uses a weighting approach to position debiasing. There are other approaches as well like logit debiasing via shallow tower that we recommend trying out.

This extends TwoTowerWithPositionDebiasedWeights from above and implements a pointwise ranking module, broadly described in Revisting neural accelerators. We chose this since implementing a light ranking in the retrieval system is an intuitive way to increase consistency with main ranking.

Then we extend this to adding knowledge distillation from ranking models. This is similar to the approach described in How to reduce the cost of ranking by knowledge distillation

In this step we add a further layer of funnel consistency where, inspired by RLHF, we use the ranking model as a "reward model" and learn how to make the retrieval more aligned with the ranking model.

References

  1. Seminal paper on two tower models in Youtube

Disclaimer: These are the personal creations/opinions of the author(s). Any artifacts, opinions stated here are theirs and not representative of their current or prior employer(s). Apart from publicly available information, any other information here is not claimed to refer to any company including ones the author(s) may have worked in or been associated with.

About

We write sample code for two tower models for retrieval and add RLHF/RLAIF style alignment with a ranking model to make the retrieval more aligned with the ranking model on top

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 100.0%