🤖 machine learning

RelFlexformer: Efficient Attention 3D-Transformers for Integrable Relative Positional Encodings

The paper introduces RelFlexformer, a class of efficient 3D-Transformer models that utilize Non-Uniform Fourier Transform (NU-FFT) to integrate arbitrary integrable relative positional encodings with $O(L \log L)$ complexity, thereby enabling effective attention mechanisms for both structured grids and unstructured heterogeneous 3D data like point clouds.

Original authors: Byeongchan Kim, Arijit Sehanobish, Avinava Dubey, Min-hwan Oh, Krzysztof Choromanski

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Byeongchan Kim, Arijit Sehanobish, Avinava Dubey, Min-hwan Oh, Krzysztof Choromanski

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive, chaotic party where guests are standing in a 3D room. Some are close together, some are far apart, and there's no neat grid or rows. Your goal is to figure out who should talk to whom based on how close they are to each other.

In the world of computer vision, this "party" is a 3D Point Cloud (a collection of dots representing objects in space), and the "guests" are data points. The computer uses a tool called a Transformer to decide who talks to whom.

Here is the problem the paper solves, explained simply:

The Problem: The "Quadratic" Bottleneck

Standard Transformers are like a party host who insists on checking the distance between every single guest and every other guest before anyone can speak.

If you have 100 guests, the host makes 10,000 checks.
If you have 1,000 guests, the host makes 1,000,000 checks.
If you have 10,000 guests (common in 3D scans), the host gets overwhelmed, runs out of memory, and the party stops. This is the quadratic cost ( $O(L^2)$ ).

To fix this, researchers invented "Efficient Transformers" (like Performers). These hosts use a shortcut: they guess who should talk to whom using a clever math trick, making the party run fast ( $O(L)$ ).

The Catch: These shortcuts are great at speed, but they are terrible at understanding geometry. They forget that in a 3D room, distance matters. They treat a guest standing next to you the same as one standing across the room, which ruins the accuracy for 3D tasks.

The Solution: RelFlexformer

The authors introduce RelFlexformer, a new type of efficient host that is both fast and geometrically aware.

Think of it like this:

The Shortcut: Instead of checking every pair, it uses a "magic lens" (called NU-FFT, or Non-Uniform Fast Fourier Transform) to instantly calculate how distance affects conversation.
The Flexible Modulation: Imagine the host has a special "volume knob" for every pair of guests. If two guests are close, the volume is high (they talk a lot). If they are far, the volume is low.
- Old efficient methods couldn't turn this knob without breaking their speed shortcut.
- RelFlexformer can turn this knob for any shape of room (even messy, irregular 3D spaces) without slowing down. It does this by translating the distance problem into a frequency problem (like turning a complex song into a simple melody) and solving it quickly.

The "Magic" Analogy: The Orchestra

Imagine the 3D points are musicians in an orchestra scattered randomly in a hall.

Standard Transformers ask every musician to listen to every other musician to find the harmony. This takes forever for a big orchestra.
Old Efficient Transformers ask the musicians to just guess the harmony based on a simple rule. It's fast, but the music sounds flat because they ignore who is sitting next to whom.
RelFlexformer is like a conductor who uses a special acoustic mirror. Instead of asking everyone to listen to everyone, the mirror instantly reflects the sound waves so that musicians know exactly how loud to play based on their specific distance from others. It preserves the "spatial feel" of the room but keeps the rehearsal time short.

What They Claim to Achieve

The paper claims that by using this "acoustic mirror" (the NU-FFT math):

Speed: It remains fast, scaling almost linearly ( $O(L \log L)$ ) even with huge amounts of data. It doesn't crash when the "party" gets too big.
Accuracy: It brings back the missing "distance sense." In tests on 3D object recognition (like identifying a chair from a cloud of dots) and 3D segmentation (labeling parts of a room), RelFlexformer:
- Beat the old "fast but dumb" methods (Performers).
- Often beat the slow but smart standard Transformers, even though it was much faster.
Versatility: It works on messy, real-world data like point clouds (from LiDAR scanners) and RGB-D images (cameras that see depth), where points aren't in neat grids.

Summary

RelFlexformer is a new way for computers to look at 3D shapes. It combines the speed of a shortcut with the precision of understanding physical distance. It allows computers to process complex 3D scenes (like a robot navigating a room) quickly without losing the ability to tell how far apart things are.

Technical Summary: RelFlexformer

Problem Statement

Attention-based Transformers have become dominant architectures for 3D learning, particularly for point cloud data. However, standard Transformers face two critical limitations in this domain:

Computational Bottleneck: The standard self-attention mechanism scales quadratically ( $O(L^2)$ ) with sequence length $L$ . This becomes prohibitive for large point sets and dense geometric models.
Geometric Misalignment: Standard attention does not naturally align with the unordered, irregularly spaced nature of point clouds. While Relative Positional Encoding (RPE) methods exist to inject geometric bias, integrating them into efficient, linear-time attention mechanisms (like Performers) is non-trivial. Existing efficient RPE methods often rely on structured token layouts (e.g., regular grids) or require explicit construction of $L \times L$ masks, which negates the computational benefits of linear attention. Specifically, inserting a geometric mask into kernelized attention without materializing the full attention matrix requires the mask to admit fast matrix-vector products, a constraint few existing methods satisfy for arbitrary 3D distributions.

Methodology

The authors propose RelFlexformer, a class of efficient 3D Transformers that integrates a broad family of RPEs via arbitrary integrable modulation functions $f$ . The core innovation is the application of Non-Uniform Fast Fourier Transform (NU-FFT) theory to enable sub-quadratic masked attention.

Key Technical Components:

Modulation via Convolution: The RPE mechanism is formulated as a spatial convolution. Given input token coordinates $\{r_i\}$ and a modulation function $f$ , the masked attention value $w$ for a token is computed as the convolution of $f$ with a signal $P$ (a sum of Dirac deltas weighted by input features).
NU-FFT Formulation: By the Convolution Theorem, convolution in the spatial domain corresponds to multiplication in the Fourier domain. The authors approximate the continuous Fourier integral using a quadrature rule with $S$ $S$ samples.
- Forward Step: Compute the Fourier Transform of the point cloud signal at sampled frequencies $\xi_s$ using NU-FFT. This takes $O(L \log L)$ time.
- Modulation: Multiply the transformed signal by the Fourier Transform of the modulation function $Ff(\xi_s)$ .
- Inverse Step: Compute the inverse NU-FFT to evaluate the modulated values at the original token coordinates. This also takes $O(L \log L)$ time.
Integration with Linear Attention: This "FastMultM" algorithm replaces the explicit $L \times L$ mask matrix multiplication in linear attention frameworks (like Performers). The overall complexity becomes $O(L \log L)$ , preserving the efficiency of linear attention while incorporating arbitrary geometric biases.
Generalization: The framework generalizes existing methods. For instance, if the modulation function is parameterized as discrete harmonic functions, the method analytically recovers Rotary Position Embedding (RoPE). It also subsumes the $\otimes$ -STRING encoding as a special case where frequencies are drawn from a specific spectral density.

Key Contributions

General Formulation: The authors formulate a general class of 3D RPE-modulated efficient attention mechanisms based on arbitrary $L_1$ -integrable functions.
Efficient Implementation: They derive an $O(L \log L)$ implementation using NU-FFT, enabling fast masked attention without materializing the attention matrix. This allows geometric information to be incorporated directly into Performer-style attention.
Versatility and Performance: The paper validates RelFlexformer as a drop-in replacement across diverse backbones (PCT, PTv3, DFormer). It demonstrates that RelFlexformer not only bridges the performance gap between efficient linear attention and dense quadratic attention but, in several cases, surpasses standard $O(L^2)$ Transformers.
Complementarity: The authors show that their distance-based RPE modulation is orthogonal to token-level encodings like PointRoPE and can be effectively combined for further gains.

Experimental Results

The authors evaluated RelFlexformer on a large portfolio of 3D datasets for classification and segmentation:

Object Classification (ModelNet40, ScanObjectNN):
- On ModelNet40, RelFlexformer achieved 92.94% accuracy, significantly reducing the gap to the dense Transformer baseline (93.2%) compared to vanilla Performer (92.34%).
- On ScanObjectNN, RelFlexformer (84.45%) surpassed the dense Transformer baseline (84.0%).
Semantic Segmentation (ScanNet, ScanNet200, ScanNet++, nuScenes, S3DIS):
- RelFlexformer consistently improved upon vanilla Performer baselines. For example, on ScanNet, mIoU improved from 74.8% (Performer) to 76.8% (RelFlexformer), approaching the Transformer baseline (77.6%).
- On nuScenes, RelFlexformer achieved 80.3% mIoU, surpassing the Transformer baseline (80.4% is the Transformer score, but the text notes RelFlexformer often outperforms it; specifically, on nuScenes, RelFlexformer + PointRoPE reached 81.2%, exceeding the Transformer's 80.4%).
- On S3DIS 6-fold cross-validation, RelFlexformer closed the mIoU gap from 73.18% (Performer) to 76.25%, nearly matching the Transformer (77.70%).
RGB-D Segmentation (NYU Depth v2, SUN RGB-D):
- Using the DFormer backbone, RelFlexformer achieved 55.32% mIoU on NYU Depth v2 and 51.04% on SUN RGB-D, outperforming both the Performer baseline and the recently proposed STRING method, while avoiding STRING's parameter overhead.

Significance and Claims

The paper claims that RelFlexformer provides a unified geometric operator that subsumes recent multidimensional encodings (like STRING) as special cases while extending sub-quadratic efficiency to dense RGB-D data via coordinate lifting.

The primary significance lies in closing the accuracy gap between efficient linear attention and standard dense attention. The authors argue that by framing geometric modulation via NU-FFT, it is possible to incorporate a broad family of integrable RPE functions without sacrificing the computational benefits of kernelized attention. This establishes RelFlexformer as a highly scalable architecture for processing large-scale, irregular point clouds that were previously bottlenecked by the quadratic cost of standard dense attention. The results suggest that efficient attention mechanisms can achieve, and sometimes exceed, the performance of dense Transformers when equipped with appropriate geometric modulation, making them viable for real-world 3D perception tasks where latency and memory are critical constraints.