Unisoma: A Unified Transformer-based Solver for Multi-Solid Systems

Sincerely thanks for insightful comments.

The experimental evaluation could be strengthened by including scaling studies to demonstrate the method's scalability.

Thanks for the insightful suggestion. We conducted a series of scaling experiments to analyze the model’s behavior under varying configurations. Specifically, we evaluated performance across different numbers of slices, processors, channels, and training samples. The default settings used in the paper is: 32 for slice number, 2 for processor number, 128 for channel number and 1000 training samples. We only adjust one variable in one experiment and keep other variables unchanged. The relative L2 recorded in the table is the average accuracy of all physical quantities.

From the results, we observe that increasing the number of slices and channels leads to improved performance up to a certain threshold (approximately slices ≥ 128, channels ≥ 192), beyond which the gains become marginal. This suggests that our model is parameter-efficient, achieving strong performance without requiring excessive capacity. Interestingly, the model shows little sensitivity to depth, with performance remaining largely stable as the number of processors increases. We attribute this to the model's inherently wide architecture, where each layer consists of multiple parallel attention-based modules. This design reduces the model's reliance on depth and makes it less prone to overfitting with increasing layers. Additionally, the relatively flat performance curve across depth may indicate that the model is reaching saturation under the current data regime.

To further investigate this, we conducted data scaling experiments. The results clearly show that model performance consistently improves as the number of training samples increases. Even at the largest scale we tested (1000 samples), the performance continues to rise, indicating that the model has not yet saturated and remains highly data-efficient.

Overall, these findings highlight the robustness and scalability of our model. It delivers competitive accuracy with moderate parameter counts, maintains stability across architectural depths, and continues to benefit from additional data—demonstrating strong practical potential in real-world applications.

Visualizations of the learned Edge Augmented Physics-Aware Tokens would help verify whether the model effectively captures the underlying physical information in the datasets.

To evaluate the effectiveness of the Edge Augmented Physics-Aware Tokens, we visualize the slice weights of two deformable solids in the Bilateral Stamping scenario. The comparison is available at https://anonymous.4open.science/r/unisoma_icml_2025-3DEF

As shown in the figures, a horizontal comparison reveals that different slices attend to different spatial regions after projecting from the original mesh space: some focus on the pressed area under the rigid solid, others on the central stretching region, and some on the fixed ends. This indicates that the tokens successfully group points with similar physical states and differentiate between slices, enabling the model to capture diverse underlying physical behaviors.

More importantly, in the vertical comparison between orderly corresponding slices of the two deformable solids, we observe that they tend to focus on similar regions—areas that are not only of high relevance to each solid individually but are also likely to come into contact. This alignment further supports the effectiveness of Edge Augmented Physics-Aware Tokens in capturing meaningful, structured physical interactions.

Thank you for your positive recognition of our work. We have carefully addressed the concerns you raised and provided additional explanations and experimental results to support our claims. We look forward to your response and feedback.

Sincerely thanks for insightful comments.

Weakness: enforcing physical correctness

We first emphasize that the “explicit modeling” we defined lies not in the explicit PDE constraints used in PINNs or hybrid models, but rather in how the model structure leverages and organizes input information. In PINNs and hybrid models (e.g., NCLaw), PDEs are added to the loss to guide training, often relying on strong assumptions such that explicit PDEs exist—e.g., NCLaw assumes elastodynamic behavior. In multi-solid systems, it is difficult to define a single global and well-studied PDE (like N-S equation in CFD) to describe the entire system accurately. Instead, the behavior is usually governed by multiple local relations, such as contact penalties, load applications and deformations. Directly embedding them into the loss leads to complex, multi-term objectives that are hard to optimize and often limited in applicability. This remains an open challenge in the field.

Therefore, we take a purely data-driven approach and propose to model the physics through the model architecture. Our framework explicitly organizes data using structured modules to capture the underlying physics. This design allows effective learning across diverse materials, object counts, and task types, with consistently strong results. We hope this view can complement PDE-based approaches, and we look forward to future work that combines both in a unified framework.

similar explicit modeling ideas

While NCLaw is a valuable hybrid model combining neural networks with traditional solvers, our method is purely data-driven, learning physical priors through architectural biases. They belong to different classes, as NCLaw depends on classical PDE solvers.

Line 178-185

The x denotes the deep features of an input object. For example, given a rigid solid $u^r_i \in \mathbb{R}^{N_i^r \times C_r}$, we project it using a linear layer: x = Linear($u^r_i)$, where $x \in \mathbb{R}^{N_i^r \times C}$. This is applied to each object individually, with a unified feature dimension $C$. We will clarify this more clearly in the revision. The edge set $E$ is constructed by k-nearest neighbors on mesh points of each object individually, i.e., $E = \text{kNN}(u^r_i)$, where each point connects to its $k$ nearest neighbors. As shown in Appendix F, we find $k = 3 \sim 5$ yields better performance.

Slicing

As shown in Eq. (1), slice weights are computed on deep features $x$ via a linear layer and softmax, and edge weights are derived from connected point weights. The slice is then formed by aggregating point and edge features. Since $x$ comes from mesh points, this applies directly to general meshes.

As noted in Line 201, slices group points with similar physical states, enabling attention to capture physically consistent interactions (Remark 3.1 in Transolver). We will emphasize this more clearly in the revision. When composing a slice from two potentially contacting solids, contact becomes important interactions, while unrelated interactions are suppressed. This focus improves attention modeling and enhances contact capture. Our results confirm that this explicit structure significantly boosts accuracy.

Equations 2, 6

Q, K, and V are computed using standard Transformer linear projections: $Q = W_q(g), K = W_k(g), V = W_v(g)$, with learnable matrices $W_q, W_k, W_v$. In Eq. (2) and (6), $x$ is the composed slice from contacting solids or deformation triplets. As mentioned in the previous response, attention in the slice domain enables the model to learn physically consistent interactions, containing contact and deformation.

Training and Loss

We apologize for the brevity due to space constraints. As shown in Appendix C, our loss functions follow standard setups in related works (e.g., relative L2 for long-time prediction, RMSE for autoregressive simulation). As noted in the first response and Lines 076–089, Unisoma does not use PDE-based losses but learns physical interactions structurally: the contact module, adaptive interaction allocation, and deformation triplet, all built on slice composition. These modules are trained end-to-end, and their physical roles are learned through architecture design.

Physical Meaning of the Contact Module

Unlike PINNs or hybrid models that use PDEs as explicit constraints, we treat contact interactions as learnable physical relationships encoded in high-dimensional features. The model learns to minimize loss by focusing on key interactions, especially contact. We verify its effectiveness through improved accuracy across multiple datasets, particularly in complex multi-contact scenarios, where baselines that model interactions holistically tend to suffer performance degradation due to the holistic interaction mixture. Our structured, explicit modeling helps avoid such issues by organizing information around physically meaningful groupings.

Sincerely thanks for insightful comments.

How to get all the solids pairs that are likely to contact?

In most cases, such as the stamping and grasping scenarios discussed in the paper, the solid pairs that are likely to come into contact are known a priori based on the deterministic nature of the physical setup.

For systems where contact is less certain, as mentioned in Line 368-375, we include solid pairs with a high likelihood of contact in the contact module. The adaptive interaction allocation mechanism then controls the extent to which each contact influences the deformation of a solid, effectively weighting more relevant interactions. This makes the process flexible and robust across both well-defined and uncertain contact settings.

Given a specific system, how to get the "loads"

As a supplement to lines 127–130 in the paper, we clarify that for a moving solid, We treat the displacement $d = u(t+1) - u(t)$ as a load applied over that time interval—commonly referred to as a displacement load [1]. Its positions at time $t$ and $t+1$ are denoted as $u(t)\in R^{N\times 3}$ and $u(t+1)\in R^{N\times 3}$, respectively. We represent each load in a Lagrangian description as the concatenation $\text{concat}(u(t), d)\in R^{N\times 6}$, encoding both its origin position and next movement. When the moving solid comes into contact with others, this displacement load is transferred as a force onto the contacting objects, influencing their deformation.

Is it easy or hard to generalize to different time steps in the autoregressive task?

Similar to MGN[2], our framework learns to predict the next state given the current state in an autoregressive manner. This step-wise prediction scheme does not explicitly encode time intervals, so generalization across different time steps depends on the temporal distribution of the training data and the model’s capacity to learn the underlying dynamics.

To evaluate the generalizability across time steps, we designed an experiment using the Cavity Grasping dataset:

1. We directly test Unisoma, trained on original time step size (600 samples, 100 epochs), on data with doubled time intervals.
2. We fine-tune (20 epochs) the same model using a small amount of doubled-step data (120 samples), then test on data with doubled time intervals.
3. We perform full training (600 samples, 100 epochs) on the doubled-interval data and test accordingly.

In all cases, all other parameters and test data remain consistent)

These results reveal several insights:

• The drop in performance (higher RMSE) when directly applying the model to doubled-step data indicates limited generalization when time step statistics shift significantly.
• However, a small amount of fine-tuning brings considerable improvement, suggesting that the model retains useful representations of the dynamics that are transferable across step sizes.
• Full retraining yields the best performance, as expected, since the model can directly adapt its dynamics modeling to the new time scale.

About the complexity of Unisoma

While attention mechanisms indeed have quadratic time complexity with respect to the number of tokens, our attention operations in Eq. (2) and Eq. (6) are performed over the slices, not directly on the original mesh points. Each slice is corresponding to an Edge Augmented Physics-Aware Tokens. Usually, the number of slices is significantly smaller than the number mesh points.

Let the number of input mesh points be $N$, which are converted to $M$ slice tokens via Eq. (1) and $M \ll N$. This step has a complexity of $\mathcal{O}(MNC)$. The attention in Eq. (2) and Eq. (6) is then applied to sequences of length $M$, with a complexity of $\mathcal{O}(M^2C)$ and is irrelevant to the mesh points $N$. Moreover, the number of modules per layer (contact modules, deformation module, etc) is typically small and the cost is negligible compared to the embedding step. Therefore, the overall complexity is approximately $\mathcal{O}(MNC + M^2C)$. This means that as the input sequence length $N$ increases, the computational cost of our model scales almost linearly ($M \ll N$ and $M$ is usually fixed), making it significantly more efficient than standard attention over full mesh sequences.

Thank you for your positive recognition of our work. We have carefully addressed the concerns you raised and provided additional explanations and experimental results to support our claims. We look forward to your response and feedback.

[1] Mau S T. Introduction to structural analysis: displacement and force methods[M]. Crc Press, 2012.

[2] Pfaff T, Fortunato M, Sanchez-Gonzalez A, et al. Learning mesh-based simulation with graph networks[C]//International conference on learning representations. 2020.

Sincerely thanks for insightful comments.

About the convolution and local information in Transolver

From the paper and official code of Transolver, the conv is applied only to structured meshes or uniform grids (Section 3.1 in Transolver); for irregular meshes—which are the focus of us—Transolver uses linear layers. Moreover, in Transolver, when conv is used, it is for generating deep features and slice weights at point level. The projection into slices is still by globally spatial aggregation over all points, diluting local relations. In contrast, as shown in Table 8, our Edge Augmented Physics-Aware Tokens incorporate neighbor information via mesh edges into the slice projection, preserving local relations. We will revise the sentence to: “However, when tackling irregular meshes, they only consider the point-level features and spatially aggregate points on the whole domain when transforming the mesh points into the slice domain....”

Efficiency comparison with Transolver

As explained in Appendix E, we ensured fair model capacity by aligning the total parameter counts of all baselines (except OOM error), with Unisoma. Efficiency was compared under similar capacity, making the evaluation meaningful. While both Unisoma and Transolver use slice, they differ in real computation. Transolver slices and deslices at every layer and applies FFN() after deslicing on full mesh points, leading to high memory usage as point count grows. Unisoma slices/deslices only once and performs most FFN() within the slice space, greatly reducing memory.

As explained in Line 406-412, Transolver treats all points as a single large tensor of size $N \times C$, applying slicing and FFN() directly, which leads to high memory usage. In contrast, Unisoma processes each object separately with smaller tensors $N_i \times C$, where $\sum N_i = N$. This splits large matrix operations into smaller ones, reducing memory overhead. Frameworks like PyTorch allocate memory for intermediate activations, so large tensors increase peak usage.

We include kNN time at efficiency test, using a KDTree algorithm with O(N log N) complexity. Since kNN is computed per object individually, not over all points, its cost remains small compared to the model’s runtime which is related to the number of all points.

Fairness of comparison

As described in Appendix E, we made our best effort to ensure fairness .We aligned the parameter count and carefully tuned each baseline to achieve better accuracy. Importantly, all models were provided with the same input information, including solid types and load data. For domain-wise models, we treated all deformable solids, rigid solids, and loads as mesh points and concatenated them into a single input sequence. For graph-wise models, we generated mesh edges using the same kNN parameters as Unisoma, ensuring consistency.

We provide more results below (batch size: 1 for deforming plate and 50 for others). It is worth noting that the advantages of Unisoma become more evident as the number of solids and points increases, due to its modular explicit modeling.

"implicit" and "explicit"

We first confirm that all models received the same inputs, except for edges, which were used only where supported. The distinction between "explicit" and "emplicit" we defined lies not in the input itself, but in how the model structure leverages and organizes that information. For example, Transolver concatenates all objects as a single domain-level sequence input and learns interactions implicitly via attention. The model does not explicitly structure the pairwise physical relations (e.g., contact constraints or force) that drive deformation.

In contrast, Unisoma adopts an explicit modeling paradigm, where physical interactions are structurally represented in the model architecture: the contact modules handle potential contacts, the adaptive interaction allocation computes equivalent load and constraint, and the deformation triplet encodes their influence on deformation. This structured design decomposes dynamics into controllable components. We will clarify this distinction more clearly in the revision.

Writing issues

Thanks again for insightful suggestions and we will revise accordingly.

Finally, as stated in the Software and Data (Line 448), we will open data and code upon acceptance to ensure the reproducibility and advance the field.

Sincerely thanks for insightful comments.

Equations 4, 9

Thanks for pointing this out—Equations (4) and (9) do contain errors, and we apologize for the typos. We give a mathematically correct formulation here:

$$z^\xi_j=\frac{\sum_{i=1}^{N^\alpha} w^\alpha_{i,j}x^\alpha_i + \sum_{i=1}^{N^\beta} w^\beta_{i,j}x^\beta_i}{\sum_{i=1}^{N^\alpha}w_{i,j}^\alpha +\sum_{i=1}^{N^\beta}w_{i,j}^\beta}=\frac{(\sum_{i=1}^{N^\alpha}w^\alpha_{i,j})z_j^{\alpha} + \sum_{i=1}^{N^\beta} (w^\beta_{i,j})z_j^{\beta}}{\sum_{i=1}^{N^\alpha}w_{i,j}^\alpha + \sum_{i=1}^{N^\beta}w_{i,j}^\beta}=\theta z_j^{\alpha}+(1-\theta)z_j^{\beta}$$

We make $\theta$ learnable. We test the correct form but it yielded no noticeable performance gain over direct element-wise addition. Hence, we chose direct element-wise addition for slice composition due to its simplicity, efficiency, and performance.

The slice composition is a simple yet effective fusion strategy that works well in practice. Importantly, it allows each solid to maintain its own slice projection parameters, rather than concatenating the two solids and projecting them jointly within each contact module (without parameter sharing across modules). We experiment with this more complex approach, but it results in no significant improvement, despite parameters increasing.

We will correct Equations (4) and (9) and clarify this design in the revision. Once again, we sincerely apologize for it and thanks for the careful observation.

Selection of baselines

We first emphasize the tasks. As claimed in the bold text and corresponding citations in Section 3, our tasks contain:

• Long-time prediction follows Transolver, where the model directly predicts the target state from the given state and loads, usually spanning long time and skipping intermediate steps (e.g., using the 1st frame to predict the 105th frame in Cavity Grasping). It can be formulated as: $P(\hat{x}(t+T)|x{(t)})$, where $T$ spans many time steps. This task needs long-range dependence and is important for fast inference in many industrial scenarios..

• Autoregressive simulation aligns with MGN, focusing on step-by-step rollout. It can be formulated as: $P(\hat{x}(t+1)|x(t))$, $P(\hat{x}(t+2)|\hat{x}(t+1))$, …, $P(\hat{x}(t+T)|\hat{x}(t+T-1))$. It focus more on the rollout trajectory and is suitable for scenarios where intermediate states are essential.

We select baselines based on their recency and relevance to the tasks. For long-time prediction, we compare Unisoma with prevalent domain-wise and GNN-based models. Domain-wise methods are better suited for global inference, while GNNs like MGN suffer from limited receptive fields, though we still include them for completeness. For autoregressive simulation, the task is mainly addressed by GNNs (e.g., MGN, HCMT, HOOD), and domain-wise models have rarely been extended to it. Thus, we focus on the most competitive baselines. For the efficiency comparison (Table 5), we selected models that are highly efficient or widely used in operator learning. GINO and GNO represent domain-wise neural operator models, while OFormer, ONO, and Transolver are recent efficient linear attention-based solvers.

We provide full efficiency comparison here. Notably, due to the difference in input shapes between Euler [B, X, Y, Z, C] and Lagrange [B, (X*Y*Z), C] views, the former is generally more memory-friendly for the same number of points (regular grid), but less suitable for solid system. We use a view mapping for part of baselines (Line 992-999). The main paper compares Lagrange-compatible models and we include Euler-only models (tagged “E”) here. Despite this disadvantage, Unisoma maintains high memory efficiency as point count increases.

RiceGrip Dataset

The DPI-Net repo includes the data generation code scripts. We generate the data on Ubuntu 18.04 with GTX 1080 (CUDA 9.1).

Long-time prediction and autoregressive simulation

As clarified before in Selection of baselines, long-time prediction aims to directly infer the target without intermediate steps, while autoregressive simulation rolls out step-by-step (105 steps for Cavity Grasping/Tissue Manipulation, 120 for Cavity Extruding). These simulation steps align with prior works (e.g., HOOD, HCMT), but our cases involve more solids (Cavity Extruding), making them more challenging and less explored.

Video results

We provide videos at https://anonymous.4open.science/r/unisoma_icml_2025-3DEF

Discussion of TIE

We acknowledge the relevance of TIE and will include a discussion.