Contextual Memory Trees
Wen Sun
CMU —> MSR NYC
wensun@cs.cmu.edu
Joint work with Alina Beygelzimer, Hal Daumé III, Paul Mineiro, and John Langford

External Memory

External Memory(many query-answer pairs)

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

Answer(A Chinese sentence)

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

K memories:(K English-Chinese pairs)
Answer(A Chinese sentence)

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

K memories:(K English-Chinese pairs)
Answer(A Chinese sentence)
An external
Inference/learning
algorithm

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

K memories:(K English-Chinese pairs)
Answer(A Chinese sentence)

Answer(A Chinese Sentence)

An external
Inference/learning
algorithm

External Memory
Query (e.g., An English sentence)

External Memory(many query-answer pairs)

K memories:(K English-Chinese pairs)
Answer(A Chinese sentence)

Answer(A Chinese Sentence)
User provides
rewards (BLUE score) &ground-truth (A Chinese sentence)

An external
Inference/learning
algorithm

External Memory
Query (e.g., An English sentence)
Insert &Reinforcement Learn
External Memory(many query-answer pairs)

K memories:(K English-Chinese pairs)
Answer(A Chinese sentence)

Answer(A Chinese Sentence)
User provides
rewards (BLUE score) &ground-truth (A Chinese sentence)

An external
Inference/learning
algorithm

Desired Properties

Desired Properties
Online: every operation works one example at a timeDesired Properties
Online: every operation works one example at a time
Linear Space: O(# of examples)Desired Properties
Online: every operation works one example at a time
Linear Space: O(# of examples)
Fast: Logarithmic Read & Write (Log(# of examples))Desired Properties
Online: every operation works one example at a time
Linear Space: O(# of examples)
Fast: Logarithmic Read & Write (Log(# of examples))Learning-based: do not assume, e.g., Euclidean space

Desired Properties
Online: every operation works one example at a time
Linear Space: O(# of examples)
Fast: Logarithmic Read & Write (Log(# of examples))Learning-based: do not assume, e.g., Euclidean space
Self-consistency: identify the item seen before

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)(m1 , . . . , mk )

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)(m1 , . . . , mk )2. If reward ri is given,

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)(m1 , . . . , mk )2. If reward ri is given,

UPDATE( q, mi , ri )

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)(m1 , . . . , mk )2. If reward ri is given,

UPDATE( q, mi , ri )3. If ground-truth
value v is given,

Learning Protocol ofContextual Memory Tree (CMT)
A memory (m): a pair of Query (q) & Value (v)e.g., (A English Sentence, A Chinese Sentence)

1. Given a query q,

QUERY(q)(m1 , . . . , mk )2. If reward ri is given,

UPDATE( q, mi , ri )3. If ground-truth
value v is given,

INSERT( q, v )Datastructure of CMT:A Nearly Balanced Binary Tree

f1
f2

f
g

m1 , . . . mn

Datastructure of CMT:A Nearly Balanced Binary Tree
Internal node

f1

nL , nR# of memories inthe Left/Right
subtreeRouter (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Datastructure of CMT:A Nearly Balanced Binary Tree
Internal node

f1

nL , nR# of memories inthe Left/Right
subtreeRouter (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1

nL , nR# of memories inthe Left/Right
subtreeRouter (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1 (q) < 0

f1

nL , nR# of memories inthe Left/Right
subtreeRouter (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1 (q) < 0

f1

nL , nR# of memories inthe Left/Right
subtreeRouter (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1 (q) < 0

f1

nL , nR# of memories inthe Left/Right
subtreef2 (q) > 0

Router (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1 (q) < 0

f1

nL , nR# of memories inthe Left/Right
subtreef2 (q) > 0

Router (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mn

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Datastructure of CMT:A Nearly Balanced Binary Treeq
Internal node

f1 (q) < 0

f1

nL , nR# of memories inthe Left/Right
subtreef2 (q) > 0

Router (binary classifier):f : Q ! { 1, +1}

e.g.,f (q) = sign( > (q))

f2

f
g

m1 , . . . mnm⇤ = arg max g(q, m)m
Select the most relevant memory

Leaf

Memories:
m1 , . . . mn
Scorer Function:g(q, x) = wT (q, x)

Update Routers:
Using reward signals to update routers
f1
f2

f

Update Routers:
Using reward signals to update routers
q
f1
f2

f

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

f1
f

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

f1
f

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

f1
f

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

rL

f1
f

rR

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

f

Binary Classification{q, sign(rR

rL )}
rL

f1

rR

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2

f

Binary Classification{q, sign(rR

rL )}
rL

f1

rR

✏ -Greedy for exploration

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2
Binary Classification{q, sign(rR

f
p

rL )}
rL

f1

rR

✏ -Greedy for exploration

Update Routers:
Using reward signals to update routersq
f1 (q) < 0

f2
Binary Classification{q, sign(rR

f1
f

p

rL )}
rL

rR

✏ -Greedy for exploration +Importance Weight (Propensity Score)Update Routers:
Using reward signals to update routersq
f1

f1 (q) < 0

f2
Binary Classification{q, sign(rR

f
p

rL )}
rL

rR

✏ -Greedy for exploration +Importance Weight (Propensity Score)
rR
⇡ rR
p

rL

Update Routers:
Using reward signals to update routersq
f1

f1 (q) < 0

f2
Binary Classification{q, sign(rR

f
p

rL )}
rL

rR

✏ -Greedy for exploration +Importance Weight (Propensity Score)
rR
⇡ rR
p

rL

Update Routers:
Add Regularization for BalanceqInternal node

f1 (q) < 0

nL , nR

f1

f2

f# of memories inthe Left/Right
subtreerLrR
Binary Classification{q, sign(rRrL )}Update Routers:
Add Regularization for BalanceqInternal node

f1 (q) < 0

nL , nRf2

f1
f# of memories inthe Left/Right
subtreerLrR
Binary Classification

Update Routers:
Add Regularization for BalanceqInternal node

f1 (q) < 0

nL , nR

f1

f2

f# of memories inthe Left/Right
subtree

rR

rL
⇢

✓

Binary Classificationq, sign (1

↵)(rR

nL
rL ) + ↵ lognR◆

Update Routers:
Add Regularization for BalanceqInternal node

f1 (q) < 0

nL , nR

f1

f2

f# of memories inthe Left/Right
subtree

rR

rL
⇢

✓

Binary Classificationq, sign (1

↵)(rR

nL
rL ) + ↵ lognR◆

Update Routers:
Add Regularization for BalanceqInternal node

f1 (q) < 0

nL , nR

f1

f2

f# of memories inthe Left/Right
subtreenL >> nR : +1 (right)

rR

rL
⇢

✓

Binary Classificationq, sign (1

↵)(rR

nL
rL ) + ↵ log
nR◆

Insertion:
Self-Consistency

f1
f2

f

m = {q, v}

Insertion:
Self-Consistency
Same query: q

f1
f2

f

m = {q, v}

Insertion:
Self-Consistency
Same query: q

f1
f2

f

m = {q, v}

Insertion:
Self-Consistency
Same query: q

f1
f2

f

m = {q, v}

Insertion:
Self-Consistency
Same query: q

f1
f2

f

m = {q, v}
If a query has been seen before, we want to just retrieve it!

Insertion:
Encourage Self-Consistency

f1

Internal nodenL , nR# of memories inthe Left/Right
subtreef2

fInsertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nR# of memories inthe Left/Right
subtreef2

fInsertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nR# of memories inthe Left/Right
subtreef2

fInsertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nRf2f# of memories inthe Left/Right
subtreem = {q, v}

Insertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nRf2f# of memories inthe Left/Right
subtreem = {q, v}

Insertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nR# of memories inthe Left/Right
subtreef2

fEnhance its
predictionm = {q, v}

Insertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nR

f2

Enhance its
predictionf# of memories inthe Left/Right
subtree⇢✓q, sign (1Binary Classificationm = {q, v}↵)f (g) + ↵ log

✓

nL
nR

◆◆

Insertion:
Encourage Self-Consistency
A memory: m = {q, v}

f1

Internal nodenL , nR

f2

Enhance its
predictionf# of memories inthe Left/Right
subtree⇢✓q, sign (1Binary Classificationm = {q, v}↵)f (g) + ↵ log

✓

nL
nR

Router’s prediction◆◆

Amortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

f

Amortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

1. Randomly sample
a memoryfAmortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

m = {q, v}

1. Randomly sample
a memoryfAmortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

1. Randomly sample
a memoryf
2. Delete it & its trace

Amortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

1. Randomly sample
a memoryf
2. Delete it & its trace

m = {q, v} 3. Re-insert itAmortized Experience Replay
Routers are updated online and may result in a lack ofself-consistency for previous insertions

f1
f2

1. Randomly sample
a memoryf
2. Delete it & its trace

m = {q, v} 3. Re-insert it

Apply replay constant times per insertion

Empirical Results

Empirical Results
Extreme Multi-Class Classification:(feature, label), zero-one lossEmpirical Results
Extreme Multi-Class Classification:(feature, label), zero-one loss

Extreme Multi-Label Classification:(feature, set of labels), Hamming lossEmpirical Results
Extreme Multi-Class Classification:(feature, label), zero-one loss

Extreme Multi-Label Classification:(feature, set of labels), Hamming loss

Image Retrieval(Caption, Image), Cosine Similarity

Extreme Multi-class Classification:
Online Progressive PerformanceExtreme Multi-class Classification:
Online Progressive Performance

ALOI (1K classes, 100K examples)
WikiPara-3 shot (10K classes, 30K examples)Extreme Multi-class Classification:
Online Progressive Performance

ALOI (1K classes, 100K examples)
WikiPara-3 shot (10K classes, 30K examples)
Extreme: very few examples per class!

Extreme Multi-class Classification:
Online Progressive Performance
WikiPara-3 shot

Higher Better

ALOI# of examples

ALOI (1K classes, 100K examples)
WikiPara-3 shot (10K classes, 30K examples)
Extreme: very few examples per class!

Extreme Multi-class Classification:
Online Progressive Performance
ALOI

WikiPara-3 shot
LOMTree(Choromanska et.al, 15)(Log(# of classes))Higher Better

One-Against-All( O(# of classes))
Recall Tree (Daumé et.al, 17)(Log(# of classes))# of examples

ALOI (1K classes, 100K examples)
WikiPara-3 shot (10K classes, 30K examples)
Extreme: very few examples per class!

Extreme Multi-class Classification:
Online Progressive PerformanceOurs: CMTALOI

WikiPara-3 shot
LOMTree(Choromanska et.al, 15)(Log(# of classes))Higher Better

One-Against-All( O(# of classes))
Recall Tree (Daumé et.al, 17)(Log(# of classes))# of examples

ALOI (1K classes, 100K examples)
WikiPara-3 shot (10K classes, 30K examples)
Extreme: very few examples per class!

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }N
i=1(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }Ni=1k{qi , vi }i=1

OAS(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }Ni=1k{qi , vi }i=1

OAS

(OAA on the subset of labelsfrom returned memories)(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }Ni=1k{qi , vi }i=1

OAS

(OAA on the subset of labelsfrom returned memories)(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }Ni=1k{qi , vi }i=1

OAS

(OAA on the subset of labelsfrom returned memories)

Computation: faster(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Extreme Multi-Label Classification:
Helping an External Inference Algorithm (One-Against-Some)

A newquery qCMT

{qi , vi }Ni=1k{qi , vi }i=1

OAS

(OAA on the subset of labelsfrom returned memories)

Computation: faster

Statistical Performance: similar, sometimes even better(1) RCV1 (1K labels), (2) AmazonCat 13K (13K labels), and (3) Wiki10-30K (30K labels)
Extreme: very large number of labels (hence linear infer = slow)

Image Retrieval:Comparison to Nearest Neighbor Approach

CMT(u):
Unsupervised version
i.e., NN (w/ Euclidean dis)at leaf level(1) Pascal (1K examples), (2) Flickr8k (8k examples), (2) MSCOCO (80K examples)
Image feature: HoG for Pascal & Flickr, Pre-trained VGG-19 for MSCOCO
Captions feature: Token Occurrences with hashing (high-dim, very sparse)Image Retrieval:Comparison to Nearest Neighbor Approach

CMT(u):
Unsupervised version
i.e., NN (w/ Euclidean dis)at leaf level# of times faster than NN(1) Pascal (1K examples), (2) Flickr8k (8k examples), (2) MSCOCO (80K examples)
Image feature: HoG for Pascal & Flickr, Pre-trained VGG-19 for MSCOCO
Captions feature: Token Occurrences with hashing (high-dim, very sparse)

How does computation scale
wrt the size of memory tree?
On ALOI, we range in dataset size from 1K to 100K

How does computation scale
wrt the size of memory tree?

Inference Time (ms)On ALOI, we range in dataset size from 1K to 100K

IncreaselogarithmicallySize of the memory tree (x 1K)

How does Self-Consistency improve?
1-shot dataset: 10K classes, 10K examples
Prob of missing itself

Self-Consistency
improves over time
Zero training err
means
perfect self-consistency# of passes of training data

How does performance scale
wrt classification difficulty?
ALOI S-shot, with S from 1 to 100(s-shot means s examples per class)How does performance scale
wrt classification difficulty?
ALOI S-shot, with S from 1 to 100(s-shot means s examples per class)How does performance scale
wrt classification difficulty?
ALOI S-shot, with S from 1 to 100(s-shot means s examples per class)

CMT (orange)
significantly outperforms
Recall Tree (Green)
when s is small,Revisit our Desired Properties

Revisit our Desired Properties
Online: Reduction to Online Classification

Revisit our Desired Properties
Online: Reduction to Online Classification
Linear Space: Store examples in leaves
O(N/ log(N )) many nodes

Revisit our Desired Properties
Online: Reduction to Online Classification
Linear Space: Store examples in leaves
O(N/ log(N )) many nodes

Logarithmic time: Regularization ensures a (provable) nearbalanced tree

Revisit our Desired Properties
Online: Reduction to Online Classification
Linear Space: Store examples in leaves
O(N/ log(N )) many nodes

Logarithmic time: Regularization ensures a (provable) nearbalanced tree
Learning-based: Reinforcement

Revisit our Desired Properties
Online: Reduction to Online Classification
Linear Space: Store examples in leaves
O(N/ log(N )) many nodes

Logarithmic time: Regularization ensures a (provable) nearbalanced tree
Learning-based: Reinforcement
Self-consistency: an asymptotic guarantee (due to replay)Thanks!

CMT is in the latest Vowpal Wabbit (VW)
https://github.com/VowpalWabbit/vowpal_wabbit
Try out CMT demos!