Exploring Priors and Long-Term Memory
in Transformers
Yuwei Sun
RIKEN AIP
Araya Research

Can Multimodal LLMs Learn by Memorizing More Data?
Probably not..

[Tong, 2024][Jones, 2023]

SJTU CS Joint Workshop

2

Neuro-symbolic is not very helpful for complex dataØ Transformers are not good at learning discrete information from higher-dimensional
perception data such as images, causing hallucination, inefficiency in training, and being
data-hungry.
Ø Neuro-symbolic approach offers a more stable way to learn discrete symbols and theirrelations. However, the brain can learn without any annotated data, and real-world image
data cannot be fully structured with a set of symbols.
1. Background(5) Challenges

Unknown Mechanism of Induction

1. Background

ColorShap
MateAzim[Mao 2019](5) Challenges

Unknown Mechanism of Induction

Color：black，brown，…
Shape：square, round, …. •Material：wooden，…
•Azimuth: left，right, …
•Complexity of Image data
Occlusion
Complex Concept
Complex Scene

Complexity of Image data•••

Occlusion
Complex Concept
Complex SceneSJTU CS Joint Workshop

3Difficulty of obtaining labels：there are many concepts aAbsence of inductive biases such as convolution operations
for localized knowledge in Transformers[Amir, 2022]

Low-level

AlexNet [Krizhevsky, 2012]

High-level

Do not learn structuralfeatures that align withthe input dataØ Difficult to generalize
with limited samples[Dosovitskiy, 2021]

Ø Unlike convolution operations in CNNs, Transformers do not learn structural features thatalign with the input data and usually perform worse than CNNs with limited samples.
SJTU CS Joint Workshop

4Vision Transformer and attention mechanism
Multi-Head Attention

Learn relations between patches
V(Value) K(Key) Q(Query)[Dosovitskiy, 2021]

Ø Transformers use pairwise attention to establish
correlations among disparate input segments.
SJTU CS Joint Workshop* Assign different weights to inputtokens.
5Analysis of attention weights in pretrained Transformers
Vision Transformer

BERT (Natural language Transformer)

Layer 12Highlysparseattention
Sparsity in token relations is learned
implicitly fromweights
Highlytraining data, which is inefficient and data-hungry.

sparseattentionweights

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.ForForeacheachheadheadin each layer, 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. The
valuesstate.
SJTU
CS Joint Workshop
valuestoto0.90
0.90isisdisplayed
displayedasasaaviolin
violinplot
plotininaapanel.
panel.kkindicatesindicatesthethesizesizeofofaametastablemetastablestate.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
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
middle layers (3, 4, and 5). Class (I) in red: Very large metastable state or global fixed point, which
middlemiddlelayers
layers(3,(3,4,4,andand5).
5).Class
Class(I)(I)ininred:
red:VeryVerylarge
largemetastablemetastablestate
stateororglobal
globalfixed
fixedpoint,
point,which
which
is predominant in the first layer. These heads can potentially be replaced by averaging operations.

The attention sparsity is measured by the minimal number of required tokens whose
6attention scores add up to 0.90
[Ramsaue 2021; Sun 2023]

Inducing an information bottleneck in the attention mechanism

A tte ntio n

Patches

Transformer
Prior-guided
attention

Priors[Goyal et al. ICLR’22]

Transformer + Shared
WorkspaceØ Priors are general assumptions about samples, such as the aggregated featuresfrom different samples of the same object.
Ø Competition through a bottleneck results in naturally emerging specialized priors.
SJTU CS Joint Workshop

7

Specialized neural modules
Definition: correspondence between strongly interconnected structuralcomponents of a network (modules) and the specialized functions they perform.
In animal brains, modularity favors evolvability, the ability to adapt tochanging environments with common sub-problems [Clune, 2013]
Modular network

Non-modular network

SJTU CS Joint Workshop

8Working memory in the biological brain and Global Workspace Theory
When we use working memory: (1) Learning in a novel situation(2) Taking a different approach in familiar situations
Capacity: Limited amount of information it can hold at one time
Function: Acts as a mental workspace where information can be
manipulated: multiple specialized modules (potientially, multimodal) compete to write to the shared space; information in theshared space is broadcast to all modules afterwards.

Long-term memory(Associative memory)

Reuse priors[Baars, 1988; Butlin, 2023]

Learning priors

Working memory
and Long-term memory(Semantic memory)SJTU CS Joint Workshop

9X

Multi-head
selfattention

FF

Globalworkspacelayer(0.45M)

Smaller computational cost

Global Workspace Layer

XMulti-head
selfattention

FF…

Globalworkspace
layer

Enhanced
ViT Blocks (12) layer
efficiency AiT Blocks (6)(0.45M)(1) Collecting patchesfrom all batch samples(2) Computing the bottleneck attention
and selecting relevant…patches…
priors…

(3) Information broadcast
from the global workspace

Enhancedlayer
ViT Blocks (12)efficiency AiT Blocks (6)patchPriors
X

Sun et al., Associative Transformer, NeurIPS workshop; arXiv:2309.12862
SJTU CS Joint Workshop

10Global Workspace Layer(1) Collecting patches fromall batch samples(2) Computing the bottleneck attention
and selecting relevant patches(3) Information broadcast
from the global workspace

Bottleneckattention
balance lossØ The explicit memory stores and updates a set of priors (as queries) by attending todifferent patches based on the multi-head cross attention.
Ø The sparsity is enabled through a bottleneck using the top-k hard attention.
SJTU CS Joint Workshop

11Global Workspace Layer(1) Collecting patches fromall batch samples(3) Information broadcast
from the global workspace(2) Computing the bottleneck attention
and selecting relevant patches

Enhanced
representations
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 thememory enables effective retrieval of knowledge within the memory.
SJTU CS Joint Workshop

12

Continuous Hopfield Network

Multiple interactions of energy
reduction to reconstruct patterns[Ramsauer, 2021]

SJTU CS Joint Workshop

13

d points

Token retrieval with a continuous Hopfield network
EnergyØ Energy of the continuous Hopfield network

Time stepsof attractionØBasinInversetemperature 𝛽 and basins of attraction

Reconstructed features from memory

𝜷=𝟏[Ramsauer, 2021]

metastable states

𝛽 = 0.2

Identicalpatches
retrieved
𝛽 = 0.05

global fixed point𝛽 = 1.0

𝛽 = 4.0

Ø A smaller 𝛽 results in a metastable stateØ A very large or small 𝛽 both can lead towithin the basin of multiple attractors.
local minima.
SJTU CS Joint Workshop
14

Problem 1: monolithic priors select repeated tokens:
Introducing Bottleneck Attention Balance LossØ Cascading multiple bottleneck attentions leads todifficulty in forming specialized priors
jiPriorsTokens (position)

SJTU CS Joint Workshop

Bottleneck

Ai,j15Diversity in patch selection with the new loss

Diversity 0.24

Diversity 0.70

SJTU CS Joint Workshop

Diversity 0.77

16

Problem 2: Computational load with the squash operation
The squash layer concatenates all tokens in the batch, Ξ ∈ ℝ !×# ×$ , allowing for
across-sample learning but also increasing the computational cost for the attention
mechanism.
To reduce its the computational load:(1) a low rank memory, where the squashed representations are projected to a latent
space of dimension 𝐷 ≪ 𝐸(2) an attention bottleneck with capacity k ≪ 𝐵×𝑁, e. g. , 1.6% ~ 3.2% of all the tokens
in our experiments
Ø Less than a 3% increase in computationcompared to Vision Transformers of similar
size.

SJTU CS Joint Workshop

17Enhanced efficiency in image classification tasks
Parameters (M)6 layers

ImageNet100

12 layers

Our study demonstrates that AiT outperforms existing sparse Transformer
models including the variants of Coordination [Goyal 2022] and Vision
Transformers, without pretraining on external data.
SJTU CS Joint Workshop

18

Vision-language relational reasoning tasks

ConcatenatedImagetokens

Texttokens

Sort-of-CLEVR dataset [Santoro, 2017]

Cross-modal priors

SJTU CS Joint Workshop

19Conclusions
Ø Associative Transformer enhances parameter efficiency in thetraining of Transformer-based models, making them more accessible
and cost-effective
Ø Implementing the cognitive science theory of the Global Workspace is
crucial for a better understanding of human-like relational reasoningØ Other tasks and domains, such as audio and videoØ Safe deployment in real-world applications.
Privacy of neural module learning

Adversarial attacks

Bidirectional Contrastive Split Learning
Sun et al. AAAI 2024

Attacking Distance-aware Attack
Sun et al. Transactions on AI 2023

SJTU CS Joint Workshop

20

Exploring Priors and Long-Term Memory
in Transformers
Yuwei Sun
RIKEN AIP
Araya Research

Associative
Transformer paper