Associative Transformer Is A Sparse
Representation LearnerYuwei Sun, Hideya Ochiai, Zhirong Wu, Stephen Lin, Ryota Kanai

Inductive bias allowing a model to generalize more
effectively to unseen samples
AlexNet [Krizhevsky, 2012]Ø Transformers use pairwise attention to establish correlations
among disparate input segments
Ø Unlike convolution operations in CNNs, self-attention inTransformers does not possess an inductive bias that is
consistent with the underlying input data structure
NeurIPS 2023 AMHN workshop[Dosovitskiy, 2021]
2

Emergent inductive biases through competitionA tte ntio nPatches

Inductive biases

PriorsØ Absence of a bottleneck for localized, and contextual representation learningØ Competition results in naturally emerging specialized priors
Ø Self-attention in Transformers becomes expensive with scaling
NeurIPS 2023 AMHN workshop

3

Does competition exist within the self-attention mechanism?
Attention heads in Vision Transformer (Ours)

Attention heads in BERT [Ramsaue, 2021]

Layer 12Layer 1Figure A.3: Analysis of operating modes of the heads of a pre-trained BERT model. For each headFigureFigureA.3:
A.3:Analysis
Analysisofofoperating
operatingmodes
modesofofthe
theheads
headsofofaapre-trained
pre-trainedBERT
BERTmodel.
model.For
Foreacheachheadhead
NeurIPS
2023 AMHN
workshopin eachlayer, the distribution of the minimal number k of patterns required to sum up the softmax
inineacheachlayer,
layer,thethedistribution
distributionofofthe
theminimal
minimalnumber
numberkkofofpatterns
patternsrequired
requiredtotosum
sumup
upthe
thesoftmax
softmax
values to 0.90 is displayed as a violin plot in a panel. k indicates the size of a metastable state. Thevalues
valuestoto0.90
0.90isisdisplayed
displayedasasaaviolin
violinplot
plotininaapanel.
panel.kkindicatesindicatesthethesizesizeofofaametastablemetastablestate.
state.The
The
bold number in the center of each panel gives the median k̄ of the distribution. The heads in eachboldboldnumber
numberininthe
thecenter
centerofofeach
eachpanel
panelgives
givesthethemedian
mediank̄k̄ofofthe
thedistribution.
distribution.TheTheheads
headsinineach
eachlayer are sorted according to k̄. Attention heads belong to the class they mainly operate in. Classlayer
layerarearesortedsortedaccording
accordingtotok̄.k̄.Attention
Attentionheads
headsbelong
belongtotothe
theclassclassthey
theymainly
mainlyoperate
operatein.
in.ClassClass(IV) in blue: Small metastable state or fixed point close to a single pattern, which is abundant in(IV)(IV)ininblue:
blue:Small
Smallmetastablemetastablestate
stateororfixed
fixedpoint
pointclose
closetotoaasingle
singlepattern,pattern,which
whichisisabundant
abundantinin
the middle layers (6, 7, and 8). Class (II) in orange: Large metastable state, which is prominent inthethemiddle
middlelayers
layers(6,(6,7,7,andand8).
8).ClassClass(II)(II)ininorange:orange:Large
Largemetastablemetastablestate,state,which
whichisisprominentprominentinin

4

Does competition exist within the self-attention mechanism?
Competition, which reveals naturally sparser interactions among attention heads inpairwiseattention,
is inimportantfor learning meaningfulrepresentations.
Attention
heads
Vision Transformer
Attention
heads in BERT
Layer 12

Learning a set of priors that can guide attention in the learning process.

Layer 1Figure A.3: Analysis of operating modes of the heads of a pre-trained BERT model. For each head
FigureFigureA.3:
A.3:Analysis
Analysisofofoperating
operatingmodes
modesofofthe
theheads
headsofofaapre-trained
pre-trainedBERT
BERTmodel.
model.ForForeacheachheadhead
NeurIPS
2023 AMHN
workshopin eachlayer, the distribution of the minimal number k of patterns required to sum up the softmax
inineacheachlayer,
layer,thethedistribution
distributionofofthe
theminimal
minimalnumber
numberkkofofpatterns
patternsrequired
requiredtotosum
sumup
upthe
thesoftmax
softmax
values to 0.90 is displayed as a violin plot in a panel. k indicates the size of a metastable state. Thevalues
valuestoto0.90
0.90isisdisplayed
displayedasasaaviolin
violinplot
plotininaapanel.
panel.kkindicatesindicatesthethesizesizeofofaametastablemetastablestate.
state.TheThebold number in the center of each panel gives the median k̄ of the distribution. The heads in eachboldboldnumber
numberininthe
thecenter
centerofofeach
eachpanel
panelgives
givesthethemedian
mediank̄k̄ofofthe
thedistribution.
distribution.TheTheheads
headsinineach
eachlayer are sorted according to k̄. Attention heads belong to the class they mainly operate in. Classlayer
layerarearesorted
sortedaccording
accordingtotok̄.k̄.Attention
Attentionheads
headsbelong
belongtotothe
theclassclassthey
theymainly
mainlyoperate
operatein.
in.ClassClass(IV) in blue: Small metastable state or fixed point close to a single pattern, which is abundant in(IV)(IV)ininblue:
blue:Small
Smallmetastablemetastablestate
stateororfixed
fixedpoint
pointclose
closetotoaasingle
singlepattern,pattern,which
whichisisabundant
abundantinin
the middle layers (6, 7, and 8). Class (II) in orange: Large metastable state, which is prominent inthethemiddle
middlelayers
layers(6,(6,7,7,andand8).
8).Class
Class(II)(II)ininorange:orange:Large
Largemetastablemetastablestate,state,which
whichisisprominentprominentinin

5Inspired by the Global Workspace Theory (GWT)• Guided attention in a limited capacity workspace• Modules compete to write into the shared space• Availability of information in the workspace through broadcast
Ø The Coordination method functions as a Global Workspace
based on cross-attention over input segments of a sample[Goyal, ICLR 2022][Baars, 1988; Dehaene, 1998]

patchesNeurIPS 2023 AMHN workshop

6Inductive biases through prior specialization in global workspace(1) Collecting patchesfrom all batch samples(2) Computing the bottleneckattention and selecting patches(3) Information broadcast
from the global workspacepriorsPriorsPriorsX

NeurIPS 2023 AMHN workshop

7

Inductive biases through prior specialization in global workspace(1) Collecting patches fromall batch samples(2) Computing the bottleneckattention and selecting patches

Bottleneckattentionbalance loss

(3) Information broadcast
from the global workspace

PriorsØ The explicit memory stores and updates a set of priors by attending to different patches,
which are smoothed by Exponential Moving Average (EMA)Ø The balance loss encourages diverse patch selection of the Top-k hard attention
NeurIPS 2023 AMHN workshop

8Selected patches by the specialized priors

Diversity 0.24

Diversity 0.70

NeurIPS 2023 AMHN workshop

Diversity 0.77

9Inductive biases through prior specialization in global workspace(1) Collecting patches fromall batch samples(3) Information broadcast
from the global workspace(2) Computing the bottleneckattention and selecting patches

Enhancedrepresentations
LinearTransformation

Ø A continuous Hopfield network [Ramsauer 2021] uses the learned priors
in the explicit memory as attractors to reconstruct input patchesØ Iteratively decreasing the energy of a patch with respect to theattractors enables effective retrieval of knowledge within the memory
NeurIPS 2023 AMHN workshop

10d points

Ø Energy of the continuous Hopfield network

EnergyAssociative memory in a continuous Hopfield network

Batch time step

Inversetemperature 𝛽 and basins of attraction
Basin of attraction

Reconstructed sample representations

𝜷=𝟏[Ramsauer, 2021]

metastable states

𝛽 = 0.2

global fixed pointØ A smaller 𝛽 is more likely to result ina metastable state

Identicalpatches
retrieved
𝛽 = 0.05

NeurIPS 2023 AMHN workshop

𝛽 = 1.0

𝛽 = 4.011

Associative Transformer
Ξ"!
Explicit
memory ! !
priorsQK

SquashBottleneck

Ξ!

V

Multi-headcrossattention

EMA

Explicit
memory ! !"#Hopfield
networkLT

priors
The same small set of upscaled
priors is utilized as attractors toadapt each image.

Top-k hard attentionenergyImages are represented by differenthues. Patches from the same image
are distinguished by differences inbrightness.

XtimeMulti-head
selfattention

FFGlobalworkspacelayerattractors

Ξ!

State

Bottleneck attention
balance loss(a) Global workspace layer

NeurIPS 2023 AMHN workshop(b) Associative Transformer block12

Prior specialization
Ø Patches across samples are attended sparsely by different priors,with emergent specialization in a prior’s attention
Image 1

Image 2

Slot 1-16
Prior
1-16

Image 3

Image 4
Slot 17-32
17-32
Prior
Slot 11 Prior
Slot 22 Prior
Slot 33 Prior
Slot 44
Prior(b) The attention maps for all 32 slots in the memory bank,2023 AMHN
13(a) A glimpse of the attention maps at Slot 1 to SlotNeurIPS4
appliedtoworkshop
four distinct input images. Each memory slotof four distinct input images.
learned to attend to different regions of pixels in input images.

AiT outperforms latent memory-based Transformers

NeurIPS 2023 AMHN workshop

14

AiT is a layer-efficient model

6 layers

12 layers

NeurIPS 2023 AMHN workshop

15

Multi-modal relational reasoning tasks
Sort-of-CLEVR dataset [Santoro, 2017](upper bound)Ø AiT achieved competitive performance with the upper bound non-Transformer basedmodels that rely on a built-in inductive bias, through the efficient association of disparate
information fragments.
NeurIPS 2023 AMHN workshop

16

Conclusions•

Associative Transformer (AiT): a biologically plausible sparse attention mechanism
based on the Global Workspace Theory and associative memory•AiT is a sparse representation learner, leveraging bottleneck attention to acquire
specialized priors•

Adaptive low-rank priors increase memory capacity, allowing AiT to learn up to 128
specialized priors from a diverse pool of 32.8k patches (refer to the paper)•The learned priors serve as attractors in a Hopfield network: The first work to incorporate
the Hopfield network as an integral element in a sparse attention mechanism forinductive biases

NeurIPS 2023 AMHN workshop

17

Associative Transformer Is A Sparse
Representation LearnerYuwei Sun, Hideya Ochiai, Zhirong Wu, Stephen Lin, Ryota Kanai