Beyond Self-Attention: Sub-Quadratic Vision Transformers for Fast Image Captioning

2026-06-07

AI Summary

This work addresses the high computational complexity and inadequate local feature modeling of vision Transformers in image captioning by proposing a sub-quadratic visual Transformer architecture based on Gaussian Mixture Model (GMM) soft clustering. Instead of conventional self-attention, the method employs an Expectation-Maximization (EM) algorithm to perform semantics-aware soft clustering of image patches, enabling linear-complexity feature aggregation. The resulting visual representations are fed into a GPT-based autoregressive decoder to generate captions. Experimental results on the Flickr30K dataset demonstrate that the proposed model significantly reduces computational overhead while maintaining or even enhancing semantic expressiveness, achieving a dual improvement in both efficiency and performance.

Abstract

Image captioning is a challenging and significant task that aims to generate coherent and semantically meaningful textual descriptions for given images. To accomplish this task, it requires a deep understanding of visual content along with the ability to express that understanding in natural language. Despite remarkable progress with transformer-based architectures, existing approaches often suffer from limitations, such as a lack of rich local feature representations and the high computational cost of quadratic self-attention. The proposed model focuses on improving computational efficiency by restructuring the vision transformer architecture. In designing this approach, the standard self-attention mechanism in Vision Transformers is replaced with a probabilistic transformer approach based on a Gaussian Mixture Model (GMM), a soft-clustering technique. Instead of computing pairwise attention among all image patches, the model groups similar patches into a fixed number of clusters using an Expectation-Maximization (EM) algorithm. This clustering-based mechanism reduces the computational complexity from quadratic O(n^2) to linear O(nK), where K << n. The autoregressive GPT-based decoder is used for caption generation. The model is evaluated on the Flickr 30K dataset, demonstrating competitive and significant improvement over existing works.

Problem

• image captioning
• vision transformers
• self-attention
• computational complexity
• local feature representation

Innovation

• Sub-Quadratic Attention
• Gaussian Mixture Model
• Vision Transformer
• Soft Clustering
• Image Captioning