HyperAI

Transformer models have achieved remarkable success across diverse deep learning fields, including natural language processing (NLP) and computer vision (CV). One drawback of these models is that the computational cost of the softmax attention, the core component of the transformer, exhibits quadratic complexity in both time and memory. As data scales up various attempts have been reported to overcome this bottleneck. The objective of this study is to propose a novel attention mechanism, "Cumulant Attention", that systematically balances efficiency and accuracy. This proposal introduces a statistical-mechanics perspective and a reliable approximation based on cumulant expansion into the attention layer. The low-order variant reduces computational complexity to linear order, similar to the linear attention, while keeping nonlinearity of the softmax attention. We evaluate several variants on CV tasks, including image classification with ViT on ImageNet-100 and video classification with ViViT on UCF-101. Experimental results demonstrate that the cumulant attention outperforms the linear attention and achieves accuracy comparable to the softmax attention. These findings validate the effectiveness of our approach and highlight future directions, including scaling to larger models, extending to other modalities, and optimizing implementations for GPU hardware.

One-sentence Summary

Researchers at Nagoya University's Institute of Materials and Systems for Sustainability propose Cumulant Attention, an attention mechanism that applies cumulant expansion from statistical mechanics to approximate softmax attention at linear computational complexity while preserving its nonlinearity, outperforming linear attention and matching softmax accuracy on ViT-based image and video classification tasks.

Key Contributions

The paper introduces Cumulant Attention, an attention mechanism that incorporates a statistical-mechanics perspective and cumulant expansion into the transformer attention layer to balance efficiency and accuracy.
A low-order variant of Cumulant Attention reduces attention computation to linear order while preserving nonlinear behavior associated with softmax attention.
Experiments on image classification with ViT on ImageNet-100 and video classification with ViViT on UCF-101 show that Cumulant Attention outperforms linear attention and achieves accuracy comparable to softmax attention.

Introduction

Transformer attention powers modern models across language, vision, speech, graphs, and multimodal tasks, but its standard softmax attention scales quadratically with sequence length, creating a serious efficiency bottleneck for large-scale data. Prior efficient attention methods, such as linear attention, reduce cost through kernelized approximations, yet they can struggle with spatial selectivity and local feature extraction in computer vision. The authors leverage ideas from statistical mechanics to design an efficient attention mechanism that interprets attention through physical concepts and cumulants, aiming to retain expressive relevance modeling while reducing computational overhead. Their method achieves strong performance on CV tasks, outperforming linear attention and approaching softmax attention accuracy.

Method

The authors recast standard scaled dot product attention as a gradient of a statistical mechanical free energy. They consider the tokens $x^l$ as possible states of a random variable $X$ , uniformly drawn from the set of $L$ keys. The query side is treated as an external field, with the field direction determined by the query-key interaction through the usual key and query projections. The partition function is defined as

Z(p) = \langle \exp(p \cdot X) \rangle,

and the log partition function

F(p) = \log Z(p)

plays the role of a negative free-energy potential. Differentiating this quantity with respect to the external field gives the normalized expectation under a Boltzmann-type distribution,

\frac{\partial F(p)}{\partial p} = \frac{\sum_{l=1}^L x^l \exp(p \cdot x^l)} {\sum_{l=1}^L \exp(p \cdot x^l)}.

After applying the value projection and the usual query-key scaling, this derivative reproduces the attention output,

\operatorname{attn}(q) = \frac{\sum_{j=1}^L \exp(q \cdot k^j) v^j} {\sum_{j=1}^L \exp(q \cdot k^j)}.

Thus, attention is interpreted not merely as a learned mixture of value vectors, but as the response of a physical system to an external field. The softmax denominator corresponds to the partition function, and attention corresponds to the derivative of the free-energy potential.

The second central idea is to approximate the free-energy potential using a cumulant expansion. Cumulants are coefficients in the power series expansion of $\log Z(p)$ and represent correlations among random variables. The first cumulant is the mean, the second cumulant is the covariance, and higher cumulants describe genuine higher-order statistical dependence. The expansion takes the form

\begin{aligned} F(p) = F(0) &+ \sum_i p_i \langle X_i \rangle + \sum_{ij} \frac{p_i p_j}{2} \left( \langle X_i X_j \rangle - \langle X_i \rangle \langle X_j \rangle \right) \\ &+ \sum_{ijk} \frac{p_i p_j p_k}{3!} \langle X_i X_j X_k \rangle_c + \cdots \end{aligned}

where the third central cumulant is

\begin{aligned} \langle X_i X_j X_k \rangle_c = &\langle X_i X_j X_k \rangle - \langle X_i X_j \rangle \langle X_k \rangle - \langle X_j X_k \rangle \langle X_i \rangle \\ &- \langle X_k X_i \rangle \langle X_j \rangle + 2 \langle X_i \rangle \langle X_j \rangle \langle X_k \rangle. \end{aligned}

The method truncates this expansion to obtain a controlled approximation to attention. Keeping only the first and second cumulants yields the second-order cumulant attention

\operatorname{attn}_{\mathrm{2nd}} = L \bar{v} + \sum_{l=1}^L \left[ q \cdot (k^l - \bar{k}) \right] (v^l - \bar{v}),

where $\bar{k}$ and $\bar{v}$ denote the average key and value vectors, respectively. This expression replaces the softmax normalization with an additive correction built from centered key-value statistics. The leading correction to this approximation comes from the diagonal part of the third cumulant, producing

\operatorname{attn}_{\mathrm{3rd}} = \operatorname{attn}_{\mathrm{2nd}} + \frac{1}{2} \sum_{l=1}^L \sum_i q_i^2 (k_i^l - \bar{k}_i)^2 (v^l - \bar{v}).

In this formulation, the approximation operates through moments of the key and value embeddings. Second-order attention corresponds to retaining mean and covariance information, which is equivalent to assuming that the underlying statistical variables follow a multivariate normal distribution. The third-order term introduces additional dependence on higher-order embedding correlations and serves as a systematic correction to that Gaussian approximation. A notable property of the cumulant formulation is that the linked-cluster theorem removes the denominator, yielding an additive approximation whose structure is determined by centered key-value moments.

Experiment

Experiments evaluated Cumulant Attention on image and video classification using ViT on ImageNet-100 and ViViT on UCF-101. The results show that cumulant attention consistently outperforms linear attention and approaches softmax attention, reflecting its selective capability and more localized attention patterns. A similar trend is observed for video classification, although the model tends to overfit because of the small dataset size.

Yuto Morimoto Zhipeng Wang Koji Yasuda

Abstract

One-sentence Summary

Key Contributions

Introduction

Method

Experiment

Build AI with AI

HyperAI Newsletters

Command Palette

Cumulant Attention in Vision Transformers (Student Abstract)

Yuto Morimoto Zhipeng Wang Koji Yasuda

Abstract

One-sentence Summary

Key Contributions

Introduction

Method

Experiment

Build AI with AI

HyperAI Newsletters