10-605 / 10-805

Machine Learning from Large Datasets

Recap of Monday

• Recap of contrastive learning and retrieval
• Some retrieval systems
• DPR and Contriever (recap)
• Combining retrieval and generation
• The original RAG paper
• REALM (earlier work)
• Fusion in Decoder (later work)
• Some recent retrieval tricks
• HyDE and LameR
• ExpandR
• Guest lecture: Michiel de Jong, Cursor

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

Performance analysis: Why MQA

Notation (one query):

• d: token encoding dim
• m: #things to maybe attend to
• h: #heads
• k: dim of queries and keys

Fast Transformer Attention: … Shazeer 2019 https://arxiv.org/pdf/1911.02150

construct the query q for x

aside: P_q = what we called W_Q

construct keys and values from M

dot prod similarities and attn scores

attn-weighted sum of values

project back to dimension of x

parallel over all heads

Performance analysis

Notation:

• b: batch size
• n: different positions in seq
• d: token encoding dim
• m: #things to maybe attend to
• h: #heads
• k: dim of queries and keys

Encoder / prefilling: batched MHA with batches of different positions in a sequence

construct the query Q for X

construct keys and values from M

dot prod similarities and attn scores

attn-weighted sum of values

project back to dimension of x

Performance analysis

Notation:

• b: batch size
• n: different positions in seq
• d: token encoding dim
• m: #things to maybe attend to
• h: #heads
• k: dim of queries and keys

Assume

• m=n
• k=v=d/h
• n < d

Encoder / prefilling: batched MHA with batches of different positions in a sequence

FLOPS: O(bnd²)

Matmul: mk,kn 🡪 mn uses mnk² FLOPS

Q: bnd,dkh 🡪 bhnk

= bnd,d(d/h)h 🡪 bhn(d/h)

= bnd,dd matmul uses bnd² FLOPS

Performance analysis

Notation:

• b: batch size
• n: different positions in seq
• d: token encoding dim
• m: #things to maybe attend to
• h: #heads
• k: dim of queries and keys

Assume

• m=n
• k=v=d/h
• n < d

Encoder / prefilling: batched MHA with batches of different positions in a sequence

Memory:

​

O(bnd) +

O(bhn²) +

O(d²)

…because

​

X,M,Q,K,V,O,Y

logits, weights

P’s

FLOPS: O(bnd²)

Performance analysis

Encoder / prefilling: batched MHA with batches of different positions in a sequence

Memory:

​

O(bnd) +

O(bhn²) +

O(d²)

…because

​

X,M,Q,K,V,O,Y

logits, weights

P’s

FLOPS: O(bnd²)

Claim: FLOPs/sec <= FLOPs/byte * bytes/sec

moving data from HBM to Matmul units is “slow”

This is nice and low!

Recap: LLM inference

• Inference with an LLM has two phases
• Prefilling / Encoding
• process the prompt and context
• for RAG and Encoder/Decoder models the context is large and structured
• recap: I said “some ways to parallelize”
• by batching inputs for different positions
• Decoding
• generate the response, word by word
• this is sequential and uses the keys-values from encoding/prefilling

Performance analysis

Notation:

• b: batch size
• n: different positions in seq
• d: token encoding dim
• m: #things to maybe attend to
• h: #heads
• k: dim of queries and keys

Decoding: extend KY cache K and V

construct the query q for x

compute new keys and values and

update key-value cache

attn-weighted sum of values

project back to dimension of x

Performance analysis

Decoding: extend KY cache K and V

construct the query q for x

compute new keys and values and

update key-value cache

attn-weighted sum of values

project back to dimension of x

More memory transfer

Less compute

This is not small if n is large … which it is for RAG/FiD

Fusion in Decoder (FiD)

FLOP analysis: Encoder is 6x as expensive as decoder!

…but at inference time decoding is slowest. Why?

Predicted by counting FLOPS for all the matmuls and assuming n_t << n_s and n_t << d

• n_t tokens in decoder output
• n_s tokens in all encoder input

Predicted by memory/FLOPs!

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

LUMEN: FiD with caching

• An attractive seeming idea
• Index every document in Wikipedia
• with DPR, Contreiver, …
• Store the BERT embeddings for every doc along with its vector index
• Given a question
• retrieve the first K doc embeddings
• use these together with the question to generate an answer

LUMEN: FiD with caching

• Store the BERT embeddings for every doc along with its vector index
• Given a question
• retrieve the first K doc embeddings
• use these together with the question to generate an answer

​

• Problem: it doesn’t work well to encode documents independently of the question for generation

question

lookup

LUMEN: FiD with caching

The first N-K layers of the FiD encoder

The last K layers of the FiD encoder

FiD decoder

Passages are encoded and stored off-line for every document

LUMEN: FiD with caching

As you scale up you can do a larger fraction of the work off-line

LUMEN: FiD with caching

Using 2/3 offline and 1/3 live is better than doing everything off-line

LUMEN: FiD with caching

As you scale up there are larger benefits from doing the off-line computation

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

GLIMMER: LUMEN with reranking

2023

Recap: Cross-Encoders

query

i-th candidate

c₁, c2, … c_N

relevance of i-th candidate to query

Recap: LUMEN=FiD with caching

Similar to cross-encoding

Except for using LUMEN’s trick of pre-computing the first few layers of the encodings independently

GLIMMER=LUMEN with reranking

Trained with “perplexity distillation loss”

log p_LM = perplexity of answer a for a decoder using only d_i, q

train to reduce KL between predicted scores and p_k – w a temperature in the softmax

GLIMMER=LUMEN with reranking

Trained with “perplexity distillation loss”

Only the top few documents are presented to the decoder

GLIMMER=LUMEN with reranking

select top 5

of 25 retrieved

GLIMMER=LUMEN with reranking

Average performance on multiple tasks from the KILT dataset

Retrieve 25 passages, rerank, and decode with top 5

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

not the “dumbest thing possible”

RETRIEVAL WITH DECODER-ONLY LLMS

Recap: Discussion of DPR

• Encoder used in DPR:
• BERT with [CLS] token
• The obvious choice in 2020
• Still a very good choice
• Decoder only models: intuitively
• Token representations don’t “know” about tokens appearing after them (causal attention)
• Token representations from late in the document aren’t used much

​

Decoder only

Discussion of DPR

• Encoder used in DPR:
• BERT with [CLS] token
• The obvious choice in 2020
• Still a very good choice
• Encoder-only and encoder-decoder models have a “bottleneck”

​

Encoder only

Decoder only

Encoder-decoder

But … everybody is working on improving decoder-only LLMs, so it would be great if we could use them!

How can you use a Decoder-Only LLM to Improve Retrieval?

2023

Hypothetical Document Embedding (HyDE)

• Details
• prompt model (InstructGPT) to convert question q to a query document d

​

​

​

• sample N documents d₁, …, d_N, by generation with temperature T
• query vector for Contriever is average of embeddings for q and d₁, …, d_N

Recap: Discussion of DPR

• DPR performs less well on queries about rare entities
• TFIDF weights rare words heavily
• DPR is trained on questions about (mostly) common entities

2022

RECAP

RECAP

Language Model as Retriever (LameR)

https://arxiv.org/pdf/2304.14233 2023 Shen et al

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

ExpandR

EMNLP 2025

ExpandR

• Start out like HYDE
• Prompted query expansion
• Followed by dense retrieval
• Then learn to improve both modules

​

​

• The retriever: fix expander + contrastive learning!

loosely interpreted

ExpandR

• Start out like HYDE
• Prompted query expansion
• Followed by dense retrieval
• Then learn to improve both modules

​

​

• The retriever: contrastive learning!

ExpandR

• Start out like HYDE
• Prompted query expansion
• Followed by dense retrieval
• Then learn to improve both modules

​

​

• The expander: an RL method called direct preference optimization (DPO)

RL: doesn’t need d^q but does need preferences d^q1 > d^q2

ExpandR: DPO Background

where:

• π* is current model, π_ref is pre-trained or SFT-trained model
• nothing actually computes “reward” r(x,y)

ExpandR: DPO Background

Note the gradient of the loss looks like this:

where:

• π* is current model, π_ref is pre-trained or SFT-trained model
• nothing actually computes “reward” r(x,y)

ExpandR

• Start out like HYDE
• Prompted query expansion
• Followed by dense retrieval
• Then learn to improve both modules

​

​

• The expander trains with^* rewards computed based on ranks of retrieved documents, based on candidate query expansions

​

*and some other tricks

loosely interpreted

ExpandR: Results

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

Decoder-only models as encoders?

2023

2025

PromptEOL

Echo embeddings

PromptEOL: Key ideas

• To summarize a sentence x
• Prompt the model with

​

• Take the last hidden state of Transformer as representation
• Then
• Use 300 sentence/word pairs as ICL demonstrations
• Train the representations contrastively with

This sentence: “x” means in one word:

Echo embeddings: key ideas

• To summarize a sentence x
• Prompt the model with

​

• Mean-pool the tokens for x as the representation

Rewrite the sentence: x; rewritten sentence: x

Outline

• More detail on Michiel’s work
• Complexity analysis for Transformers
• LUMEN
• GLIMMER
• New material
• Pipelines: decoder-only models for retrieval
• HyDE and LameR
• ExpandR
• Replacing encoder-* models
• PromptEOL
• Echo Embeddingd
• Decoder-only successors to FiD and RAG

“FID” WITH DECODER-ONLY LLMS

Main ideas in FiD and extensions

• In encoder
• cross-encoder q+p1, q+p2, … separately
• restricting cross-attention
• In decoder
• attend to everything as you generate
• Extensions
• FiDO: optimize performance to avoid decoder bottlenecks
• LUMEN: pre-compute encoder outputs
• GLIMMER: add reranking

FlashAttention (2023) and FlexAttention (2024)—also improve decoder bottlenecks

Parallel Context Windows (PCW) - 2023

TurboRAG, Blockwise Sparse Attention - 2024

Dynamic Blockwise Sparse Attention - 2025

Analog for decoder-only LLMs

2023

• Context tokens: retrieved document(s) for RAG, in-context examples, …
• Task tokens: question for RAG, test input for ICL, …
• I believe also output tokens are task tokens

​

Key idea: cross-attend within a context window, and cross-attend between task tokens and all context windows.

Very similar to FiD

• if question/answer are task tokens
• no post-training required

2023

• Context tokens: retrieved document(s) for RAG, in-context examples, …
• Task tokens: question for RAG, test input for ICL, …
• I believe also output tokens are task tokens

​

Key idea: cross-attend within a context window, and cross-attend between task tokens and all context windows.

Very similar to FiD

• if question/answer are task tokens

Good results for NQ (vs conventional RAG system) and ICL (especially classification with many classes)

Implemented with attention masks rather than parallel generation of keys and values

2024

• Context tokens: retrieved document(s) for RAG, in-context examples, …
• Task tokens: question for RAG, test input for ICL, …

​

Key idea: same as PCW except

• cache the independently-produced KV pairs of the documents
• load KV pairs for relevant docs directly
• fine-tune for a short time (100-1000 steps)

TTFT = Time To First Token

FLOPS also for first token

2025

Key idea: same as Block-Attention except

• experiments are on ICL instead of RAG
• cache the independently-produced KV pairs of the documents
• evaluate addition of KV pairs computed for retrieved documents
• retrieve ICL examples with BM25
• avoid fine-tuning
• clever use of StreamingLLM trick of “attention sink”
• don’t mess with position encodings

DBSA: Details

Baseline: many-shot learning

• Given n ICL examples
• encode n/k blocks of ICL examples containing
• k demonstrations
• a shared “anchor block”
• load all blocks into a KV cache
• add the usual positional encoding features
• Implemented with FlexAttention
• Doesn’t actually encode the masked blocks

​

• 50 examples/block, 30k and 90k training examples for classification tasks

​

​

DBSA: Details

Dynamic example selection

• Given a query q*
• retrieve m<n encoded blocks (with BM25)
• load retrieved blocks and anchor block into a KV cache and add positional encoding

​

• retrieve 30% of the available training examples

​

DBSA: Results

Total latency including set-up time

DBSA: Results

TurboRAG EMNLP 2025

Similar plan as DBSA

• retrieved KVs for chunks of relevant text
• no “attention sink” tricks, instead FT more for post-retrieval generation
• dense retrieval, not BM25

TurboRAG

KV Retrieval vs KV Cache Eviction

• KV Cache Eviction:
• get the best (more useful, smallest) KV cache by starting with a big one and making it smaller with evictions
• PCW, DBSA, TurboRAG
• get the best (more useful, smallest) KV cache by starting with a small one and making it larger with retrievals
• some cross-attention never computed

​