Transformer Memory as a Differentiable Search Index (Machine Learning Research Paper Explained)
#dsi #search #google

Search engines work by building an index and then looking up things in it. Usually, that index is a separate data structure. In keyword search, we build and store reverse indices. In neural search, we build nearest-neighbor indices. This paper does something different: It directly trains a Transformer to return the ID of the most relevant document. No similarity search over embeddings or anything like this is performed, and no external data structure is needed, as the entire index is essentially captured by the model's weights. The paper experiments with various ways of representing documents and training the system, which works surprisingly well!

Sponsor: Diffgram
https://diffgram.com?ref=yannic

OUTLINE:
0:00 - Intro
0:45 - Sponsor: Diffgram
1:35 - Paper overview
3:15 - The search problem, classic and neural
8:15 - Seq2seq for directly predicting document IDs
11:05 - Differentiable search index architecture
18:05 - Indexing
25:15 - Retrieval and document representation
33:25 - Training DSI
39:15 - Experimental results
49:25 - Comments & Conclusions

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

Abstract:
In this paper, we demonstrate that information retrieval can be accomplished with a single Transformer, in which all information about the corpus is encoded in the parameters of the model. To this end, we introduce the Differentiable Search Index (DSI), a new paradigm that learns a text-to-text model that maps string queries directly to relevant docids; in other words, a DSI model answers queries directly using only its parameters, dramatically simplifying the whole retrieval process. We study variations in how documents and their identifiers are represented, variations in training procedures, and the interplay between models and corpus sizes. Experiments demonstrate that given appropriate design choices, DSI significantly outperforms strong baselines such as dual encoder models. Moreover, DSI demonstrates strong generalization capabilities, outperforming a BM25 baseline in a zero-shot setup.

Authors: Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, Donald Metzler

Links:
Merch: http://store.ykilcher.com
TabNine Code Completion (Referral): http://bit.ly/tabnine-yannick
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://ykilcher.com/discord
BitChute: https://www.bitchute.com/channel/yannic-kilcher
LinkedIn: https://www.linkedin.com/in/ykilcher
BiliBili: https://space.bilibili.com/2017636191

If you want to support me, the best thing to do is to share out the content :)

If you want to support me financially (completely optional and voluntary, but a lot of people have asked for this):
SubscribeStar: https://www.subscribestar.com/yannickilcher
Patreon: https://www.patreon.com/yannickilcher
Bitcoin (BTC): bc1q49lsw3q325tr58ygf8sudx2dqfguclvngvy2cq
Ethereum (ETH): 0x7ad3513E3B8f66799f507Aa7874b1B0eBC7F85e2
Litecoin (LTC): LQW2TRyKYetVC8WjFkhpPhtpbDM4Vw7r9m
Monero (XMR): 4ACL8AGrEo5hAir8A9CeVrW8pEauWvnp1WnSDZxW7tziCDLhZAGsgzhRQABDnFy8yuM9fWJDviJPHKRjV4FWt19CJZN9D4n
 This is a comprehensive paper review of the paper transformer memory as a differentiable search index. This paper is pretty crazy. It takes an entire data set and puts it into the weights of a transformer. Essentially, it trains a search engine not to search through documents, but just to give you the index of the document that matches your query, just like that. Boom. So this video is a comprehensive review of the paper. I'll explain to you what's in the paper, what it's about. And by the end of the video, you should have a good idea of the paper itself. The next video, which I'm going to release tomorrow will be an interview with the authors will dive right into the content and any criticisms and questions that I raised during the review. As always, let me know what you think in the comments. Now let's get into the video. See you around. Does your company have a lot of people labeling data? Why would you leave such an important task to close source systems or self implemented things? Training data is your most valuable asset and human labels are really expensive. Today's sponsor is diffgram, which is an open source platform centered around training data. They handle everything to do with training data, especially collecting labeling, serving and more. And it is open source so you can self host all you want. But there's one cool thing if you let them host it for you. And that is unlimited pricing, no per label annotation, no expensive servers to run, you pay once you get as much as you want. So thanks again to diffgram for sponsoring today's video, check them out using the link in the description to let them know that I sent you. All right, let's get into the video. Hello there, today we're looking at transformer memory as a differentiable search index by researchers of Google research. This paper on high level takes a search problem where you have to index documents and retrieve them. And it puts all of the corpus essentially into the weights of a transformer. So it takes the corpus and trains the transformer. And then at the end, they can just give a query to the transformer and the transformer will output the ID of the document that matches and it turns out for some data sets that they have for some settings and with some clever training and representation of the documents that can actually work, which is really crazy. This kind of speaks to multiple things such as obviously our ability to overfit on stuff, but there is some generalization here as we'll see. On the other hand also the kind of inner workings of these transformers. And lastly, what's pretty cool is that this is completely as it says differentiable, it's a differentiable search index, which means that this can be part of larger neural network architectures because it is fully differentiable and it can be trained essentially end to end at once. And that means we can potentially employ reinforcement learning agents with kind of retrieval abilities and much more things. So we'll dive into the paper, we'll see what it's about. The idea, as I said, is pretty, pretty simple. If you like content like this, then as always leave a like and tell me what you think in the comments. That's always super helpful. So as I said, they take a search problem and the search problem is essentially I have a corpus, like I have a big database of documents, right? Here is a document, here is a document and I want to build an index and an index is some kind of data structure, some kind of thing. And at the index, I can throw a query and the index will return to me an ID, a document ID that specifies which document matches my query. Usually this is done via inverted indices. So I want to tokenize my documents, split them into little tokens, which are usually words or sub words, and I want to stem them and lemmatize them and whatnot. Then I build a reverse index. So for every word like in the word in I remember which documents it appears in like document three, document five, document 11, and so on. And then once the query rolls in, I simply tokenize it as well. I go look into my inverted index and I look up all the documents. And then there's also a ranking step, which means I have to now determine which of these documents is the most relevant. And that is usually done via techniques like TF IDF features. There is a famous technique called BM 25, which is also a baseline in this paper. So this is the classic search kind of way way of doing search. If you use any search engine at all, this is being done in the background. For the most part, newer search engines are catching on, there's neural search and so on. But BM 25 is still one of the most performant things that text search has available, and also other types of search. However, there is a new push in sort of neural search. And in neural search, you're trying to take your data set. And for each document, you try to map it to some sort of a vector in vector space. And then once the query comes in, you also map the query to a vector. And for example, you compare inner products, whichever inner product is largest, that's the document that's relevant. This is just one way. This is what we usually call a by encoder method where the documents in the queries are mapped, both mapped individually. So there would be an encoder here, and there would be an encoder here, they all would output one vector, and then the vectors are compared, this could be the same encoder or different encoders for documents in query. This is just one method, there's various methods such as cross encoders, rerankers, dense retrievers, you name it. However, this method here is even more is different. So what we want to do is we want to take the corpus as such, and map that somehow into a neural network. And we're going to talk about this somehow. But we're going to train a neural network, essentially, how do we represent this, let's say represented with its layers, such that when later I feed a query to the neural network, as I already said, the ID of the document is the output of the neural network. So it doesn't output a vector that I then go and compare, it doesn't, I don't have to go and feed in query document pairs. And then I get out a score of how well they fit together, which would I would do in a cross encoder. No, the transformer, in this case, the neural network directly gives me the ID of the document, which without seeing the data at inference time. So during training, all of the data is essentially has to be mapped somehow into the weights of the neural networks, right? So somewhere in these weights, that information is stored of what the documents are. So the entire corpus is in those weights. And once I enter a query, the correct document ID can only be output, obviously, if the transformer has somehow learned what is in those documents. So that's the setup is pretty simple setup. Once you kind of see what's going on. It's, it's like a meme, right? Instead of we've been trying to neuralize search, and we've still done this two step process where we train these encoders, but then the actual search is still done using, for example, a nearest neighbor algorithm like here. But, you know, this is just the idea of, well, why don't I just ask the neural network to output the result, right, the resulting doc ID, why don't I just do that? And it turns out that can work surprisingly well. So you can do large, a couple of things here, but that's essentially it. They say right here in the introduction, they use a sequence to sequence learning system to directly map a query to a relevant document ID. They have different corpuses where they train it on on the smallest corpus, this method improves the hits at one, which means that whether the top hit is the correct one, more than 20 points from 12.4% for a dual encoder. So the baseline here is 12.4% for a dual encoder. So the baseline here is a dual encoder, what I shown whenever the there are two encoders, and they each output an embedding to 33.9%. That's a giant gain, right? That's like a 2.5x improvement. However, on a corpus that's 30 times larger performance is improved by nearly seven points, which is less. It's also respectable that performance if it is improved at all. However, I want you to notice and that's already kind of the first indication, a little bit of obviously what's going on here. On smaller data sets, this method does super duper well on larger data sets, the method doesn't do that much better than a sort of cross encoder type setup, sorry, a bi encoder type setup or a dual encoder type setup, which is understandable, right? Because the smaller the data, the easier it is to absorb it all into your weights. If the data gets larger, that obviously gets harder and harder. There's more data to go around, which means there's more error, room for error for confusion, and so on. And a classic search engine or a dual encoder is going to have a easier time in that case. But still, it's a cool paper. It's just that it kind of gets worse with the data set scale. It does get better with the model scaled though. The really exciting thing is something that I've already mentioned and they mentioned this here, all aspects, sorry about that. They say all aspects of retrieval are mapped into well understood machine learning tasks. So for example, indexing, which is building the reverted index, or even if you have the dual encoder, you need to build the nearest neighbor index, which is a hard task in high dimensions is now a special case of model training. So it's just training and incrementally updating an index becomes just a special case of model updating. So all the tasks are just tasks that we already understand from neural network training. So here is a comparison of the dual encoder method, which is the, let's say old classic neural search method, not the BM 25 retrieval, but the neural search method, and this DSI, the differentiable search index. So in the dual encoder method, what we do is we train this encoder. And in this case, they train one encoder for both the queries, as well as the documents. And what we try to do is we are going to try to use some form of contrastive loss. If we actually have query document pairs, what we can do is we can try to get the documents, the query and the document that go with each other to be close together, while making the documents that are unrelated to each other be far apart. So this is some sort of contrastive loss, obviously, at inference time, what we're going to do is we have a query, we put it through the encoder, we get its embedding, and we do a maximum inner product search through our entire vector space of our of our indexed data set, and we get a ranked list. So it's kind of this two step approach with building these indices in between, and with the training objective, that is not directly what we want. It is a proxy objective, because of the because of the algorithm later needs it the inner product search, but it is not actually what we want. So let's just train what we want. In the DSI, in the differentiable search index, I simply feed my query along with I simply feed my query essentially to in some form to the system. And the system outputs directly which document is relevant for the query. So the way they train it, and this is one way they train it is where they feed in queries and documents into the into the system. So this is an encoder decoder setup. In fact, they use, I believe, a T five setup, if I'm not mistaken. So it's a sequence to sequence task, they feed in the queries and the documents, and they always output the document ID. So for if they feed a document, they just output the ID of the document they fed in. And if they feed a query, they output the ID of the document that the query would hit. So this is if you have supervised data, you can train the system already for giving queries to output the correct document. However, the method also works in what they call zero shot, which is if you do not have any queries, you simply input documents into the system, and then you train it to output the ID of those documents. And you hope that because the models were pre trained on language modeling and on various other tasks, you hope that through that, if you then enter a query that kind of describes the same thing as the documents that the system would still output the best document ID. I mean, after all, it's constrained to output document IDs in most cases. And therefore, it needs to give you something so it might as well give you the thing that is related the most. So that's the reasoning behind it. I've talked a lot about the different parts now of the system, the write up is actually pretty good, I can recommend reading this paper from top to bottom, because it goes in a very structured form into what they investigate, they investigate a lot of engineering choices, which I really appreciate in the system, because there are a lot of ways to do this. And not one or the other is not necessarily correct. So they say we explore a number of variations of the DSI architecture, they explore how do we represent documents as such, the naive approach they say is just to index the full document. So just input the text as such, like you can see right here, just input the text into the encoder output the document ID, that's it. But maybe that's not the best thing to do. Maybe you can throw away stop words, maybe you can do bag of words representation, maybe something is better than just inputting the first L tokens of the document. Turns out it's not, but you know, it's a good good thing to investigate. The end result is, then how do we represent document IDs? The data sets, they usually just have like some unique identifier per document. In this case, it's like doc one, three, seven. And here it's doc four, five, six. If we do this as a sequence to sequence tasks, maybe we can do something smarter. Maybe we can give the document IDs some sort of hierarchical notion, they investigate that too. And lastly, they investigate how should we index stuff. So how how should exactly should the indexing step this training go? They also do a lot of ablations on sort of the effect of sizes, the effect of model size and corpus size. And we're going to look into that as well. So the method is called, as I said, differentiable search index, the goal is to fully parameterize traditionally multi stage retrieval, then rank pipelines within a single neural model. And that encompasses two operations. First is indexing. And then the second one is retrieval. In the DSI, we've already discussed this indexing is sequence to sequence approach that takes a document that takes document tokens as input and generates identifiers as output, that is indexing its training on the document collection to output their identifiers, and optionally, optionally, fine tuning with labeled query sets labeled query doc ID pairs. The retrieval is then achieved by simply autoregressive generation, I input something and I see what document ID comes out in the sequence to sequence model. So it couldn't get easier than that. Let's look a different a little bit into the engineering choices they consider. First, the indexing method. The first indexing method is what they call inputs to target. And that is probably what I've described so far, which is the sequence to sequence task of document tokens maps to document ID. So they input the tokens of the document, and they output the document ID. That is the simplest method, the straightforward method from what we've heard so far. And as far as I've read in the paper, as I understand it, this is also what works the best. However, they proclaim that in this way, the only ever output is the document ID, there is no sort of language learning or anything like this, you fully rely on the pre training for language understanding. That is what they claim here is a potential weakness. And other methods are, you know, targeted at are targeted at in sort of leveraging or making that weakness go away. They have this targets to inputs method, which they say, we could also at training time, adding what they call indexing time, input a document ID and then have the model decode the tokens of the document. Now, this might seem a bit weird, because it doesn't train the model to do that. It doesn't train the model to produce document IDs from tokens. But the idea is that you could, for example, then fine tune on query, document ID pairs, and that by by training with the with this objective, you teach the model something about the document IDs and which tokens which document tokens are in the in the IDs, because the model has to learn to produce the document tokens. And therefore, it might make some associations or something. I'm not exactly sure what the thing is behind like what the reasoning is behind this. But, you know, it's good to try. It doesn't work turns out. There's also bi directional, which both are done. So during training, there is like a multitask setup where sometimes you do the doc ID to tokens, and sometimes you do the tokens to doc ID. Also in their experiment, the bi directional method doesn't improve much over just the plain method. And the last one is span corruption, where you essentially input, I think the the tokens, tokens, and you append the doc ID. And then you consider this entire thing as like one piece of text that you want to predict. And you have this span corruption objective, which means that you can mark out any random spans in here between, which also means that sometimes you mask out the document ID or maybe part of the document ID. And that kind of forces the model to learn it's a bit like births masked language modeling, if I understand this correctly. However, also, this doesn't seem to work super well for them, even though it has actually worked well in other tasks. So in other papers that have done things in in this sort of sequence to sequence space. Okay, so now we have the indexing method of the table. The document representation strategies are next. The first one is direct indexing, you say we take the first L tokens. Again, this seems to work the best. Just take the first L tokens of the document. Interestingly, during the experiments, L bigger isn't necessarily better for L, which is also might speak to a little bit of the quality and nature of the data set itself, but also tells us again, something about maybe this works in a way that works in particular because we're dealing with sizes and data set sizes and lengths of documents that are actually possible to absorb into weights. And it is interesting to see how as the data goes up, this becomes harder and harder, I would question, does it become like linearly harder to put this into a set of weights? Does it become exponentially harder? If there's more data, not sure it would be interesting to find out. The other methods are, for example, set indexing that which de duplicates repeated terms and remove stop words doesn't seem to help much. And, you know, naturally, one might think that, you know, if I remove stop words in my document representation, that gives me a cleaner signal. On the other hand, these models are pre trained on actual language, not on cleaned up language without stop words, they're pre trained on actual language. And therefore they I think they have a strong bias towards, you know, kind of correct grammar and so on. And might work with that data a lot better. I think that might be largely behind why the direct indexing method works better over the set indexing. And then there's the in what they call inverted index, which is a bit in the spirit of how search engines classically do this. They say we randomly sub sample a single contiguous chunk of K tokens from the document. So they're not only limited to the first L tokens, but they always kind of take a random sub string of the document that is of that length. Now, technically, this should work better than the direct indexing. I like the the inverted index in their experiment performs worse than the direct indexing. And I just don't believe it. Like, like, it doesn't, it does not make sense, right? Something's going on either. The data set is such that for some reason, I can find a lot of the answers that I'm looking for in the first in the beginning of the documents that are indexed, but this is purely a property of the data set. Or it is really like the introduction of a tiny bit of noise into this, namely, that for the same document ID, I see different substrings, I see different tokens that that already kicks the method out of its out of its comfort zone. That seems to be like the in first instance, it's kind of a bummer that this is the data set, but we'll have to take it in the second instance, it's a bit more worrisome. If that were the case, like if that fact would be already detrimental, where it actually should be beneficial. Or, yeah, maybe I'm misunderstanding something, but it seems to me that the this last method should be superior to the first one. So the last thing they or the next thing they investigate is how do we represent, by the way, I'm already I'm already telling you about the experimental results there, they'll be coming up in the next section. But I think it's, it's easier to mention them already here than to keep everything in your head, and then go to the experimental results. But we will go into it in just a bit. They investigate how should we represent the doc IDs. Again, the simplest thing you can do is to have these unstructured atomic identifiers, which essentially means that every document gets an unique identifier. And then in a sequence to sequence model, right, I have my sequence here. This is an in goes into my encoder, and then it goes into a decoder. And the decoder produces a sequence. Now, every one of those tokens is in a list in a vocabulary, the vocabulary has a certain amount of entries, if I tokenize correctly, I have no out of vocabulary words. And this has a some kind of a fixed size like a vocabulary size. And the decoder, it can have the same vocabulary or a different vocabulary. In this case, I think it's the same. But what they do in this first method is they simply extend the vocabulary for the decoder. And the extra tokens here, every single token is represents one document ID. This obviously only works if you know all the documents ahead of time that you're going to index, but in their case, they do. So they randomly initialize those embeddings and during indexing, they train the embeddings for those. And that essentially means it's a multi class classification problem. At the end of the day, every sequence prediction task is, but we're not going to predict multiple tokens, we're going to predict exactly one token. And that token comes exactly from this vocabulary. And that means this this is not a sequence to sequence task, this is just a multi class classification task. Now this has advantages being multi class classification, it means there's one prediction, there's no auto regressivity or anything like this. It's essentially a classic encoder only problem. Though this is the easy part, the hard part is of course, you don't you don't leverage anything, you introduce a lot of new classes, a lot of new embeddings. And they claim in the experiments that these things are quite brittle, even though in the zero shot case, apparently they work out super well. But we'll have some comments on that too. The next thing is not evenly structured string identifiers. They so they say, again, like here, every document will have an arbitrary unique identifier, which is just kind of an integer. However, they just say, well, we'll just put the integer as a tokenizable string. So if the integers if the integers like one, one to five, then the model needs to predict the tokens like the strings one, one, two, and five, or maybe it's tokenized differently, but it will actually have to produce this thing as a string, not as a output into an output classification bucket, but it will have to output the string. So this is now truly a sequence to sequence task, right. And the last thing they consider is these semantically structured identifiers. And they it's where they think, well, can't we do something better for the document IDs? Like can't we imbue them with some meaning? And they come up with the following procedure. So they have two, they have two principles they want to follow. They say the doc ID should capture some information about the semantics of its associated document. And second, the doc ID should be structured in a way that search space is effectively reduced after each decoding step. This results in identifiers where semantically similar documents share identifier prefixes. So essentially, they want the documents to have multiple like the ID, the IDs could be 255, which essentially means it's like a path, right? It's like a folder path. So this is group super group two, and then group five inside of super group two, and then document five inside of that. And the assumption is that all the documents that are in the same like group two slash five, they share some stuff such that the decoder, if it's not sure which exact document it is, but it can already say, well, in super group two, I find all the things that talk about, I don't know, household items. And then in two slash five, there are all the things that talk about electric appliances in the household. And then inside of that, there might be some documents, but the model could consider step by step, the model would first consider outputting sort of the super group and then condition on that in order to output the group and then condition on that in order to output the next level. So that's what they do. They do a hierarchical clustering approach, which means that they take another model. So they take some sort of a, I think it's a BERT model. A BERT, I think, I'm not sure where they mention it. But they take a BERT model, they put all of the documents through the BERT model, they train and embed, I don't know if they actively train it or if they take a pre-trained one. In any case, they have some way of embedding documents. So they embed those documents, then they use k-means clustering to divide them into clusters. If the clusters are still too large, they recursively subdivide them into clusters. And here you see exactly, so this here is document 233, because it's in super group two, it's in subgroup three, so that's 233. And then it's the third document inside of that. So that's 233. And presumably the two and the three prefixes, they are kind of like the path into the hierarchy and make it easier for the model to decode. Now this seems like a seems like a cool idea, honestly, because it kind of makes sense. There are however, two conflicting things. One is the fact that there is semantic meaning in, you know, in 255 or 233. In that case, right, there's semantic meaning in these things, and not just a random identifier. The other one is that it is in order. So the top hierarchy is first, then the second, then the third, which might interplay with the autoregressive way that we train these things. So in order to separate the two things, one would need to make an experiment where you just flip it around, right, you decode while you decode, you do you decode from the back, you decode like 332. And then you essentially still retain the semantic information of the identifier, but you drop away the autoregressivity. So the model essentially could not condition on the supergroup while decoding the lower layers. So you could tease that apart a little bit. They didn't do that. But in any case, this would, I guess, be an idea of doing further ablation and understanding into how this model works. It is interesting. They Yeah, that's that's it, essentially. Okay. Then how do they train? They say we try two strategies. One is to first train the indexing step. So first feed the documents and output their IDs, followed by a fine tuning stage, where you feed queries and map them to their IDs. Or the second strategy is to train them together in a multitask setup. That's exactly what we saw on the diagram, you feed documents and queries for documents, the output their document ID for queries, you output the corresponding document ID, and you have some ratio of how many indexing samples and how many query samples that go in. Turns out that second method is better, which I don't know if I would have guessed that. But yeah, it kind of makes sense because it's cleaner. And you can you can essentially scale and distribute there is no way that you can do that. So you can just do it in a simple way. There's no ordering effect. There's no catastrophic forgetting or anything like this. And yeah, so that makes sense. So that's what they do. All right, we'll get into the experiments. Now, the data set is natural questions. This is a question answering data set, and it can be used for retrieval, because the data set essentially is a question, a passage, which is usually called the context and an answer. This is one data point. Now, the idea is that you look at the context and the question and you find the answer inside of it. However, you can make you can make a retrieval data set out of this by forgetting about the answer and by severing the connection between the context and the query, and considering the answer. And essentially, the task is now if you if I have a given query, a given question, which context is the correct one to go with that question. So you can make a retrieval data set, which is usually quite hard because the data set is made with the fact in mind that you will get the same answer as you would get if you were to look at the context, right? So it is not necessarily the same as a user typing something into Google, where they need to look for a for a document. The question is a question about the document if you already have the document. So it is a little bit different, not a direct retrieval data set. Also, note that it's kind of like 300 there's 300 K data points, they make subset of that so they make a 10 K, a 100 K, 10 K data set, 100 K data set, and a 300 K data set. So a small, medium and large, although even the large one right is not large, you can because in a search task, 300,000 documents, it seems a lot. But if you build search applications, that is not that is not a lot of documents, right? A lot of document collections have millions of documents and more that you need to retrieve from. But it is good to observe scaling properties right here. But just keep in mind that their largest data set is still not super duper large. The other thing you can see they have train pairs and validation pairs. And that kind of Yeah, so the all of these things, they have a special notion right here, which I'm not exactly sure I have to be honest how this is exactly done. So the training pairs, I have the queries and the context both right. And for the validation pairs, I also have queries and context. Now usually I train a question answering system, I train on these things right with the answers, and then I input these things over here at inference time. However, if I train a search index, I certainly need to index at least the contexts of the validation pairs. And I simply prohibit myself from ever seeing the queries. So what I think they do, what I think they do is that I think they take these together, they this these are all the contexts, all the documents, and they take the queries from the training set. And that makes sort of the the quote unquote training set, right? This, this here would be indexing. And this here would be fine tuning. And then they evaluate this here would be eval. But this is a hypothesis of mine, I'm not exactly sure that that's what they do. Because certainly they can't just not index the data that they're going to retrieve from right. But I hope they don't actually fine tune on the queries that are in the validation set. But again, maybe they also first do this. And then as a last step, they then index the validation set, I'm not sure just honestly, and I couldn't read from the paper, maybe I've overlooked something. But it would be a good question to the authors how this exactly is done. Training regimen seems pretty decent. So this it's Google research. So they have the big chips. Yeah, t five isn't exactly a small model, right? Especially the larger ones. So here are the results. And they are all over the place, which makes me a little bit skeptical. First, you can see in general, the larger models for the differentiable search index generally outperform the smaller models by a lot, right? You can see here, for example, these are large models, these are small models on the same task, these are hits at one and hits at 10, which means if the correct answer is in the top one or the top 10, respectively, for all of the DSI models, that's the case. By the way, when it says t five here, that is a dual encoder baseline. And above here, you can see the BM 25 baseline. Now, also, I would like to I would like to draw your attention to the fact that BM 25 on the small data set, it gets like a performance of 12.4. On the large data set, it gets like 11.6, which, you know, is reasonably kind of goes down a bit if the data set is larger, because it can confuse the documents a bit more, but in general, it's constant. But then there's like a big jump in this 100k data set, like what's up? What's up? What's up with that? This seems to this seems to be weird. So you can't really see that in the dual encoder setup, there is a jump here, but that remains. Then if you if you look at if you look at the small models here, it goes up and it goes down again. Yeah, that's the same trend. But then here, if you can see, it kind of goes down in performance. And then it goes up. No, it goes it kind of remains down. All I'm saying is this is not okay, this might be to be expected. This might be expected because going down in performance is what I would expect if it goes if the data set becomes larger. Okay. But there are some inconsistencies among here. Yeah, all the weirder that here actually goes up. And as you can see the highlighted bits right here, for example, this thing, the methods that work, they seem to be all over the place. Sometimes this naive string doc ID is the best. Sometimes this semantic string doc ID is the best. The clear trend is that pretty much everywhere the larger models are better, which I think is reasonable to say because they're going to have more capacity of adopting the data into their weights. And in other trends, the larger the data set gets, the worse the models become. Like look at this, it goes down to be expected, it goes up again, what's up? So this data set is just is cursed. So we won't look at it. So let's just compare the very left and the very right things. You can also you can also see that there isn't a big improvement over BM 25, which is surprising, right? That even the dual encoders improve over BM 25. But this differentiable search index, especially if it gets large improves by quite a bit. Now, I suspect again, that that is kind of the nature of the data set right here. But it might as well be that the all the embedding techniques are very good. But yeah, lastly, what I want to point out, oh, yeah, the improvement over the dual encoders of the differentiable search index. So over this baseline right here, this gets smaller and smaller as the data set grows, right, which we discussed at the beginning and which I think is a little bit of a bad sign for these types of techniques in that, obviously, as I have more data, I cannot really save it into my weights as easily. And the dual encoders, they are not like the embedding space, high dimensional embedding space is kind of infinite, right? So I can save a lot of stuff there, no matter how much data I have. It'd be interesting though, because there are techniques in which you can, like if I have a matrix, and I want to store stuff in that matrix, as long as that stuff as long as I build like low rank matrices that I add to it, or in vector terms, if I build like vectors that are largely orthogonal to one another, I can, you know, state save a lot of stuff in a single matrix by just adding to it, or to a vector space or to a set of vectors. And maybe, maybe, you know, with a bit of trickery in how the weights are updated exactly for the different documents, one could improve this quite a bit. This here is zero shot setting, which means this models, they never seek any queries, they never learn to map queries to document IDs, they simply learn to map documents to doc IDs, which is an additional difficulty. Again, you can see that the weirdness of BM 25, right, that's exactly the same, right, BM 25 is going to perform the same because BM 25 is always zero shot, it never sees it never sees labeled queries. You can you just can't I guess you can, you can also run it through indexing. But yeah, interestingly, the dual encoder in in a zero shot fashion just sucks, it really sucks. The sentence t five, which is explicitly made for like sentence sentence similarity, it is apparently okay, it apparently outperforms BM 25. Also, I have trouble believing that, but, you know, if they say so. But then these DSI, they really shine in this especially here, this atomic doc ID method. For some reason, it really is is really good. As you can see, it outperforms the semantic string doc ID, which was kind of the best one before or one of the best one. Also, this naive string doc ID was really good before it outperforms that in a zero shot setting. So the results are kind of all over the place. And that is what worries me a little bit in that seems to be quite noisy. They themselves admit or report that training with these atomic doc IDs seems to perform well in the zero shot setting, but it's also quite unstable. So yeah, it's a it's a cool method, cool paper. And it shows some really interesting results. But it also seems that there's quite a bit of noise. And probably we haven't exactly figured out many of those things yet, which is a good thing if you're in research. Yeah, so they find a bunch of things like in general, they say structured semantic identifiers are helpful and improve over unstructured ones. However, we also note that unstructured atomic identifiers perform the best by a wide margin on the zero shot retrieval setup. Who knows why? We can I guess we can hypothesize the other methods I've already discussed a little bit, especially model size, it seems to be really important, as you can see, for dual encoders, that doesn't pay that much of a that doesn't make super duper difference. It makes much more difference for the differentiable search index. Whereas if you talk about data set size, a higher data set size seems to be much more detrimental to the differentiable search index than it is to a dual encoder. Interestingly, also, the length of the tokens you index per document seems to be better if it's kind of shorter, which is interesting. So if you index the same documents for longer for more tokens, that seems to hurt performance. And really, if you go much, much longer. And lastly, here, they investigate how much indexing versus retrieval they have to feed in during the multitask training. If they train index and labeled query pairs at the same time, turns out that's also fairly noisy, but you can't go too high. One seems to be fine, right? So you can get an improvement if you have more indexing, but one seems to be fine, which is already relieving, I think, you could just mix them together and you'd be fine. Yeah, I wanted to say one, one more thing. Yes. So in their conclusion, they talk about document identifiers. And they say it would be interesting to explore alternative strategies for representing documents and doc IDs, including end to end strategies for learning semantic identifiers. That's what they say, because they're kind of unsatisfied with the way they represent the document IDs, because the height of their method is this hierarchical clustering, which is also uses a separate encoder and so on. However, I'm thinking myself, you know, if you want this to be learned, like end to end and so on, isn't that like, isn't that exactly like regressing to cross encoder setup and dense retrieval setup? Isn't that essentially what you're doing if you're learning these things end to end? I don't know exactly how then that's going to be different in principle. And this is my a little bit of my worry about this paper as well that they didn't compare at all to any cross encoder setup to any any kind of re ranking setup that are very prevalent in neural search these days, any dense retriever setup, maybe dense retriever is buying code, I'm not even sure. But I feel these are some some baselines that are missing right here, along with the smaller size of the data set. But all in all, pretty cool. Again, I don't think this is necessarily going to be such a use in search in itself like search through document collections, but much more could be very useful as a part in, for example, a reinforcement learning agent who has to store stuff during the episode and then retrieve it later in a very differentiable manner in an addressable manner. It would also be interesting to see, yeah, whether whether outputting document IDs is better than outputting the information that I want directly, right, because you could also think of that. You could also say, you know, here is a query, just output the document itself or the part of the document that matches instead of outputting the document ID. You know, how does that perform, it, it will be equally interesting to see that. So lots of things to research, I really like this paper because it does something different. It does something weird. And it puts in the engineering effort to figure out what makes it work and what doesn't. And yeah, that's it. Let me know what you think in the comments. I'll see you around. Bye bye.
 This is a comprehensive paper review of the paper transformer memory as a differentiable search
 index. This paper is pretty crazy. It takes an entire data set and puts it into the weights of
 a transformer. Essentially, it trains a search engine not to search through documents, but just
 to give you the index of the document that matches your query, just like that. Boom. So this video is
 a comprehensive review of the paper. I'll explain to you what's in the paper, what it's about. And
 by the end of the video, you should have a good idea of the paper itself. The next video, which
 I'm going to release tomorrow will be an interview with the authors will dive right into the content
 and any criticisms and questions that I raised during the review. As always, let me know what
 you think in the comments. Now let's get into the video. See you around. Does your company have a
 lot of people labeling data? Why would you leave such an important task to close source systems or
 self implemented things? Training data is your most valuable asset and human labels are really
 expensive. Today's sponsor is diffgram, which is an open source platform centered around training
 data. They handle everything to do with training data, especially collecting labeling, serving and
 more. And it is open source so you can self host all you want. But there's one cool thing if you
 let them host it for you. And that is unlimited pricing, no per label annotation, no expensive
 servers to run, you pay once you get as much as you want. So thanks again to diffgram for sponsoring
 today's video, check them out using the link in the description to let them know that I sent you.
 All right, let's get into the video. Hello there, today we're looking at transformer memory as a
 differentiable search index by researchers of Google research. This paper on high level takes
 a search problem where you have to index documents and retrieve them. And it puts all of the corpus
 essentially into the weights of a transformer. So it takes the corpus and trains the transformer.
 And then at the end, they can just give a query to the transformer and the transformer will output
 the ID of the document that matches and it turns out for some data sets that they have for some
 settings and with some clever training and representation of the documents that can actually
 work, which is really crazy. This kind of speaks to multiple things such as obviously our ability
 to overfit on stuff, but there is some generalization here as we'll see. On the other hand also the
 kind of inner workings of these transformers. And lastly, what's pretty cool is that this is
 completely as it says differentiable, it's a differentiable search index, which means that
 this can be part of larger neural network architectures because it is fully differentiable
 and it can be trained essentially end to end at once. And that means we can potentially employ
 reinforcement learning agents with kind of retrieval abilities and much more things.
 So we'll dive into the paper, we'll see what it's about. The idea, as I said, is pretty,
 pretty simple. If you like content like this, then as always leave a like and tell me what you think
 in the comments. That's always super helpful. So as I said, they take a search problem and the
 search problem is essentially I have a corpus, like I have a big database of documents, right?
 Here is a document, here is a document and I want to build an index and an index is some kind of
 data structure, some kind of thing. And at the index, I can throw a query and the index will
 return to me an ID, a document ID that specifies which document matches my query. Usually this is
 done via inverted indices. So I want to tokenize my documents, split them into little tokens,
 which are usually words or sub words, and I want to stem them and lemmatize them and whatnot. Then
 I build a reverse index. So for every word like in the word in I remember which documents it appears
 in like document three, document five, document 11, and so on. And then once the query rolls in,
 I simply tokenize it as well. I go look into my inverted index and I look up all the documents.
 And then there's also a ranking step, which means I have to now determine which of these documents
 is the most relevant. And that is usually done via techniques like TF IDF features. There is a
 famous technique called BM 25, which is also a baseline in this paper. So this is the classic
 search kind of way way of doing search. If you use any search engine at all, this is being done
 in the background. For the most part, newer search engines are catching on, there's neural search and
 so on. But BM 25 is still one of the most performant things that text search has available, and also
 other types of search. However, there is a new push in sort of neural search. And in neural search,
 you're trying to take your data set. And for each document, you try to map it to some sort of a
 vector in vector space. And then once the query comes in, you also map the query to a vector.
 And for example, you compare inner products, whichever inner product is largest, that's the
 document that's relevant. This is just one way. This is what we usually call a by encoder method
 where the documents in the queries are mapped, both mapped individually. So there would be an
 encoder here, and there would be an encoder here, they all would output one vector, and then the
 vectors are compared, this could be the same encoder or different encoders for documents in query.
 This is just one method, there's various methods such as cross encoders, rerankers, dense retrievers,
 you name it. However, this method here is even more is different. So what we want to do is we
 want to take the corpus as such, and map that somehow into a neural network. And we're going
 to talk about this somehow. But we're going to train a neural network, essentially, how do we
 represent this, let's say represented with its layers, such that when later I feed a query to
 the neural network, as I already said, the ID of the document is the output of the neural network.
 So it doesn't output a vector that I then go and compare, it doesn't, I don't have to go and feed
 in query document pairs. And then I get out a score of how well they fit together, which would
 I would do in a cross encoder. No, the transformer, in this case, the neural network directly gives me
 the ID of the document, which without seeing the data at inference time. So during training,
 all of the data is essentially has to be mapped somehow into the weights of the neural networks,
 right? So somewhere in these weights, that information is stored of what the documents
 are. So the entire corpus is in those weights. And once I enter a query, the correct document ID can
 only be output, obviously, if the transformer has somehow learned what is in those documents.
 So that's the setup is pretty simple setup. Once you kind of see what's going on. It's,
 it's like a meme, right? Instead of we've been trying to neuralize search, and we've still done
 this two step process where we train these encoders, but then the actual search is still done
 using, for example, a nearest neighbor algorithm like here. But, you know, this is just the idea of,
 well, why don't I just ask the neural network to output the result, right, the resulting doc ID,
 why don't I just do that? And it turns out that can work surprisingly well. So you can do large,
 a couple of things here, but that's essentially it. They say right here in the introduction,
 they use a sequence to sequence learning system to directly map a query to a relevant document ID.
 They have different corpuses where they train it on on the smallest corpus,
 this method improves the hits at one, which means that whether the top hit is the correct one,
 more than 20 points from 12.4% for a dual encoder. So the baseline here is 12.4% for a dual encoder.
 So the baseline here is a dual encoder, what I shown whenever the there are two encoders,
 and they each output an embedding to 33.9%. That's a giant gain, right? That's like a 2.5x
 improvement. However, on a corpus that's 30 times larger performance is improved by nearly
 seven points, which is less. It's also respectable that performance if it is improved at all. However,
 I want you to notice and that's already kind of the first indication, a little bit of obviously
 what's going on here. On smaller data sets, this method does super duper well on larger data sets,
 the method doesn't do that much better than a sort of cross encoder type setup, sorry,
 a bi encoder type setup or a dual encoder type setup, which is understandable, right? Because
 the smaller the data, the easier it is to absorb it all into your weights. If the data gets larger,
 that obviously gets harder and harder. There's more data to go around, which means there's more
 error, room for error for confusion, and so on. And a classic search engine or a dual encoder is
 going to have a easier time in that case. But still, it's a cool paper. It's just that it
 kind of gets worse with the data set scale. It does get better with the model scaled though.
 The really exciting thing is something that I've already mentioned and they mentioned this here,
 all aspects, sorry about that. They say all aspects of retrieval are mapped into well
 understood machine learning tasks. So for example, indexing, which is building the reverted index,
 or even if you have the dual encoder, you need to build the nearest neighbor index, which is a hard
 task in high dimensions is now a special case of model training. So it's just training and
 incrementally updating an index becomes just a special case of model updating. So all the
 tasks are just tasks that we already understand from neural network training. So here is a
 comparison of the dual encoder method, which is the, let's say old classic neural search method,
 not the BM 25 retrieval, but the neural search method, and this DSI, the differentiable search
 index. So in the dual encoder method, what we do is we train this encoder. And in this case,
 they train one encoder for both the queries, as well as the documents. And what we try to do is
 we are going to try to use some form of contrastive loss. If we actually have query document pairs,
 what we can do is we can try to get the documents, the query and the document that go with each other
 to be close together, while making the documents that are unrelated to each other be far apart.
 So this is some sort of contrastive loss, obviously, at inference time, what we're going to do is we
 have a query, we put it through the encoder, we get its embedding, and we do a maximum inner product
 search through our entire vector space of our of our indexed data set, and we get a ranked list.
 So it's kind of this two step approach with building these indices in between, and with the
 training objective, that is not directly what we want. It is a proxy objective, because of the
 because of the algorithm later needs it the inner product search, but it is not actually what we
 want. So let's just train what we want. In the DSI, in the differentiable search index, I simply
 feed my query along with I simply feed my query essentially to in some form to the system. And the
 system outputs directly which document is relevant for the query. So the way they train it, and this
 is one way they train it is where they feed in queries and documents into the into the system.
 So this is an encoder decoder setup. In fact, they use, I believe, a T five setup, if I'm not mistaken.
 So it's a sequence to sequence task, they feed in the queries and the documents, and they always
 output the document ID. So for if they feed a document, they just output the ID of the document
 they fed in. And if they feed a query, they output the ID of the document that the query would hit.
 So this is if you have supervised data, you can train the system already for giving queries to
 output the correct document. However, the method also works in what they call zero shot, which is
 if you do not have any queries, you simply input documents into the system, and then you train it
 to output the ID of those documents. And you hope that because the models were pre trained on language
 modeling and on various other tasks, you hope that through that, if you then enter a query that kind
 of describes the same thing as the documents that the system would still output the best document ID.
 I mean, after all, it's constrained to output document IDs in most cases. And therefore,
 it needs to give you something so it might as well give you the thing that is related the most.
 So that's the reasoning behind it. I've talked a lot about the different parts now of the system,
 the write up is actually pretty good, I can recommend reading this paper from top to bottom,
 because it goes in a very structured form into what they investigate, they investigate a lot of
 engineering choices, which I really appreciate in the system, because there are a lot of ways to do
 this. And not one or the other is not necessarily correct. So they say we explore a number of
 variations of the DSI architecture, they explore how do we represent documents as such, the naive
 approach they say is just to index the full document. So just input the text as such, like
 you can see right here, just input the text into the encoder output the document ID, that's it. But
 maybe that's not the best thing to do. Maybe you can throw away stop words, maybe you can do bag of
 words representation, maybe something is better than just inputting the first L tokens of the
 document. Turns out it's not, but you know, it's a good good thing to investigate. The end result
 is, then how do we represent document IDs? The data sets, they usually just have like some unique
 identifier per document. In this case, it's like doc one, three, seven. And here it's doc four,
 five, six. If we do this as a sequence to sequence tasks, maybe we can do something smarter. Maybe we
 can give the document IDs some sort of hierarchical notion, they investigate that too. And lastly,
 they investigate how should we index stuff. So how how should exactly should the indexing step this
 training go? They also do a lot of ablations on sort of the effect of sizes, the effect of model
 size and corpus size. And we're going to look into that as well. So the method is called, as I said,
 differentiable search index, the goal is to fully parameterize traditionally multi stage retrieval,
 then rank pipelines within a single neural model. And that encompasses two operations. First is
 indexing. And then the second one is retrieval. In the DSI, we've already discussed this indexing
 is sequence to sequence approach that takes a document that takes document tokens as input and
 generates identifiers as output, that is indexing its training on the document collection
 to output their identifiers, and optionally, optionally, fine tuning with labeled query
 sets labeled query doc ID pairs. The retrieval is then achieved by simply autoregressive generation,
 I input something and I see what document ID comes out in the sequence to sequence model.
 So it couldn't get easier than that. Let's look a different a little bit into the engineering
 choices they consider. First, the indexing method. The first indexing method is what they call inputs
 to target. And that is probably what I've described so far, which is the sequence to sequence task of
 document tokens maps to document ID. So they input the tokens of the document, and they output the
 document ID. That is the simplest method, the straightforward method from what we've heard so
 far. And as far as I've read in the paper, as I understand it, this is also what works the best.
 However, they proclaim that in this way, the only ever output is the document ID, there is no sort
 of language learning or anything like this, you fully rely on the pre training for language
 understanding. That is what they claim here is a potential weakness. And other methods are, you
 know, targeted at are targeted at in sort of leveraging or making that weakness go away. They
 have this targets to inputs method, which they say, we could also at training time, adding what
 they call indexing time, input a document ID and then have the model decode the tokens of the
 document. Now, this might seem a bit weird, because it doesn't train the model to do that.
 It doesn't train the model to produce document IDs from tokens. But the idea is that you could,
 for example, then fine tune on query, document ID pairs, and that by by training with the with this
 objective, you teach the model something about the document IDs and which tokens which document
 tokens are in the in the IDs, because the model has to learn to produce the document tokens. And
 therefore, it might make some associations or something. I'm not exactly sure what the
 thing is behind like what the reasoning is behind this. But, you know,
 it's good to try. It doesn't work turns out. There's also bi directional, which both are done.
 So during training, there is like a multitask setup where sometimes you do the doc ID to tokens,
 and sometimes you do the tokens to doc ID. Also in their experiment, the bi directional method
 doesn't improve much over just the plain method. And the last one is span corruption, where you
 essentially input, I think the the tokens, tokens, and you append the doc ID. And then you consider
 this entire thing as like one piece of text that you want to predict. And you have this span
 corruption objective, which means that you can mark out any random spans in here between,
 which also means that sometimes you mask out the document ID or maybe part of the document ID.
 And that kind of forces the model to learn it's a bit like births masked language modeling,
 if I understand this correctly. However, also, this doesn't seem to work super well for them,
 even though it has actually worked well in other tasks. So in other papers that have done
 things in in this sort of sequence to sequence space. Okay, so now we have the indexing method
 of the table. The document representation strategies are next. The first one is direct
 indexing, you say we take the first L tokens. Again, this seems to work the best. Just take
 the first L tokens of the document. Interestingly, during the experiments, L bigger isn't
 necessarily better for L, which is also might speak to a little bit of the quality and nature
 of the data set itself, but also tells us again, something about maybe this works in a way that
 works in particular because we're dealing with sizes and data set sizes and lengths of documents
 that are actually possible to absorb into weights. And it is interesting to see how as the data goes
 up, this becomes harder and harder, I would question, does it become like linearly harder
 to put this into a set of weights? Does it become exponentially harder? If there's more data,
 not sure it would be interesting to find out. The other methods are, for example, set indexing that
 which de duplicates repeated terms and remove stop words doesn't seem to help much. And,
 you know, naturally, one might think that, you know, if I remove stop words in my document
 representation, that gives me a cleaner signal. On the other hand, these models are pre trained
 on actual language, not on cleaned up language without stop words, they're pre trained on actual
 language. And therefore they I think they have a strong bias towards, you know, kind of correct
 grammar and so on. And might work with that data a lot better. I think that might be largely behind
 why the direct indexing method works better over the set indexing. And then there's the in what
 they call inverted index, which is a bit in the spirit of how search engines classically do this.
 They say we randomly sub sample a single contiguous chunk of K tokens from the document.
 So they're not only limited to the first L tokens, but they always kind of take a random sub string
 of the document that is of that length. Now, technically, this should work better than the
 direct indexing. I like the the inverted index in their experiment performs worse than the direct
 indexing. And I just don't believe it. Like, like, it doesn't, it does not make sense, right?
 Something's going on either. The data set is such that for some reason, I can find a lot of the
 answers that I'm looking for in the first in the beginning of the documents that are indexed, but
 this is purely a property of the data set. Or it is really like the introduction of a tiny bit of
 noise into this, namely, that for the same document ID, I see different substrings, I see different
 tokens that that already kicks the method out of its out of its comfort zone. That seems to be
 like the in first instance, it's kind of a bummer that this is the data set, but we'll have to take
 it in the second instance, it's a bit more worrisome. If that were the case, like if that fact
 would be already detrimental, where it actually should be beneficial. Or, yeah, maybe I'm
 misunderstanding something, but it seems to me that the this last method should be superior to
 the first one. So the last thing they or the next thing they investigate is how do we represent,
 by the way, I'm already I'm already telling you about the experimental results there, they'll be
 coming up in the next section. But I think it's, it's easier to mention them already here than to
 keep everything in your head, and then go to the experimental results. But we will go into it in
 just a bit. They investigate how should we represent the doc IDs. Again, the simplest thing you can do
 is to have these unstructured atomic identifiers, which essentially means that every document gets
 an unique identifier. And then in a sequence to sequence model, right, I have my sequence here.
 This is an in goes into my encoder, and then it goes into a decoder. And the decoder produces
 a sequence. Now, every one of those tokens is in a list in a vocabulary, the vocabulary has a certain
 amount of entries, if I tokenize correctly, I have no out of vocabulary words. And this has a some
 kind of a fixed size like a vocabulary size. And the decoder, it can have the same vocabulary or a
 different vocabulary. In this case, I think it's the same. But what they do in this first method is
 they simply extend the vocabulary for the decoder. And the extra tokens here, every single token is
 represents one document ID. This obviously only works if you know all the documents ahead of time
 that you're going to index, but in their case, they do. So they randomly initialize those embeddings
 and during indexing, they train the embeddings for those. And that essentially means it's a
 multi class classification problem. At the end of the day, every sequence prediction task is,
 but we're not going to predict multiple tokens, we're going to predict exactly one token. And that
 token comes exactly from this vocabulary. And that means this this is not a sequence to sequence
 task, this is just a multi class classification task. Now this has advantages being multi class
 classification, it means there's one prediction, there's no auto regressivity or anything like
 this. It's essentially a classic encoder only problem. Though this is the easy part, the hard
 part is of course, you don't you don't leverage anything, you introduce a lot of new classes,
 a lot of new embeddings. And they claim in the experiments that these things are quite brittle,
 even though in the zero shot case, apparently they work out super well. But we'll have some
 comments on that too. The next thing is not evenly structured string identifiers. They so they say,
 again, like here, every document will have an arbitrary unique identifier, which is just kind
 of an integer. However, they just say, well, we'll just put the integer as a tokenizable string. So
 if the integers if the integers like one, one to five, then the model needs to predict the tokens
 like the strings one, one, two, and five, or maybe it's tokenized differently, but it will actually
 have to produce this thing as a string, not as a output into an output classification bucket,
 but it will have to output the string. So this is now truly a sequence to sequence task, right.
 And the last thing they consider is these semantically structured identifiers. And they
 it's where they think, well, can't we do something better for the document IDs? Like can't we imbue
 them with some meaning? And they come up with the following procedure. So they have two,
 they have two principles they want to follow. They say the doc ID should capture some information
 about the semantics of its associated document. And second, the doc ID should be structured in a
 way that search space is effectively reduced after each decoding step. This results in identifiers
 where semantically similar documents share identifier prefixes. So essentially, they want
 the documents to have multiple like the ID, the IDs could be 255, which essentially means it's
 like a path, right? It's like a folder path. So this is group super group two, and then group five
 inside of super group two, and then document five inside of that. And the assumption is that all
 the documents that are in the same like group two slash five, they share some stuff such that the
 decoder, if it's not sure which exact document it is, but it can already say, well, in super group
 two, I find all the things that talk about, I don't know, household items. And then in two slash five,
 there are all the things that talk about electric appliances in the household. And then inside of
 that, there might be some documents, but the model could consider step by step, the model would first
 consider outputting sort of the super group and then condition on that in order to output the group
 and then condition on that in order to output the next level. So that's what they do. They do a
 hierarchical clustering approach, which means that they take another model. So they take some sort of
 a, I think it's a BERT model. A BERT, I think, I'm not sure where they mention it. But they take a
 BERT model, they put all of the documents through the BERT model, they train and embed, I don't know
 if they actively train it or if they take a pre-trained one. In any case, they have some way
 of embedding documents. So they embed those documents, then they use k-means clustering to
 divide them into clusters. If the clusters are still too large, they recursively subdivide them into
 clusters. And here you see exactly, so this here is document 233, because it's in super group two,
 it's in subgroup three, so that's 233. And then it's the third document inside of that. So that's
 233. And presumably the two and the three prefixes, they are kind of like the path into the hierarchy
 and make it easier for the model to decode. Now this seems like a seems like a cool idea,
 honestly, because it kind of makes sense. There are however, two conflicting things. One is the fact
 that there is semantic meaning in, you know, in 255 or 233. In that case, right, there's semantic
 meaning in these things, and not just a random identifier. The other one is that it is in order.
 So the top hierarchy is first, then the second, then the third, which might interplay with the
 autoregressive way that we train these things. So in order to separate the two things, one would
 need to make an experiment where you just flip it around, right, you decode while you decode, you do
 you decode from the back, you decode like 332. And then you essentially still retain the
 semantic information of the identifier, but you drop away the autoregressivity. So the model
 essentially could not condition on the supergroup while decoding the lower layers. So you could
 tease that apart a little bit. They didn't do that. But in any case, this would, I guess, be an idea
 of doing further ablation and understanding into how this model works. It is interesting.
 They
 Yeah, that's that's it, essentially. Okay. Then how do they train? They say we try two strategies.
 One is to first train the indexing step. So first feed the documents and output their IDs,
 followed by a fine tuning stage, where you feed queries and map them to their IDs. Or the second
 strategy is to train them together in a multitask setup. That's exactly what we saw on the diagram,
 you feed documents and queries for documents, the output their document ID for queries, you
 output the corresponding document ID, and you have some ratio of how many indexing samples
 and how many query samples that go in. Turns out that second method is better, which I don't know
 if I would have guessed that. But yeah, it kind of makes sense because it's cleaner. And you can
 you can essentially scale and distribute there is no way that you can do that. So you can just
 do it in a simple way. There's no ordering effect. There's no catastrophic forgetting
 or anything like this. And yeah, so that makes sense. So that's what they do. All right,
 we'll get into the experiments. Now, the data set is natural questions. This is a question
 answering data set, and it can be used for retrieval, because the data set essentially
 is a question, a passage, which is usually called the context and an answer. This is one data point.
 Now, the idea is that you look at the context and the question and you find the answer inside of
 it. However, you can make you can make a retrieval data set out of this by forgetting about the
 answer and by severing the connection between the context and the query, and considering the
 answer. And essentially, the task is now if you if I have a given query, a given question, which
 context is the correct one to go with that question. So you can make a retrieval data set,
 which is usually quite hard because the data set is made with the fact in mind that you will get
 the same answer as you would get if you were to look at the context, right? So it is not necessarily
 the same as a user typing something into Google, where they need to look for a for a document.
 The question is a question about the document if you already have the document. So it is a little
 bit different, not a direct retrieval data set. Also, note that it's kind of like 300 there's 300
 K data points, they make subset of that so they make a 10 K, a 100 K, 10 K data set, 100 K data
 set, and a 300 K data set. So a small, medium and large, although even the large one right is not
 large, you can because in a search task, 300,000 documents, it seems a lot. But if you build search
 applications, that is not that is not a lot of documents, right? A lot of document collections
 have millions of documents and more that you need to retrieve from. But it is good to observe
 scaling properties right here. But just keep in mind that their largest data set is still not
 super duper large. The other thing you can see they have train pairs and validation pairs. And
 that kind of Yeah, so the all of these things, they have a special notion right here, which I'm
 not exactly sure I have to be honest how this is exactly done. So the training pairs, I have the
 queries and the context both right. And for the validation pairs, I also have queries and context.
 Now usually I train a question answering system, I train on these things right with the answers,
 and then I input these things over here at inference time. However, if I train a search
 index, I certainly need to index at least the contexts of the validation pairs. And I simply
 prohibit myself from ever seeing the queries. So what I think they do, what I think they do is that
 I think they take these together, they this these are all the contexts, all the documents,
 and they take the queries from the training set. And that makes sort of the the quote unquote
 training set, right? This, this here would be indexing. And this here would be fine tuning.
 And then they evaluate this here would be eval. But this is a hypothesis of mine, I'm not exactly
 sure that that's what they do. Because certainly they can't just not index the data that they're
 going to retrieve from right. But I hope they don't actually fine tune on the queries that are in the
 validation set. But again, maybe they also first do this. And then as a last step, they then index
 the validation set, I'm not sure just honestly, and I couldn't read from the paper, maybe I've
 overlooked something. But it would be a good question to the authors how this exactly is done.
 Training regimen seems pretty decent. So this it's Google research. So they have the big chips.
 Yeah, t five isn't exactly a small model, right? Especially the larger ones. So here are the results.
 And they are all over the place, which makes me a little bit skeptical. First, you can see in general,
 the larger models for the differentiable search index generally outperform the smaller models by
 a lot, right? You can see here, for example, these are large models, these are small models on the
 same task, these are hits at one and hits at 10, which means if the correct answer is in the top
 one or the top 10, respectively, for all of the DSI models, that's the case. By the way, when it
 says t five here, that is a dual encoder baseline. And above here, you can see the BM 25 baseline.
 Now, also, I would like to I would like to draw your attention to the fact that BM 25 on the small
 data set, it gets like a performance of 12.4. On the large data set, it gets like 11.6, which,
 you know, is reasonably kind of goes down a bit if the data set is larger, because it can confuse
 the documents a bit more, but in general, it's constant. But then there's like a big jump in this
 100k data set, like what's up? What's up? What's up with that? This seems to this seems to be
 weird. So you can't really see that in the dual encoder setup, there is a jump here, but that remains.
 Then if you if you look at if you look at the small models here, it goes up and it goes down
 again. Yeah, that's the same trend. But then here, if you can see, it kind of goes down in performance.
 And then it goes up. No, it goes it kind of remains down. All I'm saying is this is not okay,
 this might be to be expected. This might be expected because going down in performance is
 what I would expect if it goes if the data set becomes larger. Okay. But there are some inconsistencies
 among here. Yeah, all the weirder that here actually goes up. And as you can see the highlighted
 bits right here, for example, this thing, the methods that work, they seem to be all over the place.
 Sometimes this naive string doc ID is the best. Sometimes this semantic string doc ID is the best.
 The clear trend is that pretty much everywhere the larger models are better, which I think is reasonable
 to say because they're going to have more capacity of adopting the data into their weights. And in
 other trends, the larger the data set gets, the worse the models become. Like look at this,
 it goes down to be expected, it goes up again, what's up? So this data set is just is cursed.
 So we won't look at it. So let's just compare the very left and the very right things. You can also
 you can also see that there isn't a big improvement over BM 25, which is surprising, right?
 That even the dual encoders improve over BM 25. But this differentiable search index, especially
 if it gets large improves by quite a bit. Now, I suspect again, that that is kind of the nature
 of the data set right here. But it might as well be that the all the embedding techniques are
 very good. But yeah, lastly, what I want to point out, oh, yeah, the improvement over the dual
 encoders of the differentiable search index. So over this baseline right here, this gets smaller
 and smaller as the data set grows, right, which we discussed at the beginning and which I think is a
 little bit of a bad sign for these types of techniques in that, obviously, as I have more
 data, I cannot really save it into my weights as easily. And the dual encoders, they are not
 like the embedding space, high dimensional embedding space is kind of infinite, right? So
 I can save a lot of stuff there, no matter how much data I have. It'd be interesting though,
 because there are techniques in which you can, like if I have a matrix, and I want to store
 stuff in that matrix, as long as that stuff as long as I build like low rank matrices that I add to
 it, or in vector terms, if I build like vectors that are largely orthogonal to one another, I can,
 you know, state save a lot of stuff in a single matrix by just adding to it, or to a vector space
 or to a set of vectors. And maybe, maybe, you know, with a bit of trickery in how the weights are
 updated exactly for the different documents, one could improve this quite a bit. This here is zero
 shot setting, which means this models, they never seek any queries, they never learn to map queries
 to document IDs, they simply learn to map documents to doc IDs, which is an additional difficulty.
 Again, you can see that the weirdness of BM 25, right, that's exactly the same, right,
 BM 25 is going to perform the same because BM 25 is always zero shot, it never sees
 it never sees labeled queries. You can you just can't I guess you can, you can also run it through
 indexing. But yeah, interestingly, the dual encoder in in a zero shot fashion just sucks,
 it really sucks. The sentence t five, which is explicitly made for like sentence sentence similarity,
 it is apparently okay, it apparently outperforms BM 25. Also, I have trouble believing that, but,
 you know, if they say so. But then these DSI, they really shine in this especially here, this atomic
 doc ID method. For some reason, it really is is really good. As you can see, it outperforms
 the semantic string doc ID, which was kind of the best one before or one of the best one. Also,
 this naive string doc ID was really good before it outperforms that in a zero shot setting.
 So the results are kind of all over the place. And that is what worries me a little bit in that
 seems to be quite noisy. They themselves admit or report that training with these atomic doc IDs
 seems to perform well in the zero shot setting, but it's also quite unstable. So yeah, it's a
 it's a cool method, cool paper. And it shows some really interesting results. But it also seems that
 there's quite a bit of noise. And probably we haven't exactly figured out many of those things
 yet, which is a good thing if you're in research. Yeah, so they find a bunch of things like in
 general, they say structured semantic identifiers are helpful and improve over unstructured ones.
 However, we also note that unstructured atomic identifiers perform the best by a wide margin
 on the zero shot retrieval setup. Who knows why? We can I guess we can hypothesize the other
 methods I've already discussed a little bit, especially model size, it seems to be really
 important, as you can see, for dual encoders, that doesn't pay that much of a that doesn't make
 super duper difference. It makes much more difference for the differentiable search index.
 Whereas if you talk about data set size, a higher data set size seems to be much more detrimental
 to the differentiable search index than it is to a dual encoder. Interestingly, also, the length of
 the tokens you index per document seems to be better if it's kind of shorter, which is interesting.
 So if you index the same documents for longer for more tokens, that seems to hurt performance. And
 really, if you go much, much longer. And lastly, here, they investigate how much indexing versus
 retrieval they have to feed in during the multitask training. If they train index and labeled
 query pairs at the same time, turns out that's also fairly noisy, but you can't go too high.
 One seems to be fine, right? So you can get an improvement if you have more indexing,
 but one seems to be fine, which is already relieving, I think, you could just mix them
 together and you'd be fine. Yeah, I wanted to say one, one more thing. Yes. So in their conclusion,
 they talk about document identifiers. And they say it would be interesting to explore alternative
 strategies for representing documents and doc IDs, including end to end strategies for learning
 semantic identifiers. That's what they say, because they're kind of unsatisfied with the
 way they represent the document IDs, because the height of their method is this hierarchical
 clustering, which is also uses a separate encoder and so on. However, I'm thinking myself,
 you know, if you want this to be learned, like end to end and so on, isn't that like, isn't that
 exactly like regressing to cross encoder setup and dense retrieval setup? Isn't that essentially
 what you're doing if you're learning these things end to end? I don't know exactly how then that's
 going to be different in principle. And this is my a little bit of my worry about this paper as well
 that they didn't compare at all to any cross encoder setup to any any kind of re ranking setup
 that are very prevalent in neural search these days, any dense retriever setup,
 maybe dense retriever is buying code, I'm not even sure. But I feel these are some some baselines
 that are missing right here, along with the smaller size of the data set. But all in all,
 pretty cool. Again, I don't think this is necessarily going to be such a use in search in
 itself like search through document collections, but much more could be very useful as a part in,
 for example, a reinforcement learning agent who has to store stuff during the episode and then
 retrieve it later in a very differentiable manner in an addressable manner. It would also be
 interesting to see, yeah, whether whether outputting document IDs is better than outputting the
 information that I want directly, right, because you could also think of that. You could also say,
 you know, here is a query, just output the document itself or the part of the document that matches
 instead of outputting the document ID. You know, how does that perform, it, it will be equally
 interesting to see that. So lots of things to research, I really like this paper because it
 does something different. It does something weird. And it puts in the engineering effort to figure
 out what makes it work and what doesn't. And yeah, that's it. Let me know what you think in
 the comments. I'll see you around. Bye bye.