In-Context Retrieval with Siddharth Gollapudi - Weaviate Podcast #146!

Author: Weaviate vector database · Watch on YouTube

Video authored by the channel shown, not by IANIA; all rights belong to the creator. Transcript and summary generated automatically; the original-language content is preserved.

Original language: English

Original video: In-Context Retrieval with Siddharth Gollapudi - Weaviate Podcast #146! · Channel: Weaviate vector database · 56:06 · Watch on YouTube

Topics: AI models, RAG, search and information retrieval, Health and science

Summary

Siddharth Gollapudi explains in-context retrieval (ICR), where an LLM uses attention over documents in its context to identify relevant ones, rather than relying on a conventional nearest-neighbor index. He describes experiments with document codes, training methods, and architectural changes aimed at improving retrieval as context grows, including techniques to mitigate attention dilution. The conversation explores how ICR relates to reranking, sparse attention, vector databases, and long-context models, while emphasizing that retrieval over very large corpora remains a difficult generalization problem. They close by discussing the promise of models that can discover new connections and knowledge, and how progress might be evaluated through genuine discoveries.

Key points

ICR places documents in an LLM’s context and uses attention to find the most relevant document, rather than using an embedding-based nearest-neighbor index.
Unlike generative retrieval, ICR does not require documents to be stored in model parameters, which makes corpus updates less dependent on retraining and can support verbatim retrieval.
The researchers found that randomizing document codes helped reduce positional expectations and improve generalization at moderate corpus sizes.
On-policy training rollouts helped address cases where the model began generating a document code but became distracted or switched to another document.
Scaling from hundreds of documents to thousands and million-token contexts proved difficult; changing training objectives alone did not resolve the problem.
Attention dilution occurs as the softmax denominator grows, making attention scores increasingly similar; the paper explored length-aware temperature scaling and removing low-relevance tokens.
Sparse attention could improve both long-context efficiency and retrieval quality, but integrating retrieval mechanisms into end-to-end pretraining remains an open challenge.
Adding training data from multiple benchmark datasets helped somewhat, but the gains were limited; the model appeared to learn a broader relevance algorithm.
The speakers see potential in long-context models for uncovering connections across information, beyond conventional agentic search or benchmark performance.
A meaningful test of AI creativity may be whether models can discover genuinely new results, such as mathematical insights or scientific findings.

Chapters

Start	Chapter	Summary
00:05	Introduction and definition of in-context retrieval	The hosts introduce the topic, and Siddharth defines ICR as using an LLM’s attention over contextualized documents to retrieve relevant information.
02:10	ICR compared with generative retrieval	They distinguish ICR from parameter-based generative retrieval and discuss the latter’s inflexibility and potential difficulty returning exact documents.
03:43	Lost in the middle and long-context models	Siddharth reflects on positional bias in earlier language models and says their experiments encountered fewer lost-in-the-middle issues than expected.
06:20	Reranking and the challenge of scaling retrieval	The discussion contrasts listwise reranking with first-stage retrieval and focuses on the difficulty of generalizing beyond hundreds of documents.
11:27	Training models for retrieval	They discuss how underexplored model training and attention analysis are, then introduce the experiments and their hardware constraints.
13:31	Document codes and on-policy training	Siddharth describes how randomized codes improved performance and how rollout-based corrections helped the model stay consistent when generating codes.
21:14	Structured outputs and training data	They consider constrained generation, post-training limitations, and the value of using cleaned MS Marco data alongside other retrieval benchmarks.
27:29	Attention dilution and architectural changes	Siddharth explains how attention scores can wash out over long contexts and describes token filtering and length-aware temperature scaling.
34:16	Sparse attention and vector search	The speakers explore sparse attention as a way to reduce the tokens processed, while noting the difficulty of making retrieval fully trainable end to end.
38:23	Retrieval’s role in future language models	They connect retrieval to efficient context selection and discuss why sublinear attention could matter even when all documents fit in a context window.
40:59	Resources, rerankers, and latency	Siddharth discusses the value of more GPUs for pretraining and iteration, while the hosts consider the latency tradeoffs of LLM-based rerankers.
46:45	Long-context models and complex tasks	They discuss whether long-context models can handle multihop search and the possibility of discovering connections by processing broad collections of information.
51:26	AI research directions and discovery	Siddharth describes his interest in models that can develop new knowledge, with genuine scientific or mathematical discoveries as a meaningful measure of progress.
55:04	Verification and closing	The hosts briefly discuss verifying discoveries in domains such as biology and mathematics before wrapping up the conversation.

Quotes

we’re going to use like kind of attention as the retriever here rather than like building like a nearest neighbor index or something like that. — 01:39

once you get to that like 10,00 documents, you know, 2500 documents, 5,000 that 10,000, like a million tokens basically, you know, things really start to collapse really fast. — 09:26

Is there like a sublinear version of of attention? — 39:56

it might uncover something super interesting. — 48:50

Concepts and entities

In-Context Retrieval, generative retrieval, listwise reranking, attention dilution, sparse attention, MS MARCO, ColBERT, ReFrag, Siddharth Gollapudi, Weaviate, Qwen

Full transcript (original language)

00:05 [music] [music] [music] [music] [music]

00:36 [music] Sidarthth, welcome to the WVA podcast. Thank you so much for joining. >> Yeah, I’m really excited to be here. Um, yeah, excited to get talking. >> Awesome. Yeah. Well, I think your paper is so interesting. This topic of in context retrieval, I think there are so many different flavors of it. The son of search world has the whole like, you know, is rag dead? that is search dead and the long context LMS versus search. I’m really excited to dive into that aspect of the story. I’m also really excited about re-ranking and the idea of

01:07 using this for re-ranking and I guess generally just learning more about the models you’re building, the experiments you run and all just seems really interesting. So yeah, thanks so much for joining. >> Yeah, I’m excited excited to get chatting about this. Yeah, this is a great this is a great project. So >> awesome. Uh so can we maybe kick it off for our audience? Um how do you define in context retrieval? Yeah. So the idea behind in context retrieval is you know it’s it’s it’s kind of it’s similar to a lot of like other you know retrieval settings in that like given a bunch of documents you have a query you want to get the most kind of relevant document.

01:39 So for the query that is but in this case you know we instead of saying you know using like an encoder model to get embeddings or something we’re actually just going to put all the documents into the context of like an LLM uh and then given like a query kind of prompted into the LLM afterwards uh we’re going to just find the most relevant document within the context. So essentially we’re going to use like kind of attention as the retriever here rather than like building like a nearest neighbor index or something like that. Yeah. >> Yeah. I think Oh, sorry.

02:10 >> Yeah, that’s probably the easiest way to explain it. >> Yeah, one thing in your paper is differentiating it from uh generative search. I think this the little distinction there is really interesting because generative search is like you have to train a model to predict IDs compared to this setting where it’s more flexible and it reminds me maybe more of a reranker. >> Yeah. So the generative I would say generative retrieval is like you’re right in that it’s like super similar to the setting in that you know it also kind of spits out an ID. Um it’s based

02:41 on you know also using an LLM as well. Um but the key difference of course I think like you mentioned is that it’s like quite inflexible because you know you have to train the model uh to put like documents essentially into its parameters. Um and so every time you know you just want to update the corpus like you have to train the model again you know you have to do gradient updates which is like really really unwieldy especially at scale. Um and it’s also really really noisy like in general for parametric retrieval I guess because like generative retrieval pulls the

03:12 information from the parameters. Um like it’s it’s you can’t really like guarantee that you’re going to get the right information every single time. you might be able to like kind of vaguely get something in that ballpark, but you know like if I want to like exactly figure out which doc like which textbook I’m getting, you know, a math theorem from uh there’s no guarantee that like I could pull that exact information even though I might be able to pull the theorem. Um so in that sense, you know, contextual retrieval maybe is a bit more exact and like verb you can get things

03:43 verbatim. >> Yeah. >> Yeah. Yeah, I think maybe just like with a helicopter view before diving in further, I’m also curious how you think about the sort of lost in the middle phenomenon. I think that paper came out in 2023 and has like over 5,000 citations that showed that back then with with these 4,000 input like I think GPT3.5 models they had this U-shaped curve where basically if you gave more than five search results it was like hope or or and then particularly if the answer wasn’t if the right document wasn’t just ranked number one it was

04:14 hopeless. So maybe curious just what you think about the lost in the middle problem in here in 2026 almost 2027 and >> yeah it’s not I I would say like we didn’t actually experience like a lot of issues with this when we were playing around with uh in context retrieval um I think over the years and this is maybe like more hypothesizing than like you know uh true you know research has been verified and so on. Um, but my guess here is that like so I guess originally lost in the middle was like this as you

04:45 said like this phenomenon that kind of was appearing in some of the earlier LLMs where you know like language models are really good at getting information from really early on in their context and at towards the end of their context. Um so like there’s this kind of this like just basically this positional bias uh that came about. Um and my hypothesis there was like okay first of all you know the the the human data kind of led to this. So basically maybe like their instruction following data maybe even their pre-training data you would find

05:17 like the most relevant information for solving or doing tasks was probably towards the start of a like some document or cor samples in a corpus or towards the end and that kind of maybe led to one kind of positional bias. Um the other thing is maybe like attention itself was not you know like as heavily engineered as it is today. Um, so yeah, like I think as a result of that, you know, like late 2025, early 2026 when we were working on this project, we didn’t really see as many issues with this as we as you know, maybe we would have

05:48 expected. Uh, because like you do, like you said, you know, this paper was pretty famous. Um, but I guess it’s also maybe spurred a lot of um like improvement into like long context and you know, making sure that models can like address every part of their context um at the same level. >> Yeah. And I’m so excited to dive into the new neural architectures and what you’re finding with attention. It reminds me of a podcast I did with Opir Press about alibi attention. I think it was like 60. I can’t even remember what number podcast. This was so long ago, but it was like the way that uh length

06:20 generalization was achieved just by changing the positional encoding. It was like so straightforward sort of that it was and then what you say about the training data that also was so interesting. And I guess just maybe also I want to just kind of set the stage more on in context retrieval. Um, how do you think this differs from reranking like a listwise reranker? >> Yeah, so I would say so there’s actually been like a lot of prior work on using LLMs as like rerankers as you said like listwise reankers. Um, in fact our paper was based on like kind of built on some

06:51 earlier work uh you just like that took a 100 documents and tried to rerank them. Um, but like I’d say reranking is just generally a much easier task than you know first stage retrieval here because um, you know, you only maybe need like 50,000 to 100,000 tokens worth of documents at least in like this like standard like academic retrieval settings. Um, and it’s it’s like because of that you don’t really have that like kind of long context issue. You just have to worry about how we train the model for this like fixed width context

07:22 and that kind of makes the problem a bit easier. Um, probably the only place where I’d say that the problem might be a little bit trickier is where you know, you know that all the documents that you’re looking at are going to be hard negatives. Um, but I would say argue that like if we’re able to like truly kind of solve this like long context issue, um, an LLM should be good enough, you know, given like a sufficient number of forward passes, like some good chain of thought. Um, it should be able to do like this first stage retrieval into

07:52 reranking kind of in one fell swoop, like one one like set generation. Um, we just kind of have to like train it properly. Um, yeah. So it’s definitely like a much it’s it’s definitely like a superset problem. You know, we we also have to be able to solve this long context problem uh to do like first stage retrieval, but it’s not like it it re-ranking is definitely a part of that. >> Yeah. >> Yeah. My context of how I first became aware of your paper is I I met with Omar at a blue bottle in Cambridge and we

08:23 were talking about drowning in documents and um this work from Matthew Jacob and Omar also and >> and that that was and he you know he recommended me your paper and then I happened you happened to be at VLDB and it was just perfect to be able to meet and talk about this a bit but that’s how I became aware of this context is like this is a solution to drawing in documents and I looked at it from I’ve been looking at it from that angle that it’s like can we listwise re even can we list wise we rank up to 500. >> I I would say like the I I would actually like kind of argue that this

08:53 paper really tries to go beyond that because 500 documents I I I think like our paper actually does really really well at that 500 document scale. Um, I would actually say that a lot of projects can like really nail that 500 document scale because it’s really easy to, you know, take like a GPUs like 8A100s, H100s and like train on, you know, 256 document or 500 document corpus scale. Um, but like once you get to that like 10,00 documents, you know, 2500 documents, 5,000 that 10,000, like

09:26 a million tokens basically, um, you know, things really start to collapse really fast. I I would actually I think like that’s really where um like the exciting kind of research is you know like how do we how do we generalize because you know there’s no way even if we if we eventually wanted to get to web scale um you know billion documents there’s no really good way to like scale training to that scale. So you know we would actually want to be able we we need this generalization to to like kind

09:56 of get even to that point. So just like like even kind of like moving beyond reranking, you know, how do we get to that this like link generalization point? [laughter] >> Yeah, it sounds really ambitious. I think the intermediate goal is also really exciting. Like I think just uh list-wise I think people would like the takeaway from the Johning and documents paper that um you know cross encoders can’t really scale that well and that you’d expect to just be able to rerank more and get better search. But I’m I I don’t think there are many papers that

10:26 explore like a listwise rerank or like LLM setwise, however you want to call the problem, but like an LLM reranking like a hundred documents at a time. I think just that right there is actually where the state-of-the-art is. Yeah, I’m only aware of maybe like a couple there’s a couple of papers out there that do kind of like kind of think about list risking, but I would say it’s it’s definitely I would say like probably a lot easier

10:56 now that context windows, especially for Frontier models, um are like so good at this like 100 to like 250k context window, you know, like I could probably pass in, you know, that many documents to like a frontier model like Gemini or Claude um and get pretty good results at this point. It’s it’s still very much like that stage beyond like just not even you know like just like doubling the number of tokens like 500,000 or a million tokens and

11:27 then like trying to do reranking or trying to do like first stage retrieval even um that’s like still really really hard because that generalization problem hasn’t really been solved. >> Yeah. And and I think maybe like almost into coming into block search and the model you’ve trained, but I think like just getting you know a 600 million parameter model to do this. I think like probably you could get this to work with Opus 5.5 or like GPT6 soul. But then the question more interesting question is like six Luna with like this 10 cents

11:57 per 1 million token pricing. And so the the idea of also already ambitious to be scaling re-ranking even though I’m not sure like where the cutting edge LMS are. And then I think there’s also like little work on really training these kind of models. It’s I think it’s just like kind of interesting how underexplored this is. >> Yeah, that that was actually probably one of the first things that we we were looking at too when when we started looking at this problem. Um which is that you know nobody’s really looked at this. So like we’ve seen like a little bit of work on you know reranking as you’ve said. Um we’ve seen a little bit

12:29 of work on you know can like language models like do how good are language models I guess at search but like nothing really kind of you know end to end evaluation how do we make this work um so on and so forth like there’s there’s there’s some work from a few years ago on like how do frontier models do at search but you can’t really train or like look at attention scores for frontier models like nobody no frontier lab will ever publish that nowadays so I I think like how we how exactly the

12:59 science works here um has been like woefully underexplored as you say. Yeah. >> Yeah. And so I think probably you know without riffing too much I think our audience would love to learn more about the models you’re training and you mentioned you know the attention scores and going into it. I think this is absolutely one of the most captivating things about the this paper. >> Yeah. So I would say like something that we really kind of tried to like drill down on. So when we were first I guess like playing around with how to best kind of get past this um like get into

13:31 this like length generalization setup. So just like just to maybe kind of set the stage here, you know, we probably had enough like GPU memory to train on like 256 document corpora. So like given just like standard MS Marco data, you know, we had probably enough space to train on 256 documents. So we really needed like a but we really wanted a way to like get to that million token scale. like that million token was like this really nice kind of wow factor point. So um getting to that was like a huge challenge. So we kind of like went about

14:02 it in stages like first can we get to like 500 documents then can we get to a,000 documents like get good performance at 2500 and so on and so like just in getting to like 500 documents we found like a boatload of challenges. So like you know prior work in like listwise reranking um you know they used for example like there was some work out there that would use like positional codes. Um we actually found that you know if you use positional codes like the model learns to like expect certain documents at certain positions and you

14:33 know you get worse performance. Um and if you use like you know codes of like way too long a length you know it would like get way too noisy. you like start hallucinating, you know, answers to the codes, uh, answers to the the queries, sorry. Um, and so we found in general just like kind of just using random codes, forcing the model, like basically encouraging the model, I should say, to um, like read the documents and like you could see this actually like the attention to the documents like go kind of up versus like to the codes like once you once you

15:05 like kind of randomize them. Um so yeah we found that like randomizing the codes like initially led to like pretty good generalization like we we were able to see like much better performance at like you know this like 500 to like thousand document scale. Um moving beyond that was a bit of a challenge though like that that’s really where we started to like really have to engineer um engineer this like the model training. So I my memory is like a bit fuzzy on this. So like maybe like a,000 to 25 thou or sorry 1,000 to like 2500 documents like

15:36 this is really where we had to like start playing around with you know like making sure like all of our learning rate and like hyperparameter tuning was like super set. Um and we also had to start playing around with like some sort of on policy uh learning. So basically like you know how where are we messing up in the code generation. So maybe to like give an example here, you know, um if I had like some query and you know a thousand documents, it was pretty common for the model to like get kind of close.

16:07 Um like it would like spit out may it would spit out maybe like part of the right code, but then like the other the rest of the code would be kind of wrong. Um and this is where maybe like I thought some sort of like on policy modeling could help. um in that you know at training time we like ask the model to like roll out a couple of generations just to see how well it’s like kind of staying on track um and then like sort of like correct it at each like generation step uh in case it like made any mistakes uh on the lowit level um

16:37 and this actually helped a little bit. It didn’t help as much as I thought it would at like true like million token scale but at like that 1,00 to 25,000 2500 document scale it helped quite a bit. Um it’s it’s it was it was surprising actually I would say that just like correcting the rollouts was was so helpful. Um yeah but despite like I guess all of these changes um getting to that like 5,000 and like 10,000 documents so like a million token scale um was like really really hard. uh we we couldn’t really

17:10 figure it out and we just like using plane you know like the way we trained our model uh so far like use like kind of like post talk changes like oh you know can we change the codes you know can we like change the way we train the model like optimizer or like loss function stuff like that objectives um so like that’s that’s at the point where we like started thinking about okay you know can we change the architecture of the model um and that’s I guess like kind of what led to that like second half of the paper where we like discuss because you know what’s like going wrong

17:40 with the model internally and how can we like justify changes um to it. >> Yeah. >> Yeah. That first thing uh you’re mentioning about the ids that reminds me of like the semantic ID work that in that’s popular in generative retrieval where you use like a VQVE and you have like discrete identifiers and a hierarchy of that to to help with that problem. And um yeah, I guess um yeah and and the other thing you mentioned about the kind of bias with the first stage retriever I think ex you know expecting certain documents somewhere I

18:10 think that’s a super fascinating uh entanglement probably also with as we talked about loss in the middle that idea that the models might expect like if it’s hey here are the results from a search engine it’s probably going to expect like them to be ranked properly and like you cannot you can never underestimate the biases and priors in these models and um I guess Um maybe even before coming into the architecture, I’m a little curious further into what you how you use the term on policy and just how you’re how you’re kind of thinking about on policy versus off policy with training these

18:40 particular kinds of models. >> Uh yeah. So the the like main idea behind on policy is that like I actually wanted the model to generate codes um just to like see you know just like kind of give the like and I guess like in this sense we kind of had a vague idea of like you know what was going wrong here. So like every once in a while maybe like there’s like a super hard negatives or like a couple of hard negatives or something in the corpus and you know the model like starts out good but then like maybe gets derailed because it sees like another possibly

19:11 relevant um document. So like it kind of like gets split like stuck between a rock and a hard place sort of thing where you know it it outputs like part of the code for like the right document maybe and then sort of gets distracted by another document like midway through or something. >> Yeah. like this is it was it was pretty funny. Um and the most common like failure mode here I would say was uh like it would usually like it would at the at the very first digit it would like say oh you know I’ve got like a

19:42 couple of documents here that seem pretty good. Um but you know this one seems the best let me generate this document or let me generate the code for this document. Um, and this like usually would like cascade down like maybe it would generate another code but then like it would like sort of get like distracted and you know go off to like generate the code for another document. Another failure mode would be you know at that very first digit you know there’s like again five do like five really good-look documents. Um like a couple of them all have the same first

20:12 digit. So like it’s not really sure which one to pick. Like they all look pretty good. So you know it’ll generate like one of them and then because of that confusion at that first digit it’s it like kind of percolates to the next step even if they have separate digits. Um so like that tiebreaking also like sort of led to like some sort of internal strife I guess within the model. Um that’s like a like like a bit of a hypothesizing thing. But I guess like the kind of the the key takeaway here for us was, you know, how do we make sure the model like stays disciplined and like outputs the code

20:43 for like the document it like first latches on to even if it’s like kind of wrong, you know, if it’s problematic. Um just like staying disciplined and like outputting real codes. Um and that’s really where like the on that was like kind of the big motivation for the on this like on policy loss. So like during the training, you know, we had the model do like a couple of rollouts on like some of the the questions and the queries in the batch. Um and just based on those codes, we’re able to like correct it kind of step by step um like which parts go wrong. And that like it

21:14 actually helped quite a bit mostly at these like kind of smaller scales like,000 to 2500 documents. Um past that you know there are other issues but we actually did see like quite a bit of improvement there uh just by including this on policy change. That reminds me of like um like structured outputs or constrain generation for listwise reranking where >> Yeah. >> Yeah. I guess like you usually have a document identifier and um I’ve actually only seen pretty simple cases where it’s just like square brackets one square brackets two and then you’re just constraining that it’s a you know list

21:46 of ints and then and then the and then the really tricky thing is making sure each one is uniquely appearing once. That’s like the that’s the galaxy brain version of structured outputs there. Yeah. >> Yeah. We we played around actually with like structured outputs a little bit. Um, we actually found quite curiously in this case like we needed to minimize the generation as much as possible because especially at these like super long like like long context lengths like the the the the possibilities of you know hallucination and noisy outputs um just

22:19 keeps going up and like we couldn’t afford to like take that chance with like enforcing structure like we just needed to get the code and like kind of be done. Yeah. Uh yeah, I’m curious just further into that point. I mean, we’ve seen some things where like we do a lot like with WVA’s query agent, we parse schemas and we get some pretty complicated JSON schemas that we then ask for structured outputs with like solving text SQL that way. And um is that what you mean by that? like you have the constraint generation is so complex that or or is

22:51 it kind of like the let me speak freely idea where it’s like just constraining the model like that complicates the training like >> uh I wouldn’t say it’s like the latter it’s like it’s it’s a little bit of the former and that you know it’s it’s difficult to train it’s definitely difficult to train um I will say that like maybe like like one kind of small caveat I would say to our paper is that um it it was only post- trained. So we only took like maybe like a million tokens

23:22 worth of data. Um and as a result like we were kind of constrained in what we could do. Um which I think was like also a good thing and like led to like some interesting research. Um, but yeah, I think like at an engineering level, if you really wanted a good model for this that could, you know, check all the boxes, like do the generalization, um, like be super robust for reranking, so on and so forth, you’d actually really want to, you know, get like a billion, two billion, like way more tokens, order a magnitude more tokens, um, and like do like some sort of continued pre-training

23:53 or even like just like start from scratch pre-training. Um, that way you just like you can get exactly the behavior you want. like if you do want like a like good constrained structured output um you can expect that to be like a lot more robust and consistent over um like based on a pre-trained model rather than a post-trained one. >> Yeah, that’s a super cool like foundation model for search type of pitch. Yeah. Um >> I guess I’m just quickly curious like I you know think about the million token scale. I’m curious about overfitting and how you think about generalization

24:24 because if you have the same corpus that you’re always sort of embedding and then you have these queries and just then you swip switch switch out the corpus for a different one and just I guess just how you think about because now you have to like your data mix has to like mix in different corpora. Is that how you think about that? >> Yeah. So we actually originally were using only MS Marco. Um, so yeah, I guess like just broadly trading data. Um, we were originally only using uh MS Marco. Um,

24:55 this actually worked like decently well, but I think like the issue with like raw like just like plain MS Marco is that um it’s a bit noisy. So like there’s a lot of false negatives in it, you know, and there’s and to like maybe like give you know the community credit, uh there’s been a lot of like good work out there on like kind of cleaning it up. So there’s like been some recent work there was like some recent work at the time from um Nandan Takur on uh like RLHN I believe where like they used LLM like clean the data and it was like a pretty like rigorous system. Um, and so we used

25:27 that data and we like saw like much much better performance just like purely on like how well does retrieval do on MS Marco but then we like also saw the issue that you’re saying you know where like okay if we only use MS Marco like how do we do on like other bare data sets like hop QA and you know ENQ and stuff and like there you know performance wasn’t great so like okay let’s mix in a little bit of like data across the entire bear suite um into like our kind of RHN training mix um and then we found like okay we’re able to get like a little bit more kind of on

25:58 par performance with like just using plain MS Marco but the the performance improvement I would say from doing this was like not you know that big so like I mean we’ve seen this I guess in like training models like Colbert and um like other single vector retrievers but especially Colbert where you know like for example Colbert V2 only had like an MSM trading mix but it still performed pretty well like out of distribution on like other bare data sets Um and I suspect it was like a similar thing here and they’re like okay yeah while the

26:29 performance was like slightly worse and like in including like the extra data did help. Um at the end of the day you know the models like learn this like algorithm of relevance. So like and the bare data sets are anyways like pretty similar across the board. Um so yeah I think like it it would it definitely did help to like include extra training data but it was like not this like big you know needle moving thing. >> Oh yeah. So interesting. Yeah. I guess I think like it probably like once it learns an embedding of a word in a broad

26:59 scale like web search context it probably that probably generalizes pretty well and you can maybe argue that there’s like some contextual generalization that happens in like particular domains. I I don’t I don’t know if fine-tuning embedding models is super I heard a great argument that fine-tuning Colbert models is particularly really like a really powerful thing to do because of how you have the per token representations and >> yeah it’s just taking me into that type of thinking but um >> yeah I’d love to maybe come back into this kind of you know digging into the neural architecture I think that’s yeah

27:29 really unique you know to go that deep >> sorry actually I want to you mentioned co having this like per token representation you actually have that here as well in fact like I guess Think the whole point of Colbert is that it appro it would approximate something like this and that you know like you get you get these like representations of all these documents at the very end rather than like going through like this like huge forward pass where like you’d keep studying like interactions with the tension between all these tokens. So I’d say you know maybe Colra is like this like loose like late interaction type approximation of you know what we were

28:00 doing here actually. >> Yeah it’s a and I I wanted to bring up like the refrag model architecture. Yeah. And I guess like this late interaction concept, I guess there’s so many different flavors of it. And I think like when Omar is explaining Cobra, he’s really vocal about, you know, maxim is just one kind of late interaction. >> Yeah. >> That’s a pretty cool idea. Yeah. So I’d love to maybe start talking about the attention dilution effect you find and I guess just understanding like neural architectures for long context. >> Yeah. So attention dilution part of the

28:33 big problem I guess with like scaling models into these like ultra long context settings is that you know um there’s like this like ultimate like this this like fundamental let me put it that way um like upper bound on performance and it’s basically like hardware based you know um where you’re just limited by how much information like a BF-16 or like FP16 uh uh data can data type can like hold. Um, and to like that end, you know, attention dilution

29:05 basically says like, you know, the more information you have because of how softmax works, um, the scores are just going to get like very very similar to one another because the denominator gets so large relative to the numerator at some point that regardless of what you do, you know, you’re just going to end up with like super super similar um, looking scores for all of your tokens. Okay, I said like regardless of what you do, but that there’s a there there are some things you can do actually. Um, so usually what people do when they see

29:35 issues like this is that um they do like one of two things. The first one is um you know you can like make the like the softmax more peaky again. And so like there’s like pretty simple you know you just like divide the numerator and the denominator like in the exponential bits with with like a like a flat temperature like just a simple scaler. Um, and that’ll like make it, you know, peaky again. Uh, but like the problem with doing that is that, you know, temperature it kind of varies depending

30:06 on like how much like quoteunquote stuff is in the denominator. Like if you if you don’t have a lot of stuff, then maybe you want a different temperature versus when you have a lot of stuff. So like now you have to worry about like making it dynamic. So like that’s like kind of like something people have been looking at in the context in this like long context setting. Um I probably won’t we’ll go a little bit maybe into that. I’ll go a little bit into that in a bit. Um but the other thing that I did want to mention was um this like notion of I guess okay you’ve got a lot of stuff in your denominator.

30:37 Now what if we could just like pre-process quote unquote pre-process the denominator and like get rid of like a decent chunk of the stuff, right? And so the motivation here is that you know like for attention like the vast majority of the tokens that like a query token where they tend to probably aren’t that useful you know like you really only need like a small subset of them like very very few of them in some sense like to the like re-ranking point you know you probably only actually need to worry about a few of these tokens like you probably could just pre-process like

31:07 a whole boatload of them just worry about like a small subset of them um and this is like I guess like pretty similar to rag and that you know we could like segment the tokens document. In fact, this is actually what we did. We tested this. We you could like segment the tokens document by document um and then like take the query tokens and like before you actually run attention basically you know you could like like look at all of the tokens corresponding to each document and based on the ones that like have like these like low you know dot product scores or like low maximum scores. You could you could pick

31:39 your you know notion of similarity. Um based on that you could just like decide okay these are the ones that like have no relevance. We’re going to actually just drop them from the attention computation. These are the ones that have high relevance. You know we’re going to actually use them in the attention computation. Um and from that from that sense like you know things actually work pretty well. Uh that’s that that was like one of the big like kind of performance improvement things in like the latter half of our paper once we were like all right you know million tokens doesn’t really work with our base model. What needs to change? So

32:11 that was one of the things that that really seems to work. The other thing is of course this like messing around with temperature scaling. So like there’s been some prior work that’s like also kind of identified this attention dilution issue. Like for example like the the the the prior drowning in documents paper that you were mentioning did also like kind of allude to this problem. Um the issue that so like I guess like the issue with like like this temperature scaling compared to um so I guess like maybe I should actually first explain why temp the temp the temperature

32:42 scaling works here. Um so the idea here is that basically like you can add like a par like a parameter or like some sort of function that’s like parameterized by length into the softmax uh numerator. And so the reason you want to do this is because even though the denominator scales with the length, right? You’re just like you’re keep you’re summing up over all of your tokens, the numerator doesn’t scale with the length. And that’s what really leads to that like kind of wash out effect where like all the scores end up looking really really similar. So you could just add like some you know multiplicative or like

33:13 parameter into the exponential that’s you know as a a function of the length n, right? Um and as a result like the scores like stay like relatively consistent like with respect to length. Um and that seems to al also help performance quite a bit. So like we actually put both of these together. Um and by the end of the paper you know uh we actually found that it was like at about the same level as dense retrieval which is like still really not great you know like you’d expect like at million tokens that is you know you’d expect

33:43 like like this like all powerful you know really well-trained decoder language model to be able to like significantly outperform a you know dense retriever but uh I think that’s a that was like a it’s a pretty good start I would say like there’s definitely like a lot more interesting research to be done here about architecture because you know I feel We did like a pretty good job in this paper of establishing that you know there are like cool tricks you can do with respect to the data and the training but at some point like it’s just like the architectures of like the quen 3 generation just aren’t enough to you know um get past that like this like

34:16 certain like inflection point where attention just becomes impossible for the model to use. >> Yeah. And I I definitely want to anchor with your thoughts on just like what kind of timeline you think is reasonable for that kind of thing. But I if if you don’t mind me just nerding out on the details a little further on the kind of sparse attention you’re mentioning. Um >> so I guess like the only idea I’m really familiar like that I feel like I understand well is with refrag you have this idea where like you take a chunk of like you have a document chunk and you either represent it with one vector or

34:47 you have a policy in the middle that’s like choosing when to expand some of these chunks into a higher you know more fine grain representation from say like one vector into 32 vectors. And so that kind of like policy expansion that that’s all I really understand for like this kind of I don’t know if you call that sparse attention. I mean it kind of is the same like concept. Yes. I guess if if you don’t mind just like I know you just explained it well but if you don’t mind just kind of explaining a little further like like what do you mean by the like you you can

35:17 throw away certain chunks. >> Yeah. So I you actually kind of like brought up I guess like the high level idea behind this which is is like you know it sparse attention right and this is where I really uh you know because you mentioned that you know this like vector DB workshop where at at VLDB you know this is like really like a great opportunity I think for you know like retrieval research because this become I think like on a multiple multitude of levels like sparse not just like attention dilution sparse sparse attention is like a great opportunity

35:47 for you know speeding up attention in the first place like attention is like this quadratic like bottleneck basically on you know LLM efficiency at scale right so um the idea behind sparse attention in general is like you know for example with refrag you know what we were talking about with these like document based um document based kind of token aggregations and so on um is that we can like basically just pre-process the tokens that go into attention so we want to make sure that like attention sees like a very like limited set of

36:19 tokens because if see more than that then you know you run into this dilution issue. Um it gets really really slow. Um there’s like a whole host of issues. So we want to be able to pre-process the number of tokens that go into the attention mechanism. Um and you know there are like great sublinear methods that we you know we could come up with that like vector search has been playing around with for you know gazillions of years. Um but like I guess like the the true kind of challenge here is that you know a lot of the approaches that you

36:50 see today um are kind of post talk. They’re not really like completely and utterly baked into the model at pre-training. And the reason for that, I think, uh, and this is like again like hypothesizing. I’m not really like completely 100% sure about the the current state-of-the-art literature here, is that like there’s no like good pre-trained way to like load this these like some sort of retrieval engine into a model. Like how do we do like, you know, we could do a forward pass, right? Like just simply like do the pre-processing, but then how do we do

37:20 the backwards pass through something like this? you know, like what would a backwards pass through a retrieval engine even look like? Um, and so in that sense and like what’s really exciting here is that, you know, sparse attention could do like could like solve a lot of problems with like LLM generation in that, you know, we could make your way faster. We can like significantly improve long context capabilities, you know. Um, it could solve so like we could solve all of these issues like the both the efficiency and accuracy sides of the trade-off. Um, but like we still haven’t

37:51 really figured out how to like really make it part of like an end to end trade model. So that’s like a huge kind of open challenge I would say. >> Yeah, that inspires so many so many ideas for me. I guess just like the first thing is I remember back when hallucination with rag was such a big deal and I I’m just like citing times I’ve met people in Cambridge. I was talking to Nathan Lambert about this at the Diesel Cafe in Somerville. [laughter] But but this idea of like why don’t we train the models with rag with retrieval in the input and that would probably fix

38:23 that problem and then remember meeting with contextual AI and and they were trying to do they call they used to call rag like Frankenstein rag and we need to put gradients from the generator into the retriever but then you had this horrible problem that you’re describing of like you need to keep updating the retriever as you’re training it and then then that becomes a problem because the index is stale and and I mean right I’m I’m totally distracted myself but right now there’s a lot of cool work on uh updating the index with a new embedding without reindexing the whole thing. And so that’s an interesting tab to open. But then yeah, I guess just coming back

38:54 home, this idea of how do we inject precomputed representations into the models? It sounds like the most compelling future for vector databases. It sounds like just yeah pretty cool option >> you know even more than just like even to like so precomputed representations into the models like so that’s like maybe a bit more that’s definitely like interesting you know like rag in general is like still like I would say alive and well and you know you still need it to like issue like handle like tail queries but like even just like simpler than that you know um if we have a lot of

39:25 information in a context like how do we identify what’s the most important parts of that context like it like maybe going back to this like ICR setting of you know we’re just dumping all of our documents into the context we don’t know what’s the most important so in that sense even like retrieval suddenly becomes really important because like as you know like you and I like we we’ve worked on vector search for quite some time like this seems like a classic you know sublinear time like a&n type problem where like most of this stuff is useless we should be able to like very efficiently like get rid of it so that

39:56 we can focus on the important parts u but people really haven’t figured that out yet so In that sense, in fact, it doesn’t even matter like if we keep like an index fresh or not. Like just quite simply, is there like a sublinear version of of attention? >> Yeah. Yeah. Super cool. I mean, such a cool topic. And um yeah, it absolutely definitely I I mean, I guess for me it’s like is open AI or anthropic can they just are they just going to train one of these models? But but yeah, I mean the work that Yeah, you’re leading. seen

40:27 we’ve seen a little bit of it from Deep Seek and some of the other um like Chinese labs like Alibaba for Quen um but like it’s getting there you know it’s getting there maybe we might see something within the next year or so. Yeah. >> Yeah. We need the vector database summer [laughter] like these are coming out but um yeah awesome. So um I yeah I guess I’d love to come back to kind of you know we were talking a bit about like the current state of the experiments is it’s difficult to scale this long context ranking from 500 tokens to a thousand tokens and I guess I’d just

40:59 love to know your current sentiment on like you know you you mentioned like working with eight GPUs you think it’s just a matter of scaling up the resources like what do you think is needed to push it >> so we didn’t I think like at the time let’s see what’s what’s needed to push it. Excuse me. Um, so I guess like the idea here that I would like kind of push on is um we really want to like essentially just really make this well engineered.

41:30 So like I actually think you could squeeze out quite a bit more performance from like a properly trained model uh on a GPUs. Um but I guess like training the model um would you would want like like a bit more GPUs so that you could like like efficiently pre-train at you know like some really large token scale. So not even like like larger more parameters per se but you know if I want to like train on a trillion tokens you know doing that on a GPUs is like kind of annoying you know you could do it don’t get me wrong but it’s it’s like

42:01 you know it’ll take a couple weeks so like I don’t really want to do that. Um, on the flip side, you know, you could just like if we had more GPUs, you could like really nicely kind of efficiently prototype and iterate um, and like see, okay, like these are the architectural improvements that do help, you know, these are the architectural improvements that don’t help, [clears throat] excuse me. And that would really help in like just getting from that like 500 to,000 um, to like 2500 scale. Um, yeah. I I

42:31 mean, like, yeah, I think like the big point here that I definitely want to drive home is that, you know, less so from like a like an engineering perspective, we really just want to be able to like see like, okay, what how do things work at pre-training? So, in that sense, we would want more GPUs, but don’t get me wrong, you would also be able to like really really use them for like better engineering and stuff. >> Yeah. I almost feel Well, yeah. I remember like in the early neural architecture search papers with like Quac Lee from Google and you’d be like we figured out how to use all these GPUs and that’s like the and then it almost

43:04 reminds me of thinking like having more agents going it’s almost like this question of do you want more cloud cloud code usage and just like yeah to do more research of course like to have more resources like more resources would be great >> yeah you can never go wrong with more resources like I I could you know we could always come up with like we could always we could see sit here like all day brainstorming you know oh Yeah. What what we could do with more and more GPUs. So like yeah, if I had to pick one thing, you know, it would be like, okay, yeah, great. Now we can pre-train a little faster, that could open up some more possibilities. >> Yeah. And I guess like just um how much

43:37 would it surprise you? Maybe that maybe that’s like an interesting way to phrase the question like how surprised would you be if a re-ranker on the scale of say yeah 500,000 is comes out before the end of 2027? I’m actually like really surprised there already isn’t a re-ranker because I think because like for example you know we were talking about like you know retrieval and we were like kind of under the impression that you know like re-ranking already is like probably not too far off. So like my suspicion here is that like more so than you know like

44:10 a retrieval having not come out like people are still pretty happy with like you know how current re-rankers already work. So maybe they haven’t like put a lot of like effort into like some like big fancy like LLM based one, but um I would say that like yeah, I would say like to maybe like just like I’m actually kind of surprised that like one isn’t already out. Like it seems like it seems like relatively speaking it’s a much simpler task. You know, you really have to just worry about like the data quality um like the like how you train

44:41 the model. um you don’t have to worry too much about model architecture I feel because you you have a very fixed context length and you you know that you’re always going to be looking at you know 500 to a thousand documents at most probably for a reranker and you know that regime is already like pretty well captured by like a lot of like open LLMs like you know Gemma I think has like 128,000 context length now so like it is kind of surprising that people having like like there aren’t reankers out there already and I’m in fact I’m sure

45:11 there might be. So >> yeah. And I guess um I think like probably the search world objection would be like the latency of it is probably why not because I guess like in our world at WV8 we have we deal with a lot of like >> you know sub50 millisecond latency that that type of thing and so >> I feel like that’s probably like this kind of like what’s the application of high latency search like um we have like search mode with query agent we’re really interested in pursuing this and it’s like like with deep research agents or I I don’t think like code and code in

45:41 search is an interesting one because is like GP is so powerful there and I don’t think like vector service >> is definitely more powerful there. Yeah. >> Yeah. But maybe >> Yeah, I agree. So I think maybe the latency is like the big issue here. Um you definitely I mean I guess if you don’t mind me asking like what are like the parameter sizes of the three rankers that you use in like production at at Wev8? Uh I don’t know if we’ve been open about

46:13 that. >> Oh yeah, [laughter] I guess like on that SC like I’m I’m trying to like think like parameter scale, you know? Um like that’s probably probably at least like the latency issues for uh um for like these like LLM based reankers. But yeah, I could totally see them being like super relevant for these like harder questions. Um like these like AI mode style questions maybe to like put it in like Google terms. Um these AI mode style questions where um

46:45 you still need like to be able to recover a lot of information but the user is going to give you a little bit more leeway in terms of like how fast they want it. >> Yeah. And and then um so that motivates me. Another question I think is really interesting as we talk about like applications of high latency search or maybe like multihop search or deep research style use cases and then but then it’s like well if you have a long context LLM, can it just do this kind of multihop search in one pass? I even saw a really cool paper this morning from your adviser at Berkeley that’s um uh

47:17 titled no more free lunch corpus task complexity matters as corpora grow and it and I think it’s on this topic of like you know can LLM really like long context LM like eat the database away what like a database does for you >> yeah that was a great paper out of so the paper you’re mentioning is a like is a great paper out of my lab mate Prasan uh he he did a lot of great work on that over the past year um and I I I actually think like that’s really where we want long context LLMs to be. Um I guess like

47:48 he kind of pointed out like they’re not very good at it right now because like like long context training tends to focus a lot on like like search based tasks like NIA and you know MRCR and so on. Um and like tool use I guess and like long chats. So like there are these like these like much more complicated tasks, these like quoteunquote higher complexity tasks um that like long context language models aren’t really good at. Um, and I think that’s really where like we want long context language models to be like that’s really where they could shine because like when we

48:20 kind of train models to do things um like the way we want to do them to do it I guess um like that kind of introduces this like inductive bias um of like how humans do things the model should mimic it. Um, but like I guess like maybe like part of the allure of a long context model is that you know you could possibly just dump all a bunch of information to the language model. Um, and then like sort of let the language model go ham in an unsupervised fashion. Maybe like actually is able to like

48:50 discover something new from like all of its all the like training it’s done that you know we as humans would have not been able to notice. So, um I think that’s really the promise like it can like discover these like uncovered connections between like different pieces of data um that you might not be able to get if you know for example you were using like agentic search like you know you’re talking about GP like you’re using agentic search and like you g for one thing and you g for another thing and like there’s an expectation of like the relationship between those two pieces of greed information. Um here we

49:20 could just like give it a bunch of information and it has to like figure out excuse me it has to figure out like what’s relevant. Um, and it might in doing so it might uncover something super interesting. >> And that’s kind of a big takeaway I had from uh hosting Alex Zang on the podcast about recursive language models. And >> yeah, it’s like they learn they learn like decomposition like in the long context. Yeah. And true. >> Yeah. I actually I really like this like notion of unsupervised like learning here because you know I think we’re like very quickly like at the frontier level

49:53 you know like we’re very quickly reaching this like point where you know these models are able to do so many things really really really well and it would be super cool to see like okay now how do we like kind of move beyond that how are we like consistently going to build new and get new information out of these language models um yeah >> yeah and I guess I like I I guess like this idea of just well I guess there’s like two thing like because at Alex’s like the recursive language models you have like the decomposition that’s very interpretable whereas I think like the

50:24 way you’re thinking about like directly in the vector space of the models it’s more like you know uninterpretable if you just make it that simple and then I think about like is everything I do with intelligence broken down into like very simple things and I kind of think yes >> yeah you all but you could also argue you know that vectors, you know, they’re they’re these are pretty high dimensional vectors nowadays, you know, like it’s like it’s just more information. Like the more the higher the dimension, the more the information. And so, you know, there’s there’s so

50:55 many things like encoded in these vectors. It’s hard to like really track. Um, and I think in in in having in being able to encode that much information, there’s like they’re interesting insights that that you know these models find just by like interacting with features at such high high scale. Yeah, I love like TSN embedding visualizations. I just use them for random things. But >> yeah, >> I mean even just going my way back like still blows my mind. So [laughter] >> yeah, the king minus is it >> minus queen thing plus like men or

51:26 woman? Yeah. Really, really mind-blowing. >> Yeah. Yeah. Awesome. Well, uh Sedart, this has been such a cool conversation. I’ve loved learning more about this work and just all these things has really been just Yeah. >> really exciting. But um so a question I love to anchor the WVA podcast with is um uh like in addition to your work but as well as just outside work as an AI researcher um what directions for AI excite you the most? >> Yeah, I think we were just kind of like talking about this like these unsupervised notions of relevance. um like maybe to like kind of like motivate

51:57 this a little bit um you know if like we’re seeing a lot of news about this like when it comes to like proving theorems you know um like if I want to prove like some super ambitious like conjecture right or or sorry resolve some super ambitious conjecture or like prove some like unsolved like theorem or something or unproven theorem or something like that. Um and I just want to like give this to a language model like so prove X or like you know resolve Y something like that. Um and this requires like these like

52:29 really kind of these like brand new um directions of mathematics like you know you have to like maybe discover connections between like two kind of different fields. you know, you have to like like really come up with some sort of like new technology or like new like proof methods to like to solve this, you know, like basically we have to like come up with completely new knowledge like in that sense, you know, like how do we get language models to that point, you know, is it like just like a lot of like inference time, test time, search capabilities, you know, like do we have

52:59 to like really bake that into the model more so I can like find interesting things? Um, do we have to like like teach it new notions of relevance? Like like how do we get the model in some sense to be creative here? Um, this is probably this is something that’s like been I’ve been like kind of thinking about a little more over maybe like the past just month or so because um, you know, we’ve seen a lot of these like remarkable discoveries come out just just over the past few months. So I think like how do we move beyond that? Um, that see that that’s really really that seems really really exciting to me.

53:31 Um, yeah. >> Yeah. I guess like on the I I kind of splitting what you’re saying into like the sort of improvements in verifiable and non-verifiable. I I don’t know why I’ve just been thinking about it a lot that way, but like um like I’m like in terms of like you know a novel way to solve a conjecture. It’s like maybe you need like causal inference and like more neurosymbolic innovation there. I think to like >> Yeah. something interesting actually that you mentioned like verifiable versus non-verifiable is that like I guess here you really want the only like

54:02 good way to benchmark you know how well you’re doing here is you know are you discovering new things um and that reminds me of someone I was I think it was like actually at the vecdb workshop where someone was saying you know they were like working on this like this like healthbased like vector search task and like something about like proteins or something um and like they took it to this like workshop on like health at uh like ICML or ICLR this past year or something. Uh and you know basically those those guys were like yeah this is

54:33 cool you know you like do really well in all these benchmarks but like you need to discover new stuff and like in that sense it’s like okay like a real true test here of like okay if we are able to find this like like interesting notion of creativity or like these like like more unsupervised notions of relevance um like it comes down to you know are we discovering new things with it so >> yeah and I guess it’s maybe like the cost of rollouts when you talk about like going from coding to something like

55:04 proteins. I think that’s like one thing I I think like the biological machine is so complex that you can’t afford to just test it like how you can with coding and that’s >> but it’s still verifiable. It is absolutely. Yeah. Like and you know even in the case of math like we can always just like this is a bit dangerous but like you can always just like say okay here we’ve got a pre-print here. Can like the mathematicians you know verify it. Uh yeah they might know they’re not really too happy about that right now but >> yeah [laughter] I agree with the point from whatever

55:35 it’s worth. I think he makes good points. Terrence stout and yeah. Yeah. >> Awesome. Well uh Sedar thank you so much for joining the we podcast. has been such a fun conversation. I’ve learned so much. It’s been so awesome to meet you. Thanks again. >> Yeah, likewise. Thank you so much for having me on. [music] [music]

56:06 [music]