Towards Unified and Lossless Latent Space for 3D Molecular Latent Diffusion Modeling

(W1) On GEOM-Drugs results: JODO should be bold for the Frag metric.

Thank you for the gentle reminder. We have fixed these in the revision.

(W2) On QM9 results: JODO should be bold in Table 2 for the SNN metric. In Table 2, EQGAT-diff should be bold for the AtomStable metric (2D). Overall, most results in Table 2 are very close, especially for the 2D metrics.

Thank you for pointing this out. We have corrected the boldface formatting of Table 2 accordingly: JODO is now bolded for the SNN metric, and EQGAT-diff is bolded for AtomStable (2D). We also agree that most 2D metrics are close, which further emphasizes the necessity to explore beyond 2D geometry and evaluate holistic 3D generative fidelity. We also want to emphasize that our method shows significantly better performance on 3D metrics, especially Bond Length, Bond Angle, and Dihedral Angle on the both datasets, indicating that our method achieves higher accuracy and reliability in 3D molecular generation.

(W3) It is difficult to assess whether the model truly performs better for de novo design, as many scores are quite similar to those of JODO. If the authors are aware of another dataset that better differentiates between methods, additional experiments would be appreciated. On the other hand, since most studies use these datasets, the benchmark might be saturated. Still, the results are interesting, as they suggest that a different modeling paradigm can achieve comparable or better outcomes. I would suggest revising phrasing such as "Breakthrough in Geometric Accuracy," as this claim is not clearly supported by the results in Tables 1 and 2. It is perfectly acceptable that the method does not outperform across all metrics—what is compelling is that a different architecture and training paradigm can achieve competitive results while being more efficient. This, in itself, can be a valuable contribution.

Thank you for the appreciation. We agree and have toned down the statement of "Breakthrough in Geometric Accuracy" to "Leading Geometric Accuracy" to better reflect the empirical findings. While our method does not always dominate the benchmarks, it demonstrates improved or competitive performance under a novel modeling paradigm. We believe the introduction of a unified, lossless latent space—combined with efficient training—constitutes a meaningful contribution to the field.

We also acknowledge the potential saturation of current benchmarks and are actively exploring more challenging and discriminative datasets for future work.

(Limitations) The limitations of the work are not highlighted—either through experiments or in an explicit discussion. The authors should clarify and discuss limitations more clearly.

The limitation section is now in Appendix A due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper once have more space in the camera-ready version.

(Q1) In the section SE(3)-Equivariant Augmentations, the authors could cite other works that use similar data augmentation techniques—for instance, AlphaFold 3.

Thank you for the suggestion. We have included references to prior works such as AlphaFold 3 [2] and MCF [1] that also employ SE(3)-equivariant augmentations in Section 3.1 of the revised manuscript.

(Q2) The manuscript refers to both Lossless and near-Lossless reconstruction. I slightly prefer the term near-Lossless, as any trained method may result in minor reconstruction losses.

We appreciate this semantic clarification. In the revised version, we will consistently use the term near-lossless to acknowledge the minimal but nonzero reconstruction loss.

(Q3) For the sampling time comparisons: did you re-run each model (e.g., GeoLDM), or are the reported sampling times taken from their original papers?

All reported sampling times are obtained by re-running the official implementations of baselines (e.g., GeoLDM, EDM) under the same hardware and sampling settings to ensure fair comparison.

(Q4) I would be interested in further exploration of the learned latent space. For example, does it preserve molecular clusters?

We appreciate your insightful suggestion. We have clustered the molecules in the test set based on their Scaffold and embedded all molecules into the unified latent space, resulting in a t-sne visualization. We observe that molecules with the same scaffolds are relatively close in the latent space, showing that our unified latent space effectively captures molecular structure. Unfortunately, we are unable to provide this visualization in image form at this rebuttal stage, but we will include it in the final version.

(Q5) What variant of Maximum Mean Discrepancy (MMD) do you use? The result depends on the kernel choice.

We follow the same kernel and implementation as in baseline works (e.g., EDM, JODO), using an RBF kernel with a fixed bandwidth. For completeness, we will provide these details in the revised Appendix.

Equivariance through our tailor-made SE(3) augmentations: unless the authors introduce novel augmentations, this phrasing may be overstated, as these data augmentation strategies are fairly common in small molecule and protein design.

We appreciate this note. While our SE(3)-equivariant augmentations follow standard practice, we emphasize that their integration into both encoding and diffusion stages contributes to the equivariance of the unified latent space. We have revised the phrasing to avoid overstatement and better reflect the novelty in usage rather than in augmentation design per se.

* indicates results reproduced – this sentence should not begin with a symbol.

Thank you for pointing this out. We will rewrite these sentences as The results marked with an asterisk (*) are reproduced using official source codes.

Reference:

[1] Swallowing the Bitter Pill: Simplified Scalable Conformer Generation, Wang et al., ICML 2024.

[2] Accurate structure prediction of biomolecular interactions with AlphaFold 3, Abramson et al., Nature 2024.

(W1) The results presented in all the table do not have error bars.

Thank you for the suggestion. We have now added standard deviation as error bars to the reported results in Tables 1, 2, and 3. For clarity, we summarize the results with error bars below (based on 5 independent inference runs):

Table 1 (GEOM-Drugs)

Table 2 (QM9)

(W2) This work only compares with a few discrete graph generation models, and leaves out a few joint modeling papers.

We want to clarify that JODO[1], MiDi[2], and EQGAT[3] are all strong baselines with joint modeling capability.

• JODO propose a joint 2D and 3D graph diffusion model that generates geometric graphs representing complete molecules with atom types, formal charges, bond information, and 3D coordinates.
• MiDi jointly generates molecular graphs and their corresponding 3D arrangement of atoms.
• EQGAT takes continuous atom positions, while chemical elements and bond types are categorical variables, and uses a diffusion model to generate molecules.

(W3) The paper misses important citations in the text/table

• PONITA, Bekkers et al
• FlowMol, Dunn et al,
• MolFM,

We thank the reviewer for pointing these related works, and are happy to include results for PONITA and FlowMol in the Table below for comparison. However, MolFM is not designed for the same benchmark, thus cannot be compared directly, we have included it in the related work session for comparison.

Note that all results are copied from the original papers, and the AtomStable and MolStable reported in PONITA and FlowMol actually correspond to our 2D metrics (following the practice of EDM[4] and JODO[1]). We did not compare other evaluation metrics because they are not reported in these baselines' original paper.

(Limitations) Limitations section is not present in the main paper.

The limitation section is now in Appendix due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper in the camera-ready version.

(Formatting)

• EQGAT model is not cited or referred to in the text.
• The citations listed with a '-', like 1-4, 9-11, do not have access to all the numbers.

We thank the reviewer for catching this. In fact, EQGAT has been correctly cited in the Introduction section, with citation number 10, but we will also cite it in more appropriate places. We note that similar citation ranges are commonly used in NeurIPS papers [5,6]. That said, we agree that explicit citations improve clarity and have revised accordingly. We will correct the ambiguous dash-style citation ranges (e.g., “1-4”) to list individual references.

(Q1) How does the model truly ensure that the equivariance to SE(3) is truly present in the generated samples? In both the augmented and non-augmented cases.

Thank you for the question. You can gently refer to Table 6 in our main paper. We provide RMSD (Root Mean Square Deviation) to measure equivariance in generated samples under different augmented and non-augmented scenarios. The low error validates that our latent diffusion process retains SE(3)-equivariance, as the generated samples closely match the original 3D coordinates.

(Q2) Given that this method relies on a lossless latent space, how can to ensure that this is not overfitting on training data? Could you also report (only) uniqueness scores in Tables 2 and 5?

We appreciate your question. We report reconstruction error on the test sets, in Figure 1 (a;b) and Table 9 in Appendix. These molecules are different from those in the training dataset, eliminating the possibility that our model is overfitting the training dataset.

The uniqueness scores are already reported in Tables 1 and 2, where “V&U” denotes molecules that are both valid and unique. A molecule is defined as unique only if it is also valid (following EDM[4], JODO[1])—otherwise, the notion of uniqueness becomes ill-defined. Our model achieves uniqueness scores of 0.913 on GEOM-DRUGS and 0.973 on QM9, demonstrating its ability to generate a diverse set of valid molecular structures.

(Q3) Could you provide additional details on how augmentation is done and added to the model?

Yes. As described in Section 3.1 of the paper, given a molecule’s 3D coordinates $\mathbf{x}$, we apply a random rigid transformation by sampling a rotation $\mathbf{R} \sim \mathrm{SO}(3)$ and a translation $\mathbf{t} \sim \mathcal{N}(\mathbf{0}, 0.01\mathbf{I})$, and computing the augmented coordinates as $\tilde{\mathbf{x}} = \mathbf{R}\mathbf{x} + \mathbf{t}$. This transformation alters the molecule’s global orientation and position in space but preserves its internal structure. The augmented coordinates $\tilde{\mathbf{x}}$, along with other graph features (i.e., atom types and bonds), are then fed into UAE-3D, which is trained to reconstruct $\tilde{\mathbf{x}}$, thereby encouraging SE(3)-equivariance through augmentation.

(Q4) How to update this model if one of the modalities says bonds have missing data. Would this approach overfit on the bonds present and not generate the bonds not seen before? Is this a good model for molecule discovery, and if so, could you elaborate on this?

We appreciate the insightful question and would like to clarify the following:

• Missing Bond Information can be mostly inferred using xyz2mol or OpenBabel. Bond data can be inferred from the 3D coordinates. If we know the bonding connection is imperfect in the training dataset, we can use tools like xyz2mol and OpenBabel to pre-process the training data to verify and fill-in the missing chemical bonds. This process can handle most cases (>98%) according to the GEOM-DRUGS paper[7] These results should be good enough for a lot of cases and the little error is tolerable, given that the GEOM-DRUGs' ground truth annotation has 1.2% error of unstable molecules.
• Our model can function without bond information completely. If the bond information is missing, we can train a new model without bond information as input and reconstruction target. The VAE can reconstruct the 3D coordinates without explicit bond information. For example, in the CrossDocked2020 dataset which contains protein-ligand complexes, the bond type information is not available for protein's amino acids, but our model can still reconstruct the 3D coordinates of the protein-ligand complexes with high fidelity as shown in the Table below:

(Q5) The tsne plot is unclear to me. If this is in a Center of mass frame, how do I interpret translations? What do these plots contribute?

Thank you for your comments. The t-SNE plot is intended to show the relative positions of real molecules in latent space (not the physical 3D coordinate system) under different coordinate transformations. Before these transformations, the molecules remain in a center of mass frame, but after the translations, they no longer remain centered.

The key insight of this plot is that the latent code of UAE-3D exhibits structural changes with the real geometric motion. This illustrates the meaningful structural changes in the molecular representation in SE(3)-equivariant augmentation.

References:

[1] Learning Joint 2-D and 3-D Graph Diffusion Models for Complete Molecule Generation, Han et al., TNNLS 2024.

[2] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation, Vignac et al., ECML 2024.

[3] Navigating the Design Space of Equivariant Diffusion-Based Generative Models for De Novo 3D Molecule Generation, Le et al., ICLR 2024.

[4] Equivariant Diffusion for Molecule Generation in 3D, Hoogeboom et al., ICML 2022.

[5] Convolutional Differentiable Logic Gate Networks, Petersen et al., NeurIPS 2024.

[6] CSPG: Crossing Sparse Proximity Graphs for Approximate Nearest Neighbor Search, Yang et al., NeurIPS 2024.

[7] GEOM, energy-annotated molecular conformations for property prediction and molecular generation, Axelrod et al., Scientific Data 2022.

Thank you again for your detailed and constructive review. We deeply appreciate the time and expertise you invested in evaluating our work.

We would like to kindly invite you to take a look at our rebuttal, in which we have addressed each of your comments thoroughly. In particular:

• We have added standard deviation/error bars to all the major result tables, as you suggested, based on 5 inference runs.
• We have clarified that our method already compares with recent joint modeling baselines (JODO, MiDi, EQGAT), and we have further added comparisons with other recent baselines such as PONITA and FlowMol wherever applicable.
• We have elaborated on how our model handles SE(3)-equivariance, near-lossless latent space, the data augmentation strategy, and cases with missing bond information in detail.
• We also explained the design and intention behind the t-SNE plot in latent space.

We would be very grateful if you could share any follow-up thoughts or questions you may have—we are happy to engage further in discussion and clarify anything that remains unclear.

Thank you again for your valuable time and feedback!

(W1) This is a fast moving field and the SoA has advanced beyond the presented baselines (e.g., SemlaFlow outperforms the presented work both in terms of accuracy and speed).

(Q4) What advantages does your approach have over recently described flow matching approaches?

We thank the reviewer for pointing out recent advances. In response, we have added a comparison to SemlaFlow, a recent baseline. As shown in the following Table, our model outperforms SemlaFlow in both sampling speed and generation quality. The results are reproduced using the official code and pre-trained model checkpoints and evaluated with the same evaluation metrics as our model. Note that the missing metrics in SemlaFlow are due to the frequent NaN results in the calculation of related metrics caused by the existing errors in the generated molecules, making it difficult to compute reasonable reporting numbers. We will further organize and analyze the results of SemlaFlow, and include this comparison and analysis in the revised version.

Table: 2D Evaluation metrics on the QM9 dataset.

Table: 3D evaluation metrics on the QM9 dataset.

Table: Sampling speed comparison on the QM9 dataset.

Our approach offers two key advantages over recent flow-matching models such as SemlaFlow:

• Better Generation Quality: Our model achieves significantly better generation quality across multiple metrics, including FCD, AtomStable, MolStable, and others. This is due to our unified latent space design that captures both 2D and 3D molecular features effectively.
• Improved sampling speed: Our model achieves 10x faster sampling compared to SemlaFlow, as shown in the above table. This is due to our latent diffusion model (LDM) design, which allows for efficient sampling in a unified latent space. In contrast, flow-matching models typically require more complex sampling procedures that are slower and less efficient.

(W2) The performance of the model is repeatedly attributed to the unified latent space, but no evidence is provided supporting this assertion.

(Q3) Can you design an ablation study or other experiment that supports your claim the accuracy of your model is due to the unified latent space instead of any of the other architectural and training procedure choices you made?

Thank you for this valuable suggestion. To better validate the contribution of the unified latent space, we have conducted an additional ablation experiment to compare our model with its variant which have separate latent space for equivariant (3D) and invariant (2D) modalities. Specifically, we split the unified VAE into two independent neural networks: one maintaining the R-Trans architecture and taking 2D information (atom types, bond types) as input, while the other uses a vanilla transformer structure and takes 3D information (atom types and coordinates) as input. Each network generates a separate latent space, with dimensions set to half of the original hyperparameter. The LDM is then trained on the concatenated latent space of these two networks.

Table: 2D Evaluation metrics on the QM9 dataset.

Table: 3D evaluation metrics on the QM9 dataset.

These results support our claim that the unified latent space plays a key role in our improved generation quality. We attribute this improvement to the early fusion of 2D and 3D information in the VAE encoder, which allows the LDM to model a single coherent modality. In contrast, for the separated latent spaces, the LDM needs to learn over two distinct modalities and their interactions, which increases the complexity of generation.

We will include this ablation and discussion in the revised manuscript.

(Q1) How do you achieve near zero reconstruction with a VAE? The latent space is variational so one would expect sampling from the generated embedding to produce different outputs. Do you only evaluate the mean? What are sampled outputs like?

Thank you for the insightful question.

We evaluate the VAE posterior's sampled latents instead of mean. Our VAE setup follows the design in Stable Diffusion [1], where we use a very small KL regularization coefficient (1e-8; see Appendix Table 11) to prioritize reconstruction fidelity over latent regularization. This results in a low-variance posterior distribution ($\mathrm{e}^{-8.1786}$), where samples from the posterior remain close to the mean. Consequently, the decoder can produce stable and high-fidelity reconstructions from sampled latents.

Importantly, we do not remove the KL term entirely. While doing so still allows for accurate reconstruction, it significantly degrades the generation performance of the LDM, as shown in our ablation below. We attribute this to a key advantage of VAEs over AEs for generative modeling: the variational posterior ensures a structured and smooth latent space, which is critical for generation.

Although we did not explicitly evaluate using the posterior mean as input to the LDM, we expect this to yield similarly poor results based on the observed impact of removing KL regularization.

Table: Reconstruction error on the QM9 dataset.

Table: 2D Evaluation metrics on the QM9 dataset.

Table: 3D evaluation metrics on the QM9 dataset.

(Q2) To what extent is this model permutation invariant?

Our model is permutation equivariant, not invariant, in both the VAE and LDM components. This is because both modules produce atom-wise outputs and are implemented using Transformer-based architectures. Since the Transformer is inherently permutation equivariant—meaning that permuting the input atoms results in a corresponding permutation of the output—the overall model preserves this property.

References:

[1] High-Resolution Image Synthesis with Latent Diffusion Models. Rombach et al., CVPR 2022.

(W1) The paper relies heavily on empirical demonstrations, lacking formal theoretical analysis of why the unified latent space retains all necessary information for 3D generation, or convergence guarantees for the learning of equivariance solely via data augmentation. While the gains are empirically solid, more discussion or theoretical intuition would improve the scientific contribution.

Thank you for the constructive feedback. 3D molecular generation is based on data consisting of atomic types, chemical bonds, and 3D coordinates. Our method integrates this information into a unified latent space, allowing the model to learn all necessary geometric and chemical features in one space. While we currently lack rigorous theoretical analysis, we provide the following intuitive explanations:

1. Near-lossless Latent Space: The unified latent space is designed to be near-lossless, meaning it can perfectly reconstruct the original 3D coordinates and chemical information. As shown in Figure 2, our VAE achieves 100% atom/bond accuracy and near-zero coordinate RMSD on the test set. This ensures that all necessary information is preserved during encoding and decoding.
2. SE(3)-Equivariant Augmentation: By applying SE(3) augmentations during training, the model learns to be equivariant to rotations and translations. This means that the model can generate 3D coordinates that are consistent with the original molecular structure. To validate this, we conducted experiments on 1000 molecules in the test set, applying random SE(3) transformations (rotations and/or translations) and measuring the RMSD of the molecules reconstructed by the model. As shown in the following Table, The results show that the model maintains low RMSD under different SE(3) transformations, confirming its SE(3) equivariance.

(W2) The ablation studies in Section 4.5, Table 5, and Table 6 are thoughtful but could be further strengthened. For example, the effect of latent dimension size, number of R-Trans layers, or replacing the DiT backbone with other modern diffusion architectures could be explored to validate robustness and architectural choices beyond what is already discussed.

(Q2) What is the effect of varying the latent space dimension, depth of R-Trans encoder, DiT architecture (vs other transformer-based or non-transformer LDM backbones), or hyperparameters? Is the model robust to such design choices or sensitive in particular to one?

We thank the reviewer for the insightful questions. We provide the following studies and new results:

1. Latent Dimension Ablation

   As already shown in Appendix Table 7, we observe that too low a latent dimension fails to capture the necessary information, leading to poor performance. Increasing the dimension to 16 significantly improves performance, while further increasing it to 32 does not provide significant additional improvement. And in the updated Table 7 below, we also provide the reconstruction error for QM9 dataset for more clarity:

   Latent Dim	Atom Acc	Bond Acc	RMSD	AtomStable	MolStable	V&U&N	Bond Len$\downarrow$	Angle$\downarrow$	Dihedral$\downarrow$
   4	0.9201	0.8932	0.080	0.914	0.325	0.566	3.27E-01	2.20E-01	4.19E-03
   8	0.9998	0.9746	0.006	0.987	0.882	0.923	1.29E-01	1.77E-02	8.14E-04
   16	1.0000	1.0000	0.002	0.999	0.988	0.950	7.04E-02	9.84E-03	3.47E-04
   32	1.0000	1.0000	0.003	0.986	0.867	0.918	9.98E-02	1.42E-02	7.77E-04

   Likewise, dimensions 4 and 8 yield higher reconstruction error, indicating they are insufficient to capture molecular complexity. The latent dimension of 16 achieves near-zero reconstruction error, while 32 does not significantly improve performance and even slightly increases the error, indicating that too large a latent dimension may not be beneficial for reconstruction.

2. R-Trans Layer Depth Ablation (VAE reconstruction)

   R-Trans layers	Atom Acc	Bond Acc	RMSD
   3	1.0000	0.9999	0.00326
   6	1.0000	1.0000	0.00203
   9	1.0000	1.0000	0.00229
   12	1.0000	1.0000	0.00209

   We observe that decreasing the number of our transformer layers to 3 leads to a significant increase in reconstruction error, while increasing it to 9/12 can not further improve the reconstruction quality because the reconstruction error is already near-zero.

3. Backbone Ablation

   As shown in Table 4, we already compared our DiT backbone with a standard Transformer architecture. In the updated Table 4 below, we add a comparison with PerceiverIO, a modern transformer-based architecture designed for structured inputs and outputs. The results show that DiT consistently outperforms PerceiverIO across all metrics, which we attribute to its adaptive LayerNorm layers. These layers enable DiT to effectively handle data with different noise scales, thereby improving diffusion performance.

   Structure	3D-AtomStable	2D-AtomStable	V&C	V&U	V&U&N
   Transformer	0.983	0.997	0.938	0.922	0.922
   PerceiverIO	0.972	0.990	0.933	0.931	0.931
   DiT	0.993	0.999	0.983	0.973	0.950

   Moreover, when paired with other diffusion neural architectures, like Transformer and PerceiverIO, our model still achieves meaningful and comparable performances. The performance can also be potentially improved given more hyperparameter tuning, which is unfortunately not available given rebuttal's limited time. This observation demonstrate that our UAE-3D model is robust to different diffusion neural architectures.

(W2) It would also be instructive to see whether the benefits of unified latent space and data augmentation persist across substantially larger or more complex molecules. (W3) The paper claims broad applicability of the unified latent approach to diverse 3D molecular and multimodal datasets, but empirical validation is limited to QM9 and GEOM-Drugs. Explicit discussion of the limitations in generalization, scalability to much larger chemical spaces, or application to proteins/RNAs is somewhat deferred to future work.

Thank you for the suggestion. We agree that broader empirical validation is important. While QM9 and GEOM-Drugs are the standard benchmarks, we provide the following evidence to support scalability:

• GEOM-Drugs includes molecules up to 181 atoms, which are comparable in size to small peptides. Our method achieves perfect reconstruction on this dataset (see Table 9), showing robustness to large-scale molecular inputs.

• CrossDocked2020 is a large-scale dataset of protein-ligand complexes, containing over 100,000 protein-ligand pairs and up to over 400 atoms per complex. We trained our VAE on molecules from this dataset, including amino acids and ligands, achieving:

  Atom Acc	RMSD(Å)
  1.000	0.011

These results show that UAE-3D can scale to larger, protein-related molecules and maintain high fidelity. We did not test bond type prediction accuracy yet because the bond type information is not available in the dataset for protein's amino acids. Adapting to full protein/RNA generation will be addressed in future work, as discussed in our Appendix A.

(Q1) Could the authors provide more formal justification (or empirical evidence across much more diverse datasets) for the claim that SE(3)-equivariant augmentation paired with the unified VAE suffices to ensure geometric and chemical preservation, as opposed to specialized architectures?

Thank you for the suggestion. We are happy to provide empirical results on more diverse datasets to support our claims. As shown in the above table, we have conducted experiments on the CrossDocked2020 dataset, which contains large protein-ligand complexes. The results indicate that our model maintains low reconstruction loss and can preserve complex geometric and chemical information even on large-scale protein-ligand molecules. We are also conducting further validation experiments on larger PDB datasets to strengthen our claims.

(Q3) Are there patterns to the errors or invalid/implausible generations outside of what is summarized quantitatively in Tables 1/2? Could the authors expand with more systematic failure analysis (qualitative or quantitative)?

We appreciate the suggestion and offer the following observations:

• Error Types: During generation, failure cases often involve incorrect aromatic bonds, which are particularly sensitive to small coordinate perturbations. This occurs rarely but can lead to invalid valence assignments. Specifically, because the interatomic distance (i.e., bond length) in the generated molecules determines the bond type, for example, the carbon-carbon bond length in the benzene ring is 1.40Å, which is between a single bond (1.54Å) and a double bond (1.34Å). Therefore, during the generation process, these bonds are likely mismatched due to small coordinate shifts, resulting in incorrect bond assignments for some carbon atoms.
• Quantitative Insights: Across 10k samples, invalid molecule rate remains around 1%, and ring-distribution (Table 10) closely matches that of the training set, indicating no mode collapse.

(Limitations) The limitations section is in the appendix, but should be summarized more explicitly in the main paper.

The limitation section is now in Appendix due to our limited space in main paper. We appreciate the advice, and will revise to include it in the main paper in the camera-ready version.

[1] Perceiver IO: A General Architecture for Structured Inputs & Outputs, Jaegle et al., ICLR 2022.