Associative Transformer Is A Sparse Representation LearnerYuwei Sun1,2, Hideya Ochiai1, Zhirong Wu3, Stephen Lin3, Ryota Kanai41

The University of Tokyo, 2RIKEN, 3Microsoft Research, 4Araya

INTRODUCTION

Bottleneck Attention Balance Loss

Transformer models use pairwise attention to establish correlations among
disparate segments of input information.
Competition, which reveals naturally sparser interactions among attention heads
in pairwise attention, is important for learning meaningful representations.

Learning specialized priors in layers cascaded in depth requires a mechanism thatcounteracts the inherent loss of input specificity, as information flows throughmultiple layers.

The sparsity of attention for a specificpatch’s interactions withotherpatchesPs
is computed as arg min j=1 Ai,j ≥ 0.9.

The bottleneck attention balance loss encourages the selection of diverse patchesfrom different input positions into the shared workspace.
Accumulative attention scores: `importancei,l =s

Unlike convolution operations in CNNs,
self-attention in Transformers does notpossess inductive biases that allow it toattend to different segments of the input data. We aim to introduce architecture that can foster competition amongpatches by constraining the number ofpatches that each head can focus on,
thereby inducing an inductive bias formeaningful patch learning.
s patchess patchess patchess patchess patchess patchess patchess patchess patchess patchess patchess patches

Figure A.3: Analysis of operating modes of the heads of a pre-trained BERT model. For each head
FigureA.3:
Analysisofmodes of the heads of a pre-trained BERT model. For each headhead
Figureof operating
operatingineachA.3:
layer,Analysis
the distributionof the minimal number k of patterns required to sum up the softmax
inlayer,
distribution
theofrequired
tometastablesum upup thethestate.
softmax
in eacheachtolayer,
thedistribution
the minimal
minimalnumberof patterns
patterns
sum
softmax
values0.90theis displayed
as ofaofviolinplot in anumber
panel. kkk indicates
therequired
size of ato
Thevaluestoisdisplayed
plotinkk indicatesthedistribution.
size ofof aa metastablemetastablestate.
Thevalues
to 0.900.90inisthe
displayed
as aeacha violin
violinplotgives
in aa panel.
panel.
indicates
thesizestate.
Thebold number
center asof
panel
the mediank̄ of theThe heads
in eachboldnumber
inthecenterofeachpanelgives
themediank̄ofthedistribution.
The
heads
in
each
bold numberin thecenter ofpanel gives
thebelong
medianof class
the distribution.
headsin.inClass
each
layer
are sorted
accordingto eachk̄. Attention
heads
tok̄thethey mainlyTheoperate
layer
aresorted
accordingtok̄.
Attention
heads
belongtotheclasstheymainlyoperatein.
Classlayerinareblue:sortedaccordingto k̄. state
Attention
heads
belongclass theymainlyoperatein. Class
(IV)
Small
metastableor fixed
pointclosetotothea single
pattern,which
is abundant
in
(IV)
in
blue:
Small
metastablestate
orfixed
point
close
toa
singlepattern,which
isabundant
(IV)
in
blue:
Small
metastablestate
orfixed
point
close
to
a
singlepattern,which
isabundantinthe middle layers (6, 7, and 8). Class (II) in orange: Large metastable state, which is prominent inthelayers(6,7,8).
(II)
in
Large
metastablestate,
which
is prominentprominentin
the middle
middlelayers(6,and7, andand8). ClassClass(II)in orange:orange:Largemetastable
which
is
in
middlelayers(3, 4,5). Class(I) inred:
Very large
metastablestate state,or global
fixedpoint, which
middlelayersand5).
Class(I)metastablestate oror global
globalfixed point,point,which
middlelayers (3,(3,in4,4,theandfirst5). layer.
Class These(I) in
in red:red:
Verylarge
metastable
state
fixed
which
is
predominant
headsVery
canlarge
potentiallybe replaced
byaveraging
operations.
isClass
predominant
inthefirstlayer.
Theseheads
canpotentiallybereplaced
byaveragingoperations.
is predominant
in the Medium
first layer.
These heads
canpotentiallybe replacedby averagingoperations.
(III) in green:metastablestate,
whichis frequently
observed
in higher layers.
WeClass(III)ingreen:
Mediummetastablestate,
whichisfrequentlyobserved
inhigherlayers.
Wehypothesize
theseMedium
heads aremetastableused to collectrequired observed
to perform
respective
Class (III) inthat
green:
state,information
which is frequentlyinthe
higherlayers.task.
Wehypothesize
thattheseheads
areusedtocollect
information
requiredtoperform
therespectivetask.
Theseheads thatshouldbe heads
the maintocollectimprovetransformerand BERT
models.the respective task.
hypothesize
theseare target
used toinformation
required
to perform
Theseheads
shouldbethemaintargettoimprovetransformer
andBERTmodels.
These heads should be the main target to improve transformer and BERT models.

Figure 1. Inspecting attention heads in Vision Transformers.

METHODS

We propose the Associative Transformer (AiT) with a novel global workspace layer
building upon recent neuroscience studies of the Global Workspace Theory andassociative memory.
Modularization of knowledge
can find resonance with the neuroscientific ground74
74
74
ing of the Global Workspace
Theory (GWT). GWT explains a cognitive architecture
where diverse specialized modules compete to write information into sharedworkspace through a communication bottleneck.
When examining information retrieval in the human brain, it is evident that memory typically encompasses both working memory and long-term memory in thehippocampus. Specifically, the hippocampus operates on Hebbian learning, akinto the associative memory found in modern Hopfield networks.

(1) Collecting modules(2) Computing the bottleneckattention and selecting modules(3) Information broadcast
from the global workspaceMX

Ati,j,l ,j=1

Chosen instances: `loadsi,l =M
X(Ati,j,l > 0),

j=1
B×N
Var({`importancei,l }B×N)
Var({`}
loadsi,l l=1 )
l=1`bottlenecki = 1 PB×N
+ 1 PB×N.
( B×N l=1 `importancei,l )2 +  ( B×N l=1 `loadsi,l )2 + RESULTS
Our study demonstrates that AiT outperforms the previously employed attentionbased approaches such as the Coordination, the Set Transformer, and the Perceiver
when applied to vision-related tasks.

Table 1. Model performance comparison inimage classification tasks
Methods
CIFAR10 CIFAR100 Triangle Average Model Size
AiT-Base (Ours)
85.44
59.10
99.59
81.38
91.0
AiT-Small (Ours)
83.34
56.30
99.47
79.70
15.8
Coordination
75.31
43.90
91.66
70.29
2.2
Coordination-DH 72.49
51.70
81.78
68.66
16.6
Coordination-D74.50
40.69
86.28
67.16
2.2
Coordination-H
78.51
48.59
72.53
66.54
8.4
ViT-Base
83.82
57.92
99.63
80.46
85.7
ViT-Small
79.53
53.19
99.47
77.40
14.9
Perceiver
82.52
52.64
96.78
77.31
44.9
Set Transformer
73.42
40.19
60.31
57.97
2.2
BRIMs
60.10
31.75
45.93
4.4
Luna
47.86
23.38
35.62
77.6

Figure 4. Model size vs. accuracy.
Using the low-rank memory (LM) that has a more diverse set of priors showed benefits in both improving the performanceand decreasing the model size. The Hopfield network (HN) maintained the model performance while reducing themodel size by replacing the cross-attention with more efficient retrieval, which was effective only when either the LM or
SA component was applied. We assume that the retrieval relies on a diverse set of priors, which is enabled using theenhanced bottleneck attention and the learning through self-attention.

Increased Bottleneck Attention Distribution Sparsity
Slot position

samplesEpoch 1
Input patch position

Epoch 20modules

Epoch 40

Figure 2. Inducing global workspace for emerging module specialization.
(1) Collecting modules(2) Computing the bottleneckattention and selecting modules(3) Information broadcast
from the global workspacepriorsFigure 6. Selecting
diverse patches from
Figure 5. Selected patches by the bottleneck attention. different input positions.
Prior Specialization: patches in one image can be attended sparsely by different
priors through the bottleneck attention.

samples

Image 1

Upscalemodules

X

Image 2

Slot 1-16

Figure 3. The scheme of the global workspace layer.
Image 3The squash layer concatenates patches within one batch V ∈ RB×N ×E intovectors V ∈ R(B×N )×E .
Squashed representations are projected to a low-rank latent space of dimension
D << E and then are sparsely selected and stored in the explicit memory via afixed bottleneck k << (B × N ).
Explicit low-rank memory with limited slots learns M priors γ = RM ×D , to computethe bottleneck attentions that extract different sets of patches from the input.
tt t
headi = top-k(Ai )Ξ WtV ,γ̂ttt= LN(Concat(head1, . . . , headA)W O ).

The Hopfield network utilizes the memory to reconstruct the input, where alearnable linear transformation (LT) scales the memory contents fLT(γ t+1) to matchthe input dimension E.
1 t tT1 2 ˆttt+1 t−1t
E(ξ ) = −lse(β, fLT(γ )ξ ) + ξ ξ + β logM + ζ , ξ = arg minE(ξ ).
tξ2
2

Corresponding author: ywsun@g.ecc.u-tokyo.ac.jp

Image 4
Slot 17-32
Slot 1

Slot 2

Slot 3

Slot 4(a) A glimpse of the attention maps at Slot 1 to Slot 4
of four distinct input images.

(b) The attention maps for all 32 slots in the memory bank,applied to four distinct input images. Each memory slotlearned to attend to different regions of pixels in input images.

Figure 7. Learned bottleneck attention maps in CIFAR10.

CONCLUSIONSAssociative Transformer (AiT) leverages a diverse set of priors with the emerging
specialization property to enable sparse association among representations via theHopfield network. The comprehensive experiments demonstrate AiT’s efficacy
compared to conventional models including the previous coordination method.

arXiv:2309.12862