简单高效的线性注意力模型 FLatten Transformer - VisionTransformerusing Focused Linear Attention

FLattenTransformer:VisionTransformerusingFocusedLinearAttention

DongchenHan*XuranPan*YizengHan

ShijiSongGaoHuang

V

NXd

V

NXd

SoftmaxAttention

SoftmaxAttentiono=softmax(QKT)V

GiventheinputNtokensXeRN根c,withineachhead,self-attentioncanbewrittenas:

Q=xWQ,K=xWK,V=xWV,

oi=

j=1

Sim(Qi,kj)

Σ1Sim(Qi,kj)vj,

whereWQ,WK,WVeRc根careprojectionmatricesandSim(Q,K)=exp(QKT/d).

NXd

KT

N

d

Q

X

QKT

NXN

Highexpressivecapability

收Quadraticcomplexity0(N2d)

收Usuallyusingsparseattentionorwindowattentionpatternstoreducecomplexity

2

SoftmaxAttention

SparseGlobalAttentionPVT/PVTv2

1.WenhaiWang,EnzeXie,XiangLi,Deng-PingFan,KaitaoSong,DingLiang,TongLu,PingLuo,andLingShao.Pyramidvisiontransformer:Aversatilebackbonefordensepredictionwithoutconvolutions.InProceedingsoftheIEEE/CVFInternationalConferenceonComputerVision,pages568–578,2021.

3

2.WenhaiWang,EnzeXie,XiangLi,Deng-PingFan,KaitaoSong,DingLiang,TongLu,PingLuo,andLingShao.Pvtv2:Improvedbaselineswithpyramidvisiontransformer.ComputationalVisualMedia,8(3):415–424,2022.

SoftmaxAttention

ShiftedWindowAttentionSwinTransformer

1.ZeLiu,YutongLin,YueCao,HanHu,YixuanWei,ZhengZhang,StephenLin,andBainingGuo.Swintransformer:Hierarchicalvisiontransformerusingshiftedwindows.InProceedingsoftheIEEE/CVFInternationalConferenceonComputerVision,pages10012–10022,2021.

4

SoftmaxAttention

Cross-ShapedWindowAttentionCSwinTransformer

1.XiaoyiDong,JianminBao,DongdongChen,WeimingZhang,NenghaiYu,LuYuan,DongChen,andBainingGuo.Cswintransformer:Ageneralvisiontransformerbackbonewithcross-shapedwindows.InProceedingsoftheIEEE/CVFConferenceonComputerVisionandPatternRecognition,pages12124–12134,2022.

5

V

Q

d

N

T

K

×

N×d

N×d

Q

N×d

LinearAttention

LinearAttentiono=Φ(Q)(Φ(K)TV)

Carefullydesignedkernelsareintroducedastheapproximationoftheoriginalsimilarityfunction:

Sim(Q,K)=Φ(Q)Φ(K)T,

Nφ(Qi)φ(kj)T

oi=Σj=1Σ1φ(Qi)φ(kj)Tvj

Σ1(φ(Qi)φ(kj)T)vj

=

φ(1φ(kj)T

φ(

=

φ(1φ(kj)T).

KTV

d×d

收Inferiorperformance

LinearcomplexityΘ(Nd2)

Canenjoyalargereceptivefieldwhile

maintainingalowamountofcomputation

6

1.KatharopoulosA,VyasA,PappasN,etal.Transformersarernns:Fastautoregressivetransformerswithlinearattention[C]//InternationalConferenceonMachineLearning.PMLR,2020:5156-5165.

LinearAttention

HydraAttention0(Nd)

1.DanielBolya,Cheng-YangFu,XiaoliangDai,PeizhaoZhang,andJudyHofman.Hydraattention:Eficientattentionwithmanyheads.InComputerVision–ECCV2022Workshops:TelAviv,Israel,October23–27,2022,Proceedings,PartVII,pages35–49.Springer,2023.

7

LinearAttention

EfficientAttention0(Nd2)

1.ZhuoranShen,MingyuanZhang,HaiyuZhao,ShuaiYi,andHongshengLi.Eficientattention:Attentionwithlinearcomplexities.InProceedingsoftheIEEE/CVFWinterConferenceonApplicationsofComputerVision,pages3531–3539,2021.

8

IMotivation

Todesignlinearattentionmodulethatperforms

onparwithorevenbetterthanSoftmaxattention.

9

Linear

Attention

FocusAbility

Original

Image

Softmax

Attention

Sharp

Focusoncertainregions

收Smooth

收Closertotheaverageofalltokens

FocusedLinear

Attention

(Ours)

10

FocusAbility

FocusedFunctionfp

Sim(Qi,kj)=φp(Qi)φp(kj)T,

whereφp(x)=fp(ReLU(Q)),afp(x)=⋅x∗∗p.

FeatureDirectionAdjustmentwithfp

Letx=(x1,x2,x3,⋯,xn)∈ℝn,ay=(y1,y2,y3,⋯,yn)∈ℝn

Assume:(1)∃!m,s.t.xm=1xi),∃!n,s.t.yn=1yj),

(2)0<x,y<IxIIyI.

Forapairoffeatures{x,y}withm=n:∃p>1,s.t.φp(x),φp(y)>x,y.

Forapairoffeatures{x,y}withm≠n:∃p>1,s.t.φp(x),φp(y)<x,y.

11

FocusAbility

Proof:

φp(x)=fp(ReLU(x))=fp(x),aφp(y)=fp(ReLU(y))=fp(y),

fp(x)=x∗∗p=x,afp(y)=y.

Therefore,wehave:

〈φp(x),φp(y)〉=〈fp(x),fp(y)〉=fp(x)fp(y)〈u,V〉

=xy〈u,V〉,

where

〈u,V〉=〈,〉

12

1

xp

FocusAbility

Proof:

u,V=II,II=

=

Σ1xy

Σ1ab

1ap1bp),

ai=,bi=,ai,bi∈[0,1].

Basedonourassumption,wehave:

∃!m,s.t.am=1，∃!n,s.t.bn=1.

Therefore,

a={，b={.

13

FocusAbility

Considerthefollowingtwocases:

(1)m=n:

〈u,v〉=

==1.

1根1

〈φp(x),φp(y)〉=xy〈u,v〉=xy>〈x,y〉.

Thuswehave,

∃p>1,s.t.〈Φp(x),Φp(y)〉>〈x,y〉.

1根1

(2)m≠n:

〈u,v〉=

==0.

1根1

〈φp(x),φp(y)〉=xy〈u,v〉=0<〈x,y〉.

Thuswehave,

∃p>1,s.t.〈Φp(x),Φp(y)〉<〈x,y〉.

1根0+0根1

14

FocusAbility

Exampletoshowtheefectsoffp:

k2

k4

f3(k1)f3(k2)

f3(k3)f3(k4)

k1

k3

(b)AttentionMap

(a)QueryandKeys

15

FeatureDiversity

(a)SoftmaxAttentionMapMatrixRank:196

Softmaxattentioncanlearn

afull-rankattentionmap

(b)LinearAttentionMapMatrixRank:54

收Linearattentioncannotlearnanattentionmapwitharankgreaterthanheaddim64

LinearAttention:

O=φ(Q)φ(K)TV

Rank(φ(Q)φ(K)T)

≤min{Rank(φ(Q)),Rank(φ(K)T)}

≤iͺ{N,d}

16

RankRestorationModel(DWC)

O=φ(Q)φ(K)TV+DWC(V),

O=(φ(Q)φ(K)T+MDWC)V=MeqV.

FocusedLinearAttention(FLatten)

O=Sim(Q,K)V

=Φ(Q)Φ(K)TV+퐃W͗(V),

whereφp(x)=fp(ReLU(Q)),fp(x)=⋅x∗∗p.

FeatureDiversity

(c)Ours

MatrixRank:196

17

IFocusedLinearAttention

O=Sim(Q,K)V=Φ(Q)Φ(K)TV+퐃W͗(V)

FLatten

Transformer

{

Lowcomputationcomplexityaslinearattention.

HighexpressivecapabilityasSoftmaxattention.

Largereceptivefield.

Flexibility.

18

ExperimentResults

PerformancesonImageNetClassification:

19

ExperimentResults

PerformancesonCOCOobjectdetection:

20

ExperimentResults

Comparisonwithotherlinearattentiondesigns:

21

IExperimentResults

Accuracy-RuntimecurveonImageNet:

22

ExperimentResults

Ablationoneachmodule

basedonDeiT-T.

Ablationonwindowsize

based