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

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

Problem Statement
Image captioning requires generating coherent, semantically meaningful textual descriptions from visual content, bridging computer vision and natural language processing. While Transformer-based architectures have significantly advanced the field by capturing long-range contextual dependencies through self-attention, they face two primary limitations. First, standard self-attention mechanisms compute pairwise interactions among all image patches, resulting in quadratic computational complexity ($O(n^2)$). This creates substantial memory and computational overhead, particularly for high-resolution images and large-scale applications. Second, existing approaches often suffer from a lack of rich local feature representations and computational redundancy, as they process all patches equally even when many lack semantically critical information.

Methodology
The proposed architecture addresses these limitations by restructuring the Vision Transformer (ViT) framework to replace the conventional self-attention mechanism with a probabilistic clustering approach based on a Gaussian Mixture Model (GMM). The system operates in two main stages:

1. Image Encoding (GMM-Transformer):

   • Patch Embedding: Input images are divided into non-overlapping patches, projected into $D$-dimensional embeddings, and augmented with positional encodings.
   • Context-Weighted GMM Clustering: Instead of computing pairwise attention scores, the model treats patch embeddings as continuous random variables generated by a mixture of $K$ multivariate Gaussian distributions ($K \ll N$).
   • Learnable Context Weights: A sub-layer assigns learnable spatial weights to patches, modulated by statistical properties (mean and variance), allowing the model to prioritize semantically significant regions.
   • EM Algorithm: The Expectation-Maximization (EM) algorithm is employed to perform soft-clustering. In the E-step, the model computes "responsibilities" (soft assignment probabilities) indicating the degree to which each patch belongs to a cluster. In the M-step, cluster parameters (means, covariances, and mixing coefficients) are updated.
   • Complexity Reduction: This process replaces the $N \times N$ attention matrix with an $N \times K$ responsibility matrix, reducing computational complexity from quadratic $O(n^2)$ to near-linear $O(nK)$.
   • Aggregation: Refined patch features are processed through feed-forward networks with residual connections and layer normalization. Finally, a Global Average Pooling operation aggregates the cluster centroids into a compact global feature vector.

2. Language Decoding (Unified GPT-based Decoder):

   • The decoder utilizes an autoregressive GPT-style Transformer.
   • Unified Attention: Unlike traditional encoder-decoder models that separate masked self-attention and cross-attention, this architecture integrates them into a single multi-head attention block.
   • Mechanism: Textual tokens serve as query vectors, while the combined sequence of visual tokens (cluster centroids) and previous textual tokens serves as key and value vectors.
   • Causal Masking: A specialized causal mask ensures that text tokens attend to all visual tokens but only to previously generated text tokens, maintaining autoregressive generation properties.
   • Generation: The output is projected through a Softmax layer to predict the next token in the sequence iteratively.

Key Contributions

• Complexity Reduction: The architecture reduces the time complexity of the transformer encoder from quadratic to near-linear ($O(nK)$) by substituting self-attention with GMM-based soft-clustering, while maintaining the ability to generate contextually coherent captions.
• Efficient Decoder Design: A customized GPT-based decoder is employed that unifies masked self-attention and cross-attention into a single computational block, enhancing multimodal interaction and reducing redundancy.
• Empirical Validation: A comprehensive evaluation on the Flickr30k dataset demonstrates the model's competitive performance against existing state-of-the-art (SOTA) works, balancing computational efficiency with caption quality.

Experimental Results
The model was evaluated on the Flickr30k dataset (using the Karpathy split) with a configuration of 8 transformer layers, 4 attention heads, and a hidden dimension of 768.

• Performance Metrics: The proposed model achieved top scores in advanced semantic metrics: BLEU-4 (0.364), METEOR (0.368), and SPICE (0.199). These results indicate superior phrase structure and scene graph mapping compared to several SOTA models (e.g., TAKG, GNN-Visual, Scene Graph).
• Baseline Comparison: While the model trails the ADA model in baseline precision metrics (BLEU-1: 0.733 vs. 0.794; ROUGE: 0.568 vs. 0.608; CIDEr: 0.723 vs. 0.946), it outperforms other baselines such as TopDown and GNN-Visual. The authors note that the architecture favors deep context over simple replication.
• Efficiency: The model operates with a computational cost of approximately 21.17 GFLOPs and contains roughly 32.8 million trainable parameters. Confusion matrix analysis showed high true positive rates (45,580) and low false positive/negative counts, indicating strong classification capability.

Significance and Claims
The paper claims that the proposed framework successfully addresses the computational inefficiency inherent in standard Vision Transformers. By leveraging the Gaussian Mixture Model and the Expectation-Maximization algorithm, the model achieves a significant reduction in computational complexity from quadratic to near-linear without sacrificing the ability to capture meaningful feature relationships. The authors assert that the model exhibits strong predictive capability and scalable feature encoding, offering an improved balance between caption quality, computational efficiency, and scalability. Future work is directed toward incorporating advanced multimodal pretraining, adapting to domain-specific datasets, and deploying lightweight models for real-time edge applications.