Understanding Self-attention Mechanism via Dynamical System Perspective
Zhongzhan Huang Mingfu Liang Jinghui Qin Shanshan Zhong Liang Lin
The self-attention mechanism (SAM) is widely used in various fields of
artificial intelligence and has successfully boosted the performance of
different models. However, current explanations of this mechanism are mainly
based on intuitions and experiences, while there still lacks direct modeling
for how the SAM helps performance. To mitigate this issue, in this paper, based
on the dynamical system perspective of the residual neural network, we first
show that the intrinsic stiffness phenomenon (SP) in the high-precision
solution of ordinary differential equations (ODEs) also widely exists in
high-performance neural networks (NN). Thus the ability of NN to measure SP at
the feature level is necessary to obtain high performance and is an important
factor in the difficulty of training NN. Similar to the adaptive step-size
method which is effective in solving stiff ODEs, we show that the SAM is also a
stiffness-aware step size adaptor that can enhance the model's representational
ability to measure intrinsic SP by refining the estimation of stiffness
information and generating adaptive attention values, which provides a new
understanding about why and how the SAM can benefit the model performance. This
novel perspective can also explain the lottery ticket hypothesis in SAM, design
new quantitative metrics of representational ability, and inspire a new
theoretic-inspired approach, StepNet. Extensive experiments on several popular
benchmarks demonstrate that StepNet can extract fine-grained stiffness
information and measure SP accurately, leading to significant improvements in
various visual tasks.