Sampling Foundational Transformer: A Theoretical Perspective

TL;DR

Sampling Foundational Transformer (SFT) addresses the quadratic bottleneck and modality-specific design challenges of traditional transformers by introducing a context-aware sampling-without-replacement mechanism and a pseudoconvex transformer layer. The method uses differentiable sampling based on learnable importance scores and a leaky, Maxout-based attention with relative positional encoding to achieve linear-time complexity $O(n)$ per layer while supporting point clouds, graphs, and sequences. The paper provides theoretical results on the pseudoconvexity and gradient bounds of the SFT layer, and demonstrates competitive performance across dense low-rank and sparse high-rank tasks, including rotational invariance in point clouds and long-range benchmarks in sequences and graphs; it also discusses rank-injection phenomena and practical runtime advantages. Overall, SFT offers a scalable, multi-modality foundation for transformers with improved convergence and efficiency, enabling broader applicability of foundational models across data modalities.

Abstract

The versatility of self-attention mechanism earned transformers great success in almost all data modalities, with limitations on the quadratic complexity and difficulty of training. To apply transformers across different data modalities, practitioners have to make specific clever data-modality-dependent constructions. In this paper, we propose Sampling Foundational Transformer (SFT) that can work on multiple data modalities (e.g., point cloud, graph, and sequence) and constraints (e.g., rotational-invariant). The existence of such model is important as contemporary foundational modeling requires operability on multiple data sources. For efficiency on large number of tokens, our model relies on our context aware sampling-without-replacement mechanism for both linear asymptotic computational complexity and real inference time gain. For efficiency, we rely on our newly discovered pseudoconvex formulation of transformer layer to increase model's convergence rate. As a model working on multiple data modalities, SFT has achieved competitive results on many benchmarks, while being faster in inference, compared to other very specialized models.

Figures (8)

• Figure 1: Our proposed sampling-based transformer variant. (PreLN) indicates LayerNorm if Pre-LN Transformer variant, otherwise, identity. (PostLN) indicates LayerNorm if Post-LN Transformer variant, otherwise, identity. Red (our model) and blue (vanilla transformer) highlight the difference between our model and vanilla transformer (with attention bias). Sampler is described in Figure \ref{['fig:dswr']}.
• Figure 2: Our proposed differentiable sampling without replacement.
• Figure 3: Rank progression of token representation with 256 tokens and embedding size 512 through 100 randomly initialized single head SFT layers with leaky probability function and pairwise maxout attention nonlinearity.
• Figure 4: Comparison between a convex function $z = \frac{x^2+y^2+10}{10}$ (top-left) and a pseudoconvex function $z = \frac{10(x^2+y^2)}{x^2+y^2+1}$ (top-right) and their corresponding heatmaps with gradient vector fields (bottoms). The pseudoconvex function greatly resemble its convex counterpart regarding the search for the local minima.
• Figure 5: Inference time (in seconds) of the networks for ModelNet40 classification test split in 1 A100 GPU and 8 CPUs with a batch size of 32. The Vector Neuron VectorNeurons and Canonicalization LearntEquivariance framework is applied in PointNet PointNet and DGCNN DynamicGraphCNN. The leaky attention function is applied to our SFT model. Our model in this plot does not use compile method from PyTorch to speed up for fair comparison (when using, it halves the inference time on A100). The results of others are taken from LearntEquivariance.

Theorems & Definitions (17)

• Definition 4.1
• Theorem 4.2
• Theorem 4.3
• Definition B.1
• Theorem B.2
• proof
• Theorem B.3
• Lemma B.4
• Theorem B.5
• proof