Section 4
Retrieval-based LMs: Training:1

Retrieval-based LMs
Query

Input

LM

Index

OutputDatastore2

Training retrieval-based LMs
Query

Input

LM
Backpropagate

Index
Back-propagate

OutputDatastore3

Why is training challenging?
Query

Input

LM
Backpropagate

Index
Back-propagate⚠
Training LMs can
Output be very expensive!

Datastore4

Why is training challenging?
Query

Input

LM
Backpropagate

Index

⚠
Datastore

Back-propagate

Too large! Expensive to update
index during training!

⚠
Training LMs can
Output be very expensive!

5

Challenges of updating retrieval models
x1
x2

x3

Encoder

Index…

Datastore

We may encode a lot of (>100M)
text chunks using the encoder!

6

Challenges of updating retrieval models
x1
x2

x3

Encoder

Index…

Datastore

During training, we will update the encoder

7

Challenges of updating retrieval models
x1
x2

x3

Encoder

Index…

Datastore

Re-indexing will be very expensive!

8

Why is training challenging?
Query

Input

LM
Backpropagate

Index

⚠
Datastore

Back-propagate

Too large! Expensive to update
index during training!

⚠
Training LMs can
Output be very expensive!

9

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation

10

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation

11

Independent training
Retrieval models and language models are trained independently- Training language models
Input

LM

Output- Training retrieval models
Retriever
Datastore

Chunks/tokens
Query
12

Independent training
Retrieval models and language models are trained independently- Training language models
Input

LM

Output- Training retrieval models
Retriever
Datastore

Chunks/tokens
Query
13

Training language models
Input

LM

OutputBack-propagate

Minimize −log PLM(y | x)14

Training language models
Input

LM

OutputBack-propagate

Minimize −log PLM(y | x)
……
GPT

PaLM

LLaMA

GPT-J

15

Independent training
Retrieval models and language models are trained independently- Training language models
Input

LM

Output- Training retrieval models
Retriever
Datastore

Chunks/tokens
Query
16

Sparse retrieval models: TF-IDF / BM25
In 1997, Apple merged with NeXT,
and Steve Jobs became CEO of …

[0, 0, 0.4, 0, 0.8, 0.7, …]

Lexical overlap

Jobs returned to Apple as CEO
after the company's acquisition …

[0, 1.2, 0.4, 0, 0.8, 0, …]

Text chunks

Sparse vectors

No training needed!
Ramos, 2003. “Using TF-IDF to Determine Word Relevance in Document Queries”
Robertson and Zaragoza, 2009. “The Probabilistic Relevance Framework: BM25 and Beyond”

17Dense retrieval models: DPR (Karpukhin et al. 2020)
Inner Product Similarity

Dense vectors

Encoder

Encoderq
Query

Text chunks

Karpukhin et al., 2020. “Dense Passage Retrieval for Open-Domain Question Answering”

18Dense retrievers: Inference

Encoder
CorpusIndex
Maximum inner-product search

Query

Encoder19

Dense retrievers: Inference

Encoder
CorpusIndex
Maximum inner-product search

Query

Encoder

How to train dense retrieval models?
20

Training dense retrieval models: DPR
Inner Product Similarity

Encoder

Encoderq
Query

Text chunks

21

Training dense retrieval models: DPR
Inner Product Similarity

L(q, pEncoder

Encoder= − log+− −−, p1 , p2 , …, pn )+exp(sim(q, p ))n+exp(sim(q, p )) + ∑j=1 exp(sim(q, pj−))

q
Query

Text chunks

22

Training dense retrieval models: DPR
Inner Product Similarity

L(q, pEncoder

Encoder− −−, p1 , p2 , …, pn )+exp(sim(q, p ))n+exp(sim(q, p )) + ∑j=1 exp(sim(q, pj−))

Contrastive learning

q
Query

= − log+

Text chunks

23

Training dense retrieval models: DPR
Inner Product Similarity

L(q, pEncoder

Encoder− −−, p1 , p2 , …, pn )+exp(sim(q, p ))n+exp(sim(q, p )) + ∑j=1 exp(sim(q, pj−))

Contrastive learning

q
Query

= − log+

Positive

Positive

Text chunks
Negative

Negative24

Training dense retrieval models: DPR
Inner Product Similarity

L(q, p+− −−, p1 , p2 , …, pn )Positive passage

Encoder

Encoder= − log+exp(sim(q, p ))n+exp(sim(q, p )) + ∑j=1 exp(sim(q, pj−))

q
Query

Text chunks

25Training dense retrieval models: DPR
Inner Product Similarity

Negative passages

Too expensive to consider all negatives!
+ − −−L(q, p , p1 , p2 , …, pn )Positive passage

Encoder

Encoder= − log+exp(sim(q, p ))n+exp(sim(q, p )) + ∑j=1 exp(sim(q, pj−))

q
Query

Text chunks

26

fi

Training with “in-batch” negatives
Training batchL(q, p= − log

q1

+

− −
−
, p1 , p2 , …, pn )Who foundedApple?

+

exp(sim(q, p ))nexp(sim(q, p +)) + ∑j=1 exp(sim(q, pj−))q2

What is thename of
Spain's most
famous soccer
team?
…

qn

Who was the
rst ministryhead of state in
Nigeria?

p1+…It was incorporated
by Jobs and Wozniak
as Apple Computer,Inc. in 1977. …

p2+

+
p2

12-year-old Spanish
football club Real
Madrid is undoubtedly
the best football club
Spain has ever…

pn+…

Thomas Umunnakwe
Aguiyi-Ironsi seized
power during theensuing chaos after
the 15 January …

27

fi

Training with “in-batch” negatives
Training batchL(q, p= − logp1+
q1
positive …It was incorporated

+

− −
−
, p1 , p2 , …, pn )Who foundedApple?

+

exp(sim(q, p ))nexp(sim(q, p +)) + ∑j=1 exp(sim(q, pj−))q2

What is thename of
Spain's most
famous soccer
team?
…

qn

Who was the
rst ministryhead of state in
Nigeria?

by Jobs and Wozniak
as Apple Computer,Inc. in 1977. …

p2+

+
p2

12-year-old Spanish
football club Real
Madrid is undoubtedly
the best football club
Spain has ever…

pn+…

Thomas Umunnakwe
Aguiyi-Ironsi seized
power during theensuing chaos after
the 15 January …

28

fi

Training with “in-batch” negatives
Training batchL(q, p= − logp1+
q1
positive …It was incorporated

+

− −
−
, p1 , p2 , …, pn )Who foundedApple?

+

exp(sim(q, p ))nexp(sim(q, p +)) + ∑j=1 exp(sim(q, pj−))

by Jobs and Wozniak
as Apple Computer,Inc. in 1977. …

q2

What is thename of
Spain's most
famous soccer
team?
…

p2+

negativesqn

Who was the
rst ministryhead of state in
Nigeria?

+p2

12-year-old Spanish
football club Real
Madrid is undoubtedly
the best football club
Spain has ever…

pn+…

Thomas Umunnakwe
Aguiyi-Ironsi seized
power during theensuing chaos after
the 15 January …

Back-propagation to all in-batch negatives!

29

Contriever (Izacard et al. 2022)
Independent CroppingApple merged with NeXT
In 1997, Apple merged with NeXT,
and Steve Jobs became CEO of
his former company. He became
the saviour of his company andwas largely responsible …

Positives

Steve Jobs became CEO of

Unsupervised dense retrieval model!
Izacard et al., 2022. “Unsupervised Dense Information Retrieval with Contrastive Learning”

30Retrieval-in-context in LM (Ram et al. 2023)
x = World Cup 2022 was the last with 32 teams, before the increase to
World Cup 2022 was the last with 32 teams, before the increase to

Retrieval Model
FIFA World Cup 2026 will expand to 48 teams. World Cup 2022 was the last with 32 teams, before the increase toLM
48 in the 2026 tournament.

Ram et al. 2023. “In-Context Retrieval-Augmented Language Models”

31Retrieval-in-context in LM (Ram et al. 2023)
x = World Cup 2022 was the last with 32 teams, before the increase to
World Cup 2022 was the last with 32 teams, before the increase to

Retrieval Model

BM25, DPR, Contriever, …

FIFA World Cup 2026 will expand to 48 teams. World Cup 2022 was the last with 32 teams, before the increase toLM

GPT, OPT, LLaMA, …

48 in the 2026 tournament.

Ram et al. 2023. “In-Context Retrieval-Augmented Language Models”

32Retrieval-in-context in LM

Better retrieval model
Better base LMs

Better retrieval-based LMs

33

Retrieval-in-context in LM

Better retrieval model
Better base LMs

Better retrieval-based LMs

Each component can be improved separately
34

kNN-LM (Khandelwal et al. 2020)
Inference
PkNN(y | x) ∝∑[v = y]exp(−d(Enc(k), Enc(x))(k,v)∈

𝒟

PkNN−LM(y | x) = (1 − λ)PLM(y | x) + λPkNN(y | x)

𝕀

Khandelwal et al., 2020. “Generalization through Memorization: Nearest Neighbor Language Models”

35

kNN-LM (Khandelwal et al. 2020)
Inference
Re-use the LM encoder. No training needed!

PkNN(y | x) ∝∑[v = y]exp(−d(Enc(k), Enc(x))(k,v)∈

𝒟

PkNN−LM(y | x) = (1 − λ)PLM(y | x) + λPkNN(y | x)

𝕀

Khandelwal et al., 2020. “Generalization through Memorization: Nearest Neighbor Language Models”

36

kNN-LM (Khandelwal et al. 2020)
Inference
Re-use the LM encoder. No training needed!

PkNN(y | x) ∝∑[v = y]exp(−d(Enc(k), Enc(x))(k,v)∈

PkNN−LM(y | x) = (1 − λ)PLM(y | x) + λPkNN(y | x)

Training

𝒟

Minimize −log PLM(y | x)

𝕀

Khandelwal et al., 2020. “Generalization through Memorization: Nearest Neighbor Language Models”

37Independent training
Work with o -the-shelf models (no extra training required)
Each part can be improved independently

ff

38

Independent training
Work with o -the-shelf models (no extra training required)
Each part can be improved independently
LMs are not trained to leverage retrieval
Retrieval models are not optimized for LM tasks/domainsff39

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation40

Sequential training- One component is rst trained independently and then xed- The other component is trained with an objective that depends on the rst one

fi

fi

fi

41

Sequential training- One component is rst trained independently and then xed- The other component is trained with an objective that depends on the rst one

Retriever

LM

fi

fi

fi

42

Sequential training- One component is rst trained independently and then xed- The other component is trained with an objective that depends on the rst one

RetrieverLM

Retriever

LM

fi

fi

fi

43

Sequential training- One component is rst trained independently and then xed- The other component is trained with an objective that depends on the rst one

RetrieverLM

Retriever

LM

fi

fi

fi

44

Sequential training- Retrieval models are rst trained independently and then xed- Language models are trained with an objective that depends on the retrieval
Input
Retrieval model(pre-trained)

LM

Backpropagate

Index
Stop-gradient

DatastoreOutputfi

fi

45

RETRO (Borgeaud et al. 2021)
x = World Cup 2022 was the last with 32 teams, before the increase tox1

x2

Borgeaud et al., 2021. “Improving language models by retrieving from trillions of tokens”

x3

46

RETRO (Borgeaud et al. 2021)
x = World Cup 2022 was the last with 32 teams, before the increase tox11kp1 . . . p1

x1
x2

x3

x3

x2

Retrieval
Encoder

1
k
p
.
.
.
p
Index 22

1
kp3 . . . p3

47

RETRO (Borgeaud et al. 2021)
x = World Cup 2022 was the last with 32 teams, before the increase tox1

x3(k chunks of text per split)1k
p1 . . . p1

x1
x2

x3

x2

Retrieval
EncoderIndex

p12 . . . pk2
1kp3 . . . p3

E1
LM
Encoder

E2

E3

48

RETRO (Borgeaud et al. 2021)
x = World Cup 2022 was the last with 32 teams, before the increase to

x1

x3

x2

(k chunks of text per split)1k
p1 . . . p1

x1
Retrieval
Encoderx2

Index

p12 . . . pk2

E1
LM
Encoder1kp3 . . . p3

x3

E2

E3

E1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)
49

RETRO (Borgeaud et al. 2021)
x = World Cup 2022 was the last with 32 teams, before the increase tox1

x2

Very large!
600B tokens during training;
1.8T tokens during inference.
1k
p1 . . . p1

x1
Retrieval
Encoderx2

Index

p12 . . . pk2

x3(k chunks of text per split)

E1
LM
Encoder1kp3 . . . p3

x3

E2

E3

E1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)
50

RETRO: Training
p11 . . . pk1

x1
Retrieval
Encoder

x2

Index

1
k
p2 . . . p2

E1
LM
Encoder

p13 . . . pk3

x3

E2

E3

E1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)51

RETRO: Training
p11 . . . pk1

x1
Retrieval
Encoder

x2

Index

1
k
p2 . . . p2

E1
LM
Encoder

p13 . . . pk3

x3

E2

E3

E1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)

Back-propagate
52

RETRO: Training
p11 . . . pk1

x1
Retrievalx2
Encoder
Updating an
index with 600B is
xextremely expensive!!

Index

3

⚠

1
k
p2 . . . p2

E1
LM
Encoder

p13 . . . pk3

E2

E3

E1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)

Back-propagate
53

RETRO: Training1kp1 . . . p1

x1
Retrievalx2
Fix the retrieval encoderEncoderandthe index duringxtraining!
3

Index

1
k
p2 . . . p2

E1
LM
Encoder1kp3 . . . p3

E2

E3

Stop-gradientE1 E2 E3x1
x2 EMB

x3

ATTN

CCA

FFW

HEAD

RETRO blocks (xL)

Back-propagate
54Sequential training- One component is rst trained independently and then xed- The other component is trained with an objective that depends on the rst one

RetrieverLM

Retriever

LM

fi

fi

fi

55

Sequential training- Language models are rst trained independently and then xed- Retrieval models are trained/ ne-tuned with supervisions from LMs
Input
Retrieval modelLM
Stop-gradient

Index
Back-propagate

DatastoreOutput

fi

fi

fi

56

REPLUG (Shi et al. 2023)

Shi et al., 2023. “REPLUG: Retrieval-Augmented Black-Box Language Models”

57

REPLUG (Shi et al. 2023)

Shi et al., 2023. “REPLUG: Retrieval-Augmented Black-Box Language Models”

58

REPLUG (Shi et al. 2023)

Shi et al., 2023. “REPLUG: Retrieval-Augmented Black-Box Language Models”

59

REPLUG LSR (LM-Supervised Retrieval)

60

REPLUG LSR (LM-Supervised Retrieval)

61

REPLUG LSR (LM-Supervised Retrieval)

62

REPLUG LSR (LM-Supervised Retrieval)Updating retrieval encoder

Retrieval Index becomes “stale”

How to deal with this issue? We will talk about it soon!
“Asynchronous update”

63

REPLUG results
Bits per byte (BPB): The lower the better

64

REPLUG results
With Contriever, “independent training”

65

REPLUG results
Fine-tuning Contriever with
LM-supervised training“Sequential training”

66Sequential training
Work with o -the-shelf components (either a large index or a powerful LM)
LMs are trained to e ectively leverage retrieval results
Retrievers are trained to provide text that helps LMs the mostOne component is still xed and not trained

fi

ff

ff

67

Sequential training
Work with o -the-shelf components (either a large index or a powerful LM)
LMs are trained to e ectively leverage retrieval results
Retrievers are trained to provide text that helps LMs the mostOne component is still xed and not trained

Let’s jointly train retrieval models and LMs!

fi

ff

ff

68

Q &A

69

We’ll be back at 4:00pm!

Section 4
Retrieval-based LMs: Training (cont’d):71Why is training challenging?
Query

Input

LM
Backpropagate

Index

⚠
Datastore

Back-propagate

Too large! Expensive to update
index during training!

⚠
Training LMs can
Output be very expensive!

72

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation73

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation74

Challenges of updating retrieval models
x1
x2

x3

Encoder

Index…

Datastore

During training, we will update the encoder75

Challenges of updating retrieval models
x1
x2

x3

Encoder

Index…

Datastore

Re-indexing will be very expensive!

76

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation

77

Joint training w/ asynchronous index update- Retrieval models and language models are trained jointly- Allow the index to be “stale”; rebuild the retrieval index every T steps
Input

LM

Backpropagate

Index

Datastore

Back-propagate w/
asynchronous update

Output78

Asynchronous index update
x1
x2

x3

Encoder

Index…

Datastore

79

Asynchronous index update
x1
x2

x3

UpdatedEncoder

Stale
Index

…

Datastore

80Asynchronous index update
x1
x2

x3

UpdatedEncoder

Updated
Index…

Datastore

T updates

Refresh81

REALM (Guu et al. 2020)
x = The [MASK] at the top of the pyramid.
q (=x)The pyramidion on top … the pyramid.
…
The [MASK] at the top of the pyramid.

Index

The pyramidion on top
allows for less material
higher up the pyramid.

LMpyramidP(z | x)Guu et al., 2020. “REALM: Retrieval-Augmented Language Model Pre-Training”

P(y | x, z)

82

𝒵

REALM: Training
Objective: maximize∑z∈q (=x)Indexθ : top-K retrieved chunks

The pyramidion on top
allows for less material
higher up the pyramid.

Pθ(z | q)Pθ(y | q, z)θThe pyramidion on top … the pyramid.
…
The [MASK] at the top of the pyramid.

LM𝒵pyramidPθ(z | x)Pθ(y | x, z)

83

𝒵

REALM: Training
Objective: maximize∑z∈q (=x)Indexθ : top-K retrieved chunks

The pyramidion on top
allows for less material
higher up the pyramid.

Pθ(z | q)Pθ(y | q, z)θThe pyramidion on top … the pyramid.
…
The [MASK] at the top of the pyramid.
Back-propagationLM𝒵

pyramidPθ(z | x)Pθ(y | x, z)

84

𝒵

REALM: Training
Objective: maximize∑z∈q (=x)
Stale index;
Update every T steps

IndexPθ(z | q)Pθ(y | q, z)θThe pyramidion on top … the pyramid.
…
The [MASK] at the top of the pyramid.

θ : top-K retrieved chunks

LM

The pyramidion on top
allows for less material
higher up the pyramid.

𝒵

pyramidPθnew(z | x)

Up-to-date parametersPθnew(y | x, z)

85

REALM: Index update rate
How often should we update the retrieval index?

- Frequency too high: expensive- Frequency too slow: out-dated

86

REALM: Index update rate
How often should we update the retrieval index?

- Frequency too high: expensive- Frequency too slow: out-dated

REALM: updating the index every 500 training steps

87

REALM: Index update rate
How often should we update the retrieval index?

- Frequency too high: expensive- Frequency too slow: out-dated

REALM: updating the index every 500 training steps
40
35

EM score on NQ

38.20

30

28.70

25
20

REALM

30× slower update88

Atlas (Izacard et al. 2022)

Izacard et al., 2022. “Atlas: Few-shot Learning with Retrieval Augmented Language Models”

89Atlas (Izacard et al. 2022)
Retrieval-based encoder-decoder model

Jobs is CEO of __Jobs was raised … Jobs is CEO of

Index

Steve Jobs passed … Jobs is CEO of

Encoder

Decoder

Jobs cofounded … Jobs is CEO ofApple

Izacard and Grave., 2020. “Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering”

90Atlas (Izacard et al. 2022)
Jobs is CEO of __Jobs was raised … Jobs is CEO of

Index

Steve Jobs passed … Jobs is CEO of

Encoder

Decoder

Jobs cofounded … Jobs is CEO of

Retrieve docs & Process each doc independently using “Fusion-in-Decoder” Apple

Izacard and Grave., 2020. “Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering”

91Atlas (Izacard et al. 2022)

Adapted to a lot of downstream tasks! (Section 5)
92

Atlas: Retriever training
Perplexity Distillation
Retrieve the text that can help LM encoders improve perplexity

Pretr(z | q) =exp(s(z, q))K∑k=1 exp(s(zk, q))

How likely each document is retrievedPppl(z | q, y) =exp(logPLM(y | q, z))K∑k=1 exp(logPLM(y | q, zk))

How much each document improves the ppl

93

Atlas: Retriever training
Similarity based on
retrieval encoder

Pretr(z | q) =Prob of the gold labels if
augmenting this text chunk

KL Divergenceexp(s(z, q))K∑k=1 exp(s(zk, q))

How likely each document is retrievedPppl(z | q, y) =exp(logPLM(y | q, z))K∑k=1 exp(logPLM(y | q, zk))

How much each document improves the ppl

Perplexity Distillation

94

Atlas: Asynchronous index update
Input

LM
(Every T steps)

Backpropagate

Index
Back-propagate w/
asynchronous update

DatastoreOutputUpdate the index every T step
30% overhead for asynchronous update on Wikipedias95

Atlas: Asynchronous index update
Input

LM
(Every T steps)

Backpropagate

Index
Back-propagate w/
asynchronous update

DatastoreOutputUpdate the index every T step
30% overhead for asynchronous update on Wikipedias

How can we get rid of this?

96

Training methods for retrieval-based LMs• Independent training• Sequential training• Joint training w/ asynchronous index update• Joint training w/ in-batch approximation

97

Joint training w/ in-batch approximation- Retrieval models and language models are trained jointly- Use “in-batch index” instead of full index
Input

LM

Backpropagate

In-batch approximation
Back-propagate

Training batch

Output98

In-batch approximation
Encoder

Full index

Full corpus

99In-batch approximation>100M

Encoder

Full corpus

Full index

Re-indexing will be very expensive!

100In-batch approximation>100M

Encoder

Full corpus

Full index

Re-indexing will be very expensive!
~10K

EncoderTraining batch

In-batch index
Computed on the y for each batch!

fl

101

TRIME: Training with in-batch memory (Zhong et al. 2022)
Similar to kNN-LM

Datastore= {(k, v)}

𝒟context (chunk) next token

Zhong et al., 2022. “Training Language Models with Memory Augmentation”

Keys

Values

10/10, would buy this cheap
Item delivered broken. Very cheap
To check the version of PyTorch, you can use torch
You are permitted to bring a torch
A group of infections … one of the torch

102

TRIME: Training with in-batch memory (Zhong et al. 2022)
Similar to kNN-LM

Datastore= {(k, v)}

KeysValues

10/10, would buy this cheap
Item delivered broken. Very cheap
To check the version of PyTorch, you can use torch
You are permitted to bring a torch
A group of infections … one of the torch

context (chunk) next token

Inference⊤P(y ∣ x) ∝ exp(E f(x)) +output embedding(same as standard LMs)∑[v = y]exp(−d(Enc(k), Enc(x)))(k,v)∈

datastore𝕀

𝒟

𝒟(very large!)
103

TRIME: Training with in-batch memory (Zhong et al. 2022)
Similar to kNN-LM

Datastore= {(k, v)}

KeysValues

10/10, would buy this cheap
Item delivered broken. Very cheap
To check the version of PyTorch, you can use torch
You are permitted to bring a torch
A group of infections … one of the torch

context (chunk) next token

Inference⊤P(y ∣ x) ∝ exp(E f(x)) +output embedding(same as standard LMs)(k,v)∈

datastore

2. Aligning input chunks with all the chunks indatastore that share the same next token(very large!)

𝕀

𝒟

𝒟

1. Aligning the output representations
with static embeddings∑[v = y]exp(−d(Enc(k), Enc(x)))

104

TRIME: Training with in-batch memory (Zhong et al. 2022)⊤

P(y ∣ x) ∝ exp(E f(x)) +∑[v = y]exp(−d(Enc(k), Enc(x)))(k,v)∈

Jobs become CEO of AppleHe moves to Apple

Positive chunks → pull together
Jobs become CEO of AppleShe works at Microsoft

𝕀

𝒟

Negative chunks → push away

105

TRIME: Training with in-batch memory (Zhong et al. 2022)⊤

P(y ∣ x) ∝ exp(E f(x)) +∑(k,v)∈

Jobs become CEO of Apple[v = y]exp(−d(Enc(k), Enc(x)))
Very large!
He moves to Apple

Positive chunks → pull together
Jobs become CEO of AppleShe works at Microsoft

𝕀

𝒟

Negative chunks → push away

106

TRIME: Training
Key idea: build a temporary index from same training batch on the y⊤P(y ∣ x) ∝ exp(E f(x)) +∑(k,v)∈[v = y](−d(Enc(k), Enc(x)))
train

In-batch approximation𝒟(built from in-batch examples on the y)

fl

fl

𝕀

107

TRIME: Training
Key idea: build a temporary index from same training batch on the y⊤P(y ∣ x) ∝ exp(E f(x)) +∑(k,v)∈[v = y](−d(Enc(k), Enc(x)))
train

In-batch approximation(built from in-batch examples on the y)train!

𝒟

We can back-propagate to all the representations in datastore

fl

fl

𝕀

𝒟

108

TRIME: Full index vs. in-batch index
Keys

ValuesTo check the version of PyTorch, you can use torch
Item delivered broken. Very cheap
He moves to Apple
Apple merged with NeXT, and Jobs became CEO……

Full corpus

Full index (used during inference)109

TRIME: Full index vs. in-batch index
>100M
Keys

ValuesTo check the version of PyTorch, you can use torch
Item delivered broken. Very cheap
He moves to Apple
Apple merged with NeXT, and Jobs became CEO……

Full corpus

Full index (used during inference)
Compute on the y!

Apple merged with NeXT, and …
VS Code was developed byMicrosoft for Windows in 2015 …
He moves to Apple …
…

Training batch

Keys

Values~10K

Jobs becameApple merged with NeXT, and Jobs became CEO……
He moves to AppleVS Code was developed By

In-batch index (used during training)

fl

110

TRIME: Full index vs. in-batch index

How to batch training data —so we can have good in-batch examples?
Compute on the y!
Apple merged with NeXT, and …
VS Code was developed byMicrosoft for Windows in 2015 …
He moves to Apple …
…

Training batch

Keys

ValuesJobs becameApple merged with NeXT, and Jobs became CEO……
He moves to AppleVS Code was developed By

In-batch index (used during training)

 

fl

111

TRIME: Data batching strategy
Current token

In memory

Not in memory

Key idea: similar text chunks— more training signals fromin-batch examples!

Use BM25 scores to nd
similar text chunks

fi

112

TRIME: Results
Perplexity: The lower the better
19

18.65
18

17

16.37

15

15.41

14

Transformer

kNN-LM

TrimeLM

Perplexity on Wikitext-103

113

NPM: Nonparametric masked LMs (Min et al. 2023)

Min et al., 2023. “Nonparametric Masked Language Modeling”

114

NPM: Nonparametric masked LMs (Min et al. 2023)

1. masked language model pretained on >1B tokens
115NPM: Nonparametric masked LMs (Min et al. 2023)

2. Each mask corresponds to a phrase(instead of a token)

1. masked language model pretained on >1B tokens

116NPM: Nonparametric masked LMs (Min et al. 2023)
3. During inference, predictions are madepurely according to retrieval results

2. Each mask corresponds to a phrase(instead of a token)

1. masked language model pertained on >1B tokens

117NPM: Nonparametric masked LMs (Min et al. 2023)
Key challenges
1. How to approximate the full retrieval index during training
2. How to get training signals (positive/negatives) from the index approximation

In-batch approximation with same-doc batching118

NPM: Training
1. Sample sequences from the same document
In the 2010 NFL season, the Seattle Seahawks made history bymaking it into the playoffs despite having a 7–9 record.

… against the Seattle Seahawks as a member of (…) In the 2010season, the Seahawks became the irst team in NFL history to …

For simplicity, we assume 2 sequences in a batchf119NPM: Training
2. Identify co-occurring spans
In the 2010 NFL season, the Seattle Seahawks made history bymaking it into the playoffs despite having a 7–9 record.

… against the Seattle Seahawks as a member of (…) In the 2010season, the Seahawks became the irst team in NFL history to …

For simplicity, we assume 2 sequences in a batchf120

NPM: Training
3. One is “positive” for the otherIn the 2010 NFL season, _____________ made history by making it
into the playoffs despite having a 7–9 record.

positive… against the Seattle Seahawks as a member of (…) In the 2010season, the Seahawks became the irst team in NFL history to …

For simplicity, we assume 2 sequences in a batchf121

NPM: Training
4. The others are “negatives”
In the 2010 NFL season, _____________ made history by making it
into the playoffs despite having a 7–9 record.

positive… against the Seattle Seahawks as a member of (…) In the 2010season, the Seahawks became the irst team in NFL history to …

negatives

For simplicity, we assume 2 sequences in a batchf122

Beyond lexical clues?
In TRIME and NPM, retrieval models are trained to use lexical information

Positives: co-occurring tokens/spans

Can we do more than that?

123RPT: Retrieval-pretrained transformer(Rubin and Berant 2023)

Reference score

P(“Apple” | "Jobs become CEO of”, “NeXT merged with …”)
Reference chunk124
 Rubin and Berant, 2023. “Long-range Language Modeling with Self-retrieval”

RPT: Retrieval-pretrained transformer(Rubin and Berant 2023)

Reference score

P(“Apple” | "Jobs become CEO of”, “NeXT merged with …”)
Reference chunk

P(“Apple” | "Jobs become CEO of”, “He joined his former …”) > Reference score
Positive chunks125
 Rubin and Berant, 2023. “Long-range Language Modeling with Self-retrieval”

RPT: Retrieval-pretrained transformer(Rubin and Berant 2023)

Reference score

P(“Apple” | "Jobs become CEO of”, “NeXT merged with …”)
Reference chunk

P(“Apple” | "Jobs become CEO of”, “He joined his former …”) > Reference score
Positive chunks
P(“Apple” | "Jobs become CEO of”, “Jobs was raised …”) < Reference score
Negative chunks126
 Rubin and Berant, 2023. “Long-range Language Modeling with Self-retrieval”

Joint training
End-to-end trained — each component is optimized
Good performance
Training is more complicated(async update, overhead, data batching, etc)
Train-test discrepancy still remains

127Summary
Training method
Independent training(Ram et al 2023; Khandelwal et al 2020)

Sequential training(Borgeaud et al 2021; Shi et al 2023)✴ Easy to implement: o -the-shelf✴ Models are not end-to-end trainedmodels✴ Easy to improve: sub-module— suboptimal performancecan be separately improved

Joint training: async update(Guu et al 2020; Izacard et al 2022)

Joint training: in-batch approx(Zhong et al 2022; Min et al 2023; Rubin andBerant 2023)✴ Training may be complicated✴ End-to-end trained — very good(overhead, batching methods, etc)
performance!
✴ Train-test discrepancy still remains

How do retrieval-based language models
perform on downstream tasks? → Section 5!
ff

128