Transolver works well without transformer attention layers
Does Transolver really need a Transformer?
Machine Learning
Summary
Transolver is a type of neural network used to analyze 3D fluid flows by breaking up data into parts, mixing information, and then putting it back together. The authors found that the usual attention mechanism, which helps the network focus on important parts, can be replaced by a simple linear step without losing accuracy. Instead, the key part is how the data is sliced and recombined multiple times. They also created a faster and more memory-efficient software module to handle the slicing process.
What this means in practice
- •For computational fluid dynamics engineers: Implement efficient neural operator models that omit attention layers but maintain accuracy for 3D fluid simulations.
- •For machine learning infrastructure teams: Use the flashslice kernel to reduce memory and computation costs when processing large unstructured mesh data.
Authors
Shizheng Wen, Siddhartha Mishra
Abstract
The widely used Transolver family of neural operators is based on physics-attention, which softly assigns the points of an unstructured mesh to a small number of slices, applies self-attention among the resulting tokens, and broadcasts the result back to the points. We provide a comprehensive empirical and theoretical analysis to elucidate the mechanisms which are responsible for model performance. To this end, we perform careful ablations on a challenging suite of nine 3D fluid dynamics benchmarks to find that replacing token attention with a constant linear map does not affect the accuracy. Thus, Transolver does not need a Transformer at all. However, removing the global mixing (slicing/deslicing) or doing it only once leads to performance collapse. We leverage the theory of averaging neural operators to explain and corroborate our findings by showing that just slicing/deslicing, in conjunction with pointwise MLPs, already suffices for universal approximation of continuous operators and attention is redundant in this context. Finally, we provide a novel FlashAttention-style efficient implementation of the key slicing/deslicing module of Transolver. This flashslice kernel streams slice and deslice over the points without materializing the heavy slice-weight tensor, while reproducing the best available implementation to floating point error. At the same time, it leads to very significant memory and compute savings, particularly at large slice counts.