6. Transolver & Transolver++

A tokenizer to scale transformers for Physics.

Today’s topic. Image taken from [1].

Introduction

As you probably already know, Transformers have become the go-to method in almost all fields of machine learning, from language and vision to audio, biology, and even physics. We have discussed in previous blog posts GNN- and CNN-like models that work well for physics. However, when it comes to scale, transformers have no real competitors. We know that transformers are data-hungry, and in low-data regimes the previously cited approaches often provide better performance. However, for multiphysics settings (more than one physical task at a time) and foundation models (large “general-purpose” models), most players are now working almost exclusively with transformer architectures.

Today, after briefly recalling how a standard transformer works, we will discuss its main drawback, the quadratic complexity in the sequence length, and how Transolver, and in particular Transolver++, make it usable at scale. The secret, as we will see, lies not in a new attention mechanism, but in the way tokens are created from the input mesh and also GPU-level optimization played its role.

Quick recall of Transformers

The heart of any transformer, and the reason for its success, is the attention mechanism. Given an input sequence X of N elements embedded in dimension d,

\(X = (X_1, \ldots, X_N) \in \mathbb{R}^{N \times d}\)

standard self-attention computes updated representations of each element in the sequence through a convex combination of the sequence elements, where the coefficients are given by self-similarity scores.

A highly simplified version of Attention is as follows:

\(\text{Simplified_Att}(X) = \text{Softmax}(XX^\top) X\)

and so our simplified attention for the elements i is

\(\text{Simplified_Att}(X)_i = \sum_{j=1}^d (X_i^\top X_j) X_j\)

This shows that attention creates a new representation for each element of the input sequence by summing the representations of the other elements, weighted by the scalar product that measures how much two vectors are similar.

In practice, we have omitted a couple of details. First, the coefficients computed via the scalar products are normalized using a Softmax, which is applied row-wise as follows:

\(\text{Softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^d e^{z_j}}\)

And instead of using the original input sequence for the three roles, it is more effective to learn linear pointwise transformations of it:

\(Q = X W_Q, \quad K = X W_K, \quad V = X W_V\)

finally, the full attention becomes:

\(\text{Att}(X) = \text{Softmax}(\frac{QK^\top}{\sqrt{d}})V\)

with a square root of the embedding dimension d to ensure that the output has unit variance.

The main issue of Attention

Attention, is very powerful, but a problematic drawback is that it has quadratic complexity with respect to the sequence length L (and linear complexity with respect to the embedding dimension d)

\(\mathcal{O}(N^2d)\)

Therefore, it would be completely unfeasible to apply the attention operation to a mesh at an industrial scale or to a serious physics problem. Given that today we don’t want to modify the attention mechanism itself, the question is the following:

If we can’t pay attention between mesh points, where can we apply attention?

The answer is, at a high level, simple: we apply attention between (soft-) clusters of mesh points. Transolver does this in a way that I like!

Transolver’s tokenization

Transolver proposes the Physics-Attention mechanism that, instead of treating individual mesh points as tokens (which, according to the authors, may overwhelm the attention mechanism from learning reliable correlations, an argument I am not fully convinced by, although the computational cost is certainly an issue), softly aggregates subsets of the mesh called slices.

By “soft,” I mean that slices do not correspond to disjoint subsets: a given mesh point can be partially assigned to multiple slices, with, for example, half of its value contributing to one slice and the other half to another. If you think of the physical field as a surface over the domain, then a slice corresponds to cutting this volume with a hyperplane (it literally looks like a slice in this interpretation!).

The idea behind Transolver’s tokenization is that it makes more sense to group points into meaningful regions rather than treating each point independently. For instance, when tokenizing a car, one would naturally split it into front, middle, and back areas, instead of assigning a separate token to every single surface point.

Transolver’s tokenization of a CFD simuation on a car. Image taken from [1].

Given a mesh set containing the coordinate information of N mesh points and the observed quantities (the physical fields, such as velocity, pressure, etc.), denoted by X.

\(X = \{X_i\}_{i=1}^N\)

The slices are created by projecting the features x onto M vectors to obtain, for each of the M slices, the portion of the physical quantity that will be assigned to that slice. These values are then normalized with a Softmax:

\(\{ w_i\}_{i=1}^N = \text{Softmax}(\text{Project}(X_i))\}_{i=1}^N\)

Project() is implemented as a pointwise linear layer (nn.Linear in PyTorch). Each slice feature vector is then obtained by averaging the original features X, weighted by the coefficients w, and normalized for stability:

\(Z_j = \frac{\sum_{i=1}^N w_{i,j} X_i}{\sum_{i=1}^N w_{i,j}}\)

Now that we have some nice tokens Z we can apply a standard attention between them:

\(Q, K, V = ZW_Q, ZW_K, ZW_V, \quad Z' = \text{Softmax}(\frac{QK^\top}{\sqrt{d}})V\)

Now that we have updated the representations of the tokens, it is time to project them back onto the original mesh. Transolver does this by broadcasting them using the same coefficients w that were used for tokenization:

\(X_i' = \sum_{j=1}^M W_{i,j} Z_j'\)

A clear visualization of the 3 steps of the Physics Attention is shown in the following image:

3 steps of the Physical Attention of Transolver: Slicing, Attention, Deslicing. Image taken from [1].

Continuity of the tokenization

Since the projection onto fixed vectors is a continuous operation and the physical fields considered are sufficiently smooth, the resulting tokens are continuous rather than disconnected mesh points, as shown in the images, that I think is a good prior for physical fields.

Complexity

The overall complexity consists of computing the N projections, each with cost MC, plus the cost of applying attention to the M tokens, for a total of

\(\mathcal{O}(NMC + M^2C)\)

which is much smaller (depending on M, which is a hyperparameter) than the N^2 cost of applying attention directly to the mesh.

Transolver++

Transolver++ upgrades Transolver with the goal of building an efficient transformer that scales well and addresses a degeneration issue observed in deeper layers. As illustrated in the following image, the slices in deeper layers of Transolver tend to become increasingly uniform.

Left: Slices in deep layers: Transolver++ learns more diverse and concentrated states while Transolver tokenize uniformly. Right: Visualizations of physical quantity change ratio (difference between each point and its neighbors) and slice weights learned by models. The lighter color means faster change. Image taken from [1].

This issue is mitigated through 2 additions. First, they make the temperature of the softmax adaptive, that is, dependent on the input, as follows:

\(\text{Ada-Temp: } \tau = \{\tau_i\}_{i=1}^N = \{\tau_0 + \text{Linear}(X_i) \}_{i=1}^N\)

higher temperature forms a more uniform distribution, while a lower temperature makes the distribution more concentrated on crucial states.

The second modification consists in replacing the softmax-based assignment to the slices (i.e., the tokens) with a mechanism closer to a hard assignment. This cannot be done using an argmax, since it is not differentiable. Therefore, they employ the Gumbel-Softmax to perform differentiable sampling from the categorical distribution that governs the assignment of mesh points to tokens, as follows:

\(\text{Rep-Slice}(X, \tau) = \text{Softmax}\left(\frac{\text{Linear}(X) -\log(-\log(\epsilon))}{\tau}\right), \quad \mathcal{U} \sim (0,1)\)

where epsilon is drawn from a uniform distribution in (0,1) and part with 2 logarithms is called the Gumbel distribution that is usually used to model extreme events and helps to “break ties” between mesh points with similar embedding, effectively alleviating the degeneration problem in the slicing of Transolver.

Gpu Optimization

Transolver++ is further equipped with an extremely optimized parallelism framework. They found out that, contrary to my initial intuition, the bottleneck of Transolver is not the attention but the feedforward layer that embed million-scale points consuming huge GPU memory. As usual with GPU optimization, the bottleneck lies in the communication between different GPUs and the strategy adopted in this work to minimize it, consists in separating the input mesh into multiple GPUs for parallel computing and only communicate when calculating attention among the slices (i.e. tokens).

Conclusions

We discussed one of the most relevent Transformers for Numerical simulations. Differently from the model we discussed in previous blogs, this is not a Neural Operator, meaning that it does not generalize across resolutions. However, this model performs extremely well in every type of data, from regular grids to point clouds, and is able to scale very efficiently (for a transformer). As references, I suggest the official papers [1, 2] as well as for the code [3, 4].

References

1. Transolver: A Fast Transformer Solver for PDEs on General Geometries

2. Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries

3. Transolver Github repo

4. Transolver++ Github Repo