Cross-Axis Transformer with 3D Rotary Positional Embeddings (2311.07184v3)

Abstract: Despite lagging behind their modal cousins in many respects, Vision Transformers have provided an interesting opportunity to bridge the gap between sequence modeling and image modeling. Up until now however, vision transformers have largely been held back, due to both computational inefficiency, and lack of proper handling of spatial dimensions. In this paper, we introduce the Cross-Axis Transformer. CAT is a model inspired by both Axial Transformers, and Microsoft's recent Retentive Network, that drastically reduces the required number of floating point operations required to process an image, while simultaneously converging faster and more accurately than the Vision Transformers it replaces.

• The paper presents a novel Cross-Axis Transformer that leverages 3D rotary positional embeddings to shift from quadratic to linear complexity in attention computations.
• It employs a hierarchical embedding approach to robustly encode spatial information across image dimensions, thereby enhancing accuracy and convergence speed.
• Empirical results demonstrate that the model achieves nearly 2.8x higher accuracy and reduces training time by half compared to traditional vision transformers.

Overview of the Cross-Axis Transformer with 2d Rotary Embeddings

The paper "Cross-Axis Transformer with 2d Rotary Embeddings" introduces the Cross-Axis Transformer (CAT), a novel architecture intended to enhance the performance and efficiency of Vision Transformers (ViTs). This paper identifies two primary challenges in the current implementation of Vision Transformers: computational inefficiency due to quadratic scaling and inadequate handling of spatial dimensions. These challenges are addressed through the development of CAT, which integrates ideas from Axial Transformers and Microsoft's Retentive Network, leading to reduced computational complexity and improved accuracy and convergence speed in image processing tasks.

Key Contributions

The paper outlines several major innovations embodied in the design of the Cross-Axis Transformer:

1. Reduction in Computational Complexity: CAT achieves linear complexity with respect to the input size, circumventing the quadratic complexity typically associated with Vision Transformers due to the standard implementation of the attention mechanism. This is made possible by employing a cross-axis attention mechanism that effectively attends to each image position with an optimal pass.
2. 2d Rotary Positional Embeddings: Building on the Roformer's rotational embeddings, the paper introduces a hierarchical approach to positional encoding that operates effectively over two dimensions—height and width. This allows the model to maintain and encode spatial information robustly across different scales, advancing the state-of-the-art in image processing.
3. Efficiency Gains Demonstrated with Numerical Results: The architecture demonstrates significant numerical improvements over existing models such as Classic ViT, DINOv2, and BeiT in ImageNet 1k results. CAT achieves nearly 2.8 times higher validation accuracy than baseline vision transformers while reducing training time by half and decreasing the floating point operations required by approximately one third.
4. Empirical Validation and Practical Implications: The paper's experiments exemplify the practical feasibility of achieving state-of-the-art results using consumer-grade hardware and limited computational resources. This expands the accessibility of cutting-edge AI research to independent researchers and smaller institutions, highlighting viable pathways for high-impact research beyond large tech enterprises.

Implications and Future Directions

The implications of this work are substantial for both theory and practice in computer vision and AI:

• Practical Applications: The substantially reduced resource requirements of the CAT model open possibilities for deploying Vision Transformers in resource-constrained environments, including mobile devices and edge computing scenarios.
• Theoretical Developments: The introduction of cross-axis attention could stimulate further research into efficient attention mechanisms for other high-dimensional tasks beyond computer vision, possibly extending into fields such as language processing or multidimensional data analysis.
• Further Research Directions: This paper lays a foundation for larger-scale evaluations of CAT using extensive datasets and advanced computational resources. There's potential to explore the integration of CAT with more complex architectural designs such as hierarchical vision backbones or multimodal transformers.

In conclusion, "Cross-Axis Transformer with 2d Rotary Embeddings" provides a significant contribution to the advancement of Vision Transformer efficiency and effectiveness. By addressing fundamental challenges associated with ViTs, it carves a path for further innovation and broad adoption of transformer models in diverse settings, encouraging a democratization of research capabilities in the AI community.