Towards All-Atom Foundation Models for Biomolecular Binding Affinity Prediction
Authors: Liang Shi
Zuobai Zhang
Huiyu Cai
Santiago Miret
Zhi Yang
Jian Tang
Venue: ICLR 2026 Poster (2026)
Keywords: Biology foundation model
biomolecular interaction prediction
representation learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Biomolecular interactions play a critical role in biological processes. While recent breakthroughs like AlphaFold 3 have enabled accurate modeling of biomolecular complex structures, predicting binding affinity remains challenging mainly due to limited high-quality data. Recent methods are often specialized for specific types of biomolecular interactions, limiting their generalizability. In this work, we repurpose AlphaFold 3 for representation learning to predict binding affinity, a non-trivial task that requires shifting from generative structure prediction to encoding observed geometry, simplifying the heavily conditioned trunk module, and designing a framework to jointly capture sequence and structural information. To address these challenges, we introduce the **Atom-level Diffusion Transformer (ADiT)**, which takes sequence and structure as inputs, employs a unified tokenization scheme, integrates diffusion transformers, and removes dependencies on multiple sequence alignments and templates. We pre-train three ADiT variants on the PDB dataset with a denoising objective and evaluate them across protein-ligand, drug-target, protein-protein, and antibody-antigen interactions. The model achieves state-of-the-art or competitive performance across benchmarks, scales effectively with model size, and successfully identifies wet-lab validated affinity-enhancing antibody mutations, establishing a generalizable framework for biomolecular interactions. Our open-source implementation is available at https://github.com/VectorShi/ADiT.

Published as a conference paper at ICLR 2026
TOWARDS ALL-ATOM FOUNDATION MODELS FOR BIOMOLECULAR BINDING AFFINITY PREDICTION
Liang Shi!
Zuobai Zhang 23,
Huiyu Cai
Santiago Miret*
Zhi Yang!: 1.
Jian Tang256.1
'Scooo of Computer Science, Peking University ?Mila Québec Al Institute
UUniversité de Montréal "Intel Labs
HHEC Montréal
OCIAAR AI Chair
liang. swi@stu.uuu.s..e pku.uu.ppued.. zuobai. zhang@mila quebec yangzhi@pku.er i jian. tan.tangheec.a
ABSTRACT
Biomolecular interactions play critical role in biological processes. While recent breakthroughs like AlphaFold have enabled accurate modeling of biomolecular complex structures, predicting binding affinity remains challenging mainly due limited high-quality data. Recent methods are often specialized for specific types of biomolecular interactions, limiting their generalizability. this work, we adapt AlphaFold 3 for representation learning to predict binding affinity, non-trivial task that requires shifting from generative structure prediction to encoding observed geometry, simplifying the heavily conditioned trunk module, and designing a framework to jointly capture sequence and structural information. To address these challenges, we introduce the Atom-level Diffusion Transformer (ADiT), which takes sequence and structure as inputs, employs a unified tokenization scheme, integrates diffusion transformers, and removes dependencies multiple sequence alignments and templates. We pre-train three ADiT variants on the PDB dataset with a denoising objective and evaluate them across protein- ligand, drug-target, protein-protein, and antibody-antigen interactions. The model achieves state-of-the-art or competitive performance across benchmarks, scales effectively with model size, and successfully identifies wet-lab validated definity-enhancing antibody mutations, establishing a generalizable framework for biomolecular interactions. Our open-source implementation is available at com http://gtttttooooo/o//eeetttr/// wtt.o///eeettrr//////151.
https::/gi
INTRODUCTION
Biomolecular interactions, including protein-proter and protein-ligand interactions, are pivotal to a
wide range of biological processes, such cellular functions (Busch al.) 2025) and therapeutic
applications (Pacholarz et al., 2012) Cheng et al.) 2021). Recent breakthroughs, exemplified by
AlphaFold and related methods (Abramson al.] 2024 Krishna et al. 2024 Chai Discovery. 2024]
Chen et al.) [2025] Wohlwend et al.) 2024). have enabled accurate structure modeling of complex biological proteins. For example, AlphaFold demonstrates remarkable accuracy predicting protein complex structures from sequences, effectively resolving both the spatial arrangement and interaction mechanisms protein-protein and protein-ligand interactions. Despite recent research progress, structure prediction remains intermediate step, with the ultimate goal being the design of functional proteins with strong binding affinities for specific targets.
Binding affinity prediction remains difficult due to the scarcity of reliable experimental affinity labels. Early methods relied exclusively on sequences as input (Sun et al. 2017: Rao et al. (019b), failing to fully leverage recent advancements in structure prediction (Abramson al. 2024), even though
"quul contribution.
"ooreeponniig authors.

---

Published as a conference paper at ICLR 2026
all-atom structures are fundamental to understanding binding mechanisms and determining affinity. Recent approaches have incorporated structural information but are often specialized for specific types of biomolecular interactions. For instance, RDE-Network (Luo et al. 2023), DiffAffinity (Liu eet al. 2024), and Prompt-DDG (Wu et al. 2024) target protein-protein binding affinity, while MGraph-DTA (Yang et al., 2022). HGNN-DTA (He et al.. 2023), and ProFSA (Gao et al.) (023) focus on protein-ligand binding affinity. This narrow specialization limits the generalizability of these approaches and their applicability across diverse interaction types and data modalities.
Recent advancements in foundation models across various domains, such as BERT (Kenton & Toutanova 2019) and the GPT series (Floridi Chiriatti 2020, [Achiam al., 2023) for language. the SAM series (Kirillov et al. 2023 Ravi et (0024) for vision. and CLIP (Radford et al. (021) for multi-modality, have inspired new directions for addressing these challenges. These models are first pre-trained on broad datasets using self-supervised learning objectives and then adapted to a wide range of downstream tasks (Bommasani et al.) 202T). By reducing the need for domain- specific inductive biases and leveraging large datasets, these models learn powerful representations that significantly enhance generalization capabilities across diverse tasks. In biology, AlphaFoold exemplifies this paradigm by developing universal model capable of predicting various types of biomolecular complex structures, demonstrating the power of generalization across interaction types.
Inspired by these successes, adapt the AlphaFold 3 architecture for representation learning to design a structure-based unified foundation model for binding affinity prediction. While AlphaFold 3 is primarily generative model for predicting biomolecular structures, adopting AlphaFold 3 for affinity prediction requires non-trivial modifications. We propose a novel strategy to adapt AlphaFold 3 representation learner, guided by three key insights: (i) when the model's goal shifts from predicting geometry to encoding observed geometry, the heavy conditioning (trunk) module becomes less critical and might allow for more lightweight design; (ii) the wansformer-based atom and sequence-leve architecture offers a general framework for jointly encoding both structural and sequential information; (iii) by large-scale pre-training on structure data, it can potentially mitigate data scarcity in affinity prediction and improve generalization across diverse interaction tasks.
Building on this framework, we present the Atom-level Diffusion Transformer (ADiT), a universal, all-atom structure foundation model designed to learn transferable representations across diverse biomolecular interactions. Our model accepts sequence and structure as inputs, employs a unified scheme to tokenize proteins and molecules, and generates atom-level, token-level, and complex- level representations. We integrate key components of diffusion transformers while simplifying the architecture by removing dependencies a pre-training and fine-tuning framework, we train three versions of our ADiT, namely ADiT-S, ADiT-M, and ADIT-L,, on the PDB dataset. These pre-trained models evaluated across four different types of interactions including protein-ligand drug-target, protein-protein, and antibody- competitive performance across all benchmarks. Ablation
multiple sequence alignments and templates. Following
antigen, and demonstrate state-of-the-art studies highlight the importance of large-scale pre- pre-training and atom-level modeling, while consistent performance gains with increasing model size align with established scaling trends in other domains.
Main contributions: (i) We present a general- purpose foundation model, ADīT, designed for binding affinity prediction across various types of biomolecular interactions. (ii) Extensive experiments demonstrate that ADIT, empowered by large-scale pre-training, achieves state-of-the-art or competitive performance across diverse, well-established benchmarks. (iii) We observe consistent performance improvements across various benchmarks as model size increases, reflecting scaling trends observed in other domains. (iv) A case study on antibodies shows that ADiT can successfully identify eehanncinn mutations, with predictions validated by wet-lab experiment results.
2 RELATED WORK
Protein Representation Learning. Research in protein representation learning can be broadly categorized into two main paradigms: requence-based methods and structure-based methods. Sequence-based approaches interpret protein sequences as a form of biological language, leveraging large-scale pre-training
sequence or alignment datasets to extract functional and evolutionary information (Rao et al., 2019b a Elnaggar et al. 2021 Rives et al.) 2021 Meier et al.] 2021] Wang et al.) 2024b; Xu et al. 2023). In contrast, structure- saaettrr--bsed representation learning has gained significant attention due to the critical role of protein structures in determining function, as well as recent breakthroughs in structure prediction tools (Jumper et al. 2021). These methods explicitly
2

---

Published as a conference paper at ICLR 2026
Pre- Pre-training
Applications
Atorn Attn, Encoder DoL
Atom Attn. Decoder Jooa ofl)
All-atom siruclure
akip
Input Embedder
Men Seign
Arte dengin
TekeedWoom
dinch
K00 Sagg
Protein
Banding
Protein Protein Binding
Aggregate Predict
Cond
Token Attention 19000 DI
Protein property Atom-llaeel noise
ADIT Nath Trric
Protein Structure Denciing
Drug Tarpet
antbody-Artihh Brnding
Banding
Figure 1: Biomolecular complex representation learning with ADIT. The model inputs all- all-atom structures, processes them through multi-level attention trunk, and generates atom representations. It is pre-trained using a denoising objective and then fine-tuned for downstream application tasks.
encode protein structures using techniques such as 3D CNNs (Derevyanko et al. 2018
GNNs (Zhang They employ
ct al. 2023c Zhao et al. 2024 Fan et al. 2023), and transformers (Jiao et al. 2024
various supe-sup learning strategies, including contrastive learning, self- celf-prediction and diffusion
denoising, to train structure encoders effectively (Zhang et al.) 2023c Chen et al. 2022 Zhang et al.,
2023d). Additionally, some approaches use structure tokenizers convert 3D protein structures into token sequences, enabling the application of sequence-based methods for structural representation
learning (Su et al., 2024 Wang et al.] 2024a Heinzinger al., 2024). However, many current
approaches have two key limitations:
(I) they primarily focus on coarse-grained, residue-level structures, and (2) they are restricted to proteins alone. In this work, we address these gaps by
focusing on all-atom representation learning for general biomolecular complexes.
Biomolecular Complex Modeling. Recent advancements in molecular docking, both in predicting bound structures from unbound inputs (Ganea et al. 2022, Zhang et al.) [223al Corso et al.) 2023
Ketata et al., 2023) and in directly predicting bound structures from sequences (Abramson al.
2024, Krishna et al. 2024, Chai Discovery 2024 Chen et al.) 20255 Wohlwend et al.) 2024).
have substantially ad advanced the field of biomolecular complex structure prediction. Despite these progress, predicting structures remains an intermediate step towards designing proteins with functional properties like strong binding affinity. Numerous studies have focused on predicting these binding affinities for protein-protein and moleill interactions. For instance, prior work has explored energy-based based methods (Delgado et al. 2019 Alford et al.) (0077), evolution- approaches (Meier al., 2021 Rao al. 2021 Notin al. 2022), end-to-end learning models (Shan et al.) 2002)) and structure-informed pre-training methods (Yang et al. [2200] Luo et al. 2023: Wu et al.) [2224] Mo et al.
2024 Liu et al., 2024) for protein-protedd binding affimity. For protein-ligand binding, pre-training
methods have become the dominant approach due to data scarcity (Karimi et al. 2019 Liu et al.
2023. Wu al. 2022 Zhou al. 2023 Gao et al. 2023). Protein-ligand afmnity prediction is crucial
in drug-target prediction (Yang et al. 2022 He et al. 2023 Kroll et al. 2023a b). underscoring its importance in protein design. It is also common to benchmark against zero-shot baselines such as physics-based scoring functions (e.g.. Rosetta (Barlow et al. (018)), predicted mutational effects (e.g., NERE (Jin al. 20236) and DSMBind (Jin et al. 2023a)), and structural confidence metrics (c.g., AlphaFold 3' s's Tp (Abramson et al. 2024)), although these typically fall far short of specialized fine- tuned models. Besides, Boltz-2 has been recently proposed for grotein-ligand binding affinity prediction (Passaro et al. 2025). However, unified framework for biomolecular binding affinity prediction remains elusive, with many studies still relying on sgquence-basee methods and overlooking breakthroughs like AlphaFold 3. In this work, we introduce a general foundation model for biomolecular complexes, leveraging structures to learn robust representations.
METHOD
In this section, we introduce the Atom-level Diffuusio Transformer (ADIT), a unified foundation model for diverse biomolecular interactions. Our model employs unified scheme to tokenize and featurize proteins and molecules (Sec. 3.2) and incorporates diffusion transformers to capture token-
and alvel interactions (Sec. 3.3 Following standard practices in modern foundation model
development (Kenton& & Toutanova 2019 Achiam et al.) 2023), we train ADiT with pre-training
and fine-tuning framework (Sec. 3.4 An overview of our method is shown in Figure
3

---

Published as a conference paper at ICLR 2026
Reprosetatio
Sap
Cond
Stage Drgss
PSM
Exbbed
Frabed
Alom Doed
Tranderrue
Li Luaver Rec.
Liase BeLU
Alzm Per
Clesr
Liuner
Close
Liberr
Alae-Tobee
Takm Aomtot
SwitLu
SWISLD
Sagle
Uaee
SwiDLU
SwIOLO
Sqgidd
Trom Coea
Aeye
Diliaioo Traallena Liean
Cooate Alse Conditoos and Reprasotaais
Sagle
Cor
Auma
Mon Type
Alce Ooents
LK Leaue
Enabed
Liboe
Ccad
Alagte
Absa Cood
Abea Sirgle
Ataniie Docodar
Coad
Alin Darlle
SwicLU
Four
Diflutic Transooneer
Doed Adyphc Stale
Fimbed
I.N linwer
Area tain Bags
Logers (Atunn
Alee Stgll
Fagr
Ake: Aus./atiie
(a) Feature Initaaiiztionn
(b) Main Trunk
(c) Diffusion Tronsformer Laper
Figure The Atom-level Diffusion Transformer (ADiT) architecture. (a) Initialization of features for atom and token sequences. (b) Main trunk for hierarchical representation learning, comprising an atom attention encoder, global token attention, and an atom attention decoder. (c) Diffusion transformer blocks with multiple sparse attention update layers for atom-level modeling.
3.1 ALPHAFOLD FOR BINDING AFFINITY PREDICTION
Given the sequence A E (1,.., 20." and the structure E R Lx3 of biomolecular complex, our objective to extract atom representations and use them predict binding affinity through downstream linear or muli-layer perceptron (MLP) prediction head. As AlphaFold (Abramson et al.) 2024) provides a highly expressive, transformer-base architecture that jointly encodes sequential and structural information at the atom level across diverse types of biomolecular interactions, it is natural starting point for building general-purpose purpose representation learner.
However, adapting AlphaFold 3 for representation learning is both non-trivial and essential. AlphaFold was originally designed as generative model for structure prediction, and prior work has shown that generative models often produce suboptimal representations when applied directly to downstream tasks such as protein function or affinity prediction (Hu et al.) (002). This limitation arises because generative objectives primarily optimize for reconstructing structures rather than capturing features that are informative for functional or interaction properties. Similar trends have been observed in other domains: for instance, generative models for images text, while capable of producing high-quality outputs, do not automatically yield strong embeddings for tasks like classification, retrieval, or zero-shot generalization (Li al. 2023 Achiam al., (023)). Consequently, transforming AlphaFold into ADiT, model explicitly designed for representation learning, requires substantial modifications both its architecture and pre-training strategy.
Specifically, AlphaFold 3 is conditional diffusion model that consists of two primary components: a trunk conditioning module and a diffusion module heavily conditioned on diffusion time step. The conditioning module is multimodal encoder that integrates information from sequences, multiple sequence alignments (MSAs), and structural templates for structure prediction. However, in representation learning scenarios where complex structures provided as input, the model's ''s primary role shifts from inferring geometry to encoding the observed geometry. As a result, the reliance on heavy multimodal conditioning becomes less critical. Additionally, AlphaFold 3 depends heavily
4

---

Published as a conference paper at ICLR 2026
MSAs and templates, which are not always available and may be suboptimal for learning generalizable representations. In contrast, the diffusion module employs a modern transformer architecture that is well-suited for scalability. It incorporates advanced techniques such as SwiGLU (Shazeer 2020), gating mechanisms from AlphaFold 2 (Jumper et al., 2021), and a two-level architecture that alternates between atom-level and token-level representations. These features make the diffusion module particularly effectiie for updating representations based on both the geometry and the sequence. Based on these insights, we remove the dependencies on MSAs and templates, and introduce ADiT for representation learning across diverse biomolecular interactions and affinity prediction.
3.2 SIMPLIFIED FEATURIZATION SCHEME
We employ a generalized tokenization scheme, where each residue in proteins is represented as single token, and each heavy atom in small molecules is encoded as a single token. We initialize the multi-level input features using top-down approach. This involves first initializing features at the token level, followed by deriving atom-level features by integrating both atom-specific information and the propagated token-level features. Details of our feature embedder illustrated in Fig. a2))
Token Level Encoding. We first construct a token-level sequence embedder that encodes the chemical and biological information of complexes, producing representations for all tokens. These representations serve as the conditioning input for the subsequent diffusion transformer. Unlike AiphaFold 3, which relies on MSA and template inputs, introducing additional complexity and potential challenges when MSAs or templates are unavailable or of low quality, we use a pre-trained protein language model, ESM-2-650M (Lin et al. 2023), incorporate evolutionary information. Specifically, the token conditioning representations consist of two components: the ESM sequence feature and the token type embedding. The ESM sequence feature is only computed for residue tokens, while for small molecule tokens, this feature is set to zero, The token type embedding encodes the origin of the token, distinguishing between protein and small molecule tokens.
Next, we construct pair representations by concatenating the corresponding token conditioning representations and integrate relative position encodings that include both relative sequence and chain distance. Different from AlphaFold 3, we avoid computationally intensive Pairformer blocks, but only use a concatenation and linear layer for embedding to save compute. We use cooken and 2 ztoken to denote the token conditioning and pair representations, respectively.
Atom Level Encoding. At the atomic level, we encode both chemical and evolutionary information, along with structural information for each atom. Note that the subsequent diffusion transformer will condition on the former (chemical and evolutionary information) while extracting representations from the latter (structural information). To achieve this, distinguish between the single representation alom which includes structural information, and the conditions cl which exclude it.
To generate atom conditioning representations _ embed the atom type and name (for protein atoms) and augment them with token conditioning representations. For single atom representations gatom, encode the atom type, name, and coordinates. Then, for atom pair representations zatom combine atom conditioning representations, embed their Euclidean distance using RBF kernels, and encode whether they belong the same token. These atom pair representations are then augmented with the corresponding token pair representations to capture both local and global interactions.
3.3 HIERARCHICAL REPRESENTATION LEARNING
To capture both token- and atom- information, we adopt a hierarchical representation learning framework, utilizing Diffusion Transformer (DiT) modules (Peebles & Xie 2023). The framework begins by updating atom representations through Nattm DiT blocks, followed by "Atom2Token" average pooling operation to generate token representations. These token representations are then refined using block DiT blocks. Subsequently, a linear layer and a "Token2Atom" unpooling operation produce updated atom representations, which are combined with the previous atom representations via skip connection. The "Token2Atom" operation is non-learned broadcast operation. The atom representations are further refined through another Natoon DiT blocks. An overview is illustrated in Figure b) and outlined in Algorithm
Diffusson Transformer Block. We adapt the diffusion transformer to model both interatomic and inter-token interactions for representation learnin. The transformer comprises multiple stacked
5

---

Published as a conference paper at ICLR 2026
blocks, each containing an Adaptive LayerNorm function, multi-head self-attention module, and transition function, as detailed in Algorithm [ We add skip connections to enhance training stability. The Adaptive LayerNorm function and transition function process both single and conditioning representations, while the multi-head self-attention module computes attention weights using single and pair representations. Formally, the multi-head self-attention module is defined as:
q
k + Lineary (Ziz) + Big)
q. k 0. + Linear kt.
A + softmax,
g. - (Linearr
8, < Linear concatt,
g
where h denotes the attention head index, and is the dimension of the latent representation. The term Biy controls whether interactions between pair (i.j) are modeled. For token-level interactions, all Bij values are zero. For atom-level interactions, each group of 32 atoms attends to a nearby group of 128 atoms (based on sequence proximity). The DiT architecture is illustrated in Figure (cc) c).
Notably, ADIT utilizes only single linear layer to embed all atom positions during the initialization of atom representations, without incorporating geometric biases such as locality or SE(3) invariance. This non-equivariant design contrasts with current trends that emphasize stronger domain sppecff inductive biases, thereby facilitating the general usage on different data modalities. We approximate the required invariance through a combination of centroid-centering of the input coordinates and data augmentation via random rotations during pre- pre-training. We hypothesize that this non- equivariant approach offers greater flexibility to capture crucial non-geometric features (e.g.. electrostatic interactions and chemical semanties from RDKit and ESM-2 features) that govern binding thermodynamics, without being overly constrained by the strong geometric bias of equivariant layers. Furthermore, non-equivariant Transformers are architecturally simpler and highly scalable, enabling us to build the foundation model more efficiently than current equivariant frameworks.
3.4 PRE-TRAINING AND FINE-TUNING FRAMEWORK
Inspired by recent foundation models, we adopt a two-stage training approach: first pre-training our model with denoising objective, followed by fine-tuning it for specific downstream tasks.
Pre-training. Denoising pre- pre-training has emerged as a widely used self superviue learning approach for learning representations of 3D structures (Zaidi et al. 2023 Zhang et al. 2023b). The process involves adding Gaussian noise to each atom's 's coordinate as E ~ (O, o ), where a represents the noise scale. The noisy structure is then input into ADIT, and noise prediction head is applied to the resulting atom representations. As ADiT is a non non-cquivariant model, we incorporate random rotations of structures data augmentation. Following (Zaidi et al., 2023), scale, treating it as tunable hyperparameter. Preliminary experiments indicate that 0.5A yields the best performance. Although we also experimented with varying noise scales as in (Zhang et al.. 2023d), no practical improvements were observed. Nonetheless, we hypothesize that a carefully curated training distribution of noise scales could enable the model to better capture both coarse- and fine-grained features. Due to computational constraints, we leave this exploration to our future work.
fixed noise
Fine-tuning. To adapt our pre-trained model for downstream tasks, we fine-tune it using a reduced learning rate. For predicting properties of the bound complex, we input the structure into the model to obtain output atom representations. These representations are then sequentially aggregated through "Atem2Token" average pooling and "Token2Complex" sum pooling, followed by prediction head. In addition, our goal is to learn representations from clean samples, rather than to generate samples from noise in standard diffusion models. We do not need to condition on the diffusion time step. Accordingly, during fine-tuning, the time step is always set to zero, corresponding to a clean sample.
Difference with AlphaFold 3. While ADIT leverages architectural components similar to those in AlphaFold 3, our model represents a significant re-engineering tailored for the task of affinity prediction, leading to substantial differences in architecture, input, and objective. The key distinctions are as follows: (1) Goal: AlphaFold 3 is fundamentally a generative model designed for structure prediction, while ADIT is a representation learner specialized for affinity prediction. (2) Inputs: We remove AlphaFold 3's reliance on computationally expensive MSA and template conditioning, which we replace with pre- pre-trained protein features from ESM-2 and explicit RDKit features for small molecules. (3) Architecture: We simplify the backbone by removing the heavy trunk module and the

---

Published as a conference paper at ICLR 2026
computationally intensive Pairformer blocks, focusing instead on more scalable stack of Diffusion Transformer blocks. (4) Objective: We adapt the core diffusion transformer module to a simpler denoising objective, rather than the generative structure prediction task. We believe these adaptations are important for the strong performance of ADiT on binding affinity prediction.
EXPERIMENTS
The goal of this paper is to develop a general-purpess foundation model for predicting the binding affinity of various types of biomolecular complexes. In this section, evaluate our model four downstream tasks and compare its performance against recent strong baselines. The experimental results demonstrate that general-purpose foundation model, ADIT achieves state-of-the-art competitive results across these tasks, even compared with specialized models.
4.1 PRE-TRAINING SETUP
For pre-training, curate datasets from the Protein Data Bank (PDB) (Berman et al. 2000), resulting in 433,297 protein single chains, 481 1,322 protein-protein interaction examples, and 427,947 protein- ligand interaction examples. Here the interaction example refers to the interaction between two specific chains that form chain-chain interfaces within a single PDB entry. Two chains are considered interact if the minimum heavy-atom distance between them is less than angstroms. To accelerate pre-training, we leverage the clustering results provided by the RCSB PDB, yielding a total of 150,009 clusters. Importantly, we do not use any function labels during pre-training avoid potential data leakage and ensure fair evaluation on downstream tasks.
To explore the scaling effects, we pre-trained three model variants: ADIT-S (12M params), ADIT-M (35M params), and ADiT-L (253M params). The largest variant, ADÎT-L, matches the number of layers and hidden dimensions used in AlphaFold 3. The differences between these variants lie the configurations of their atom and token transformer layers. For details, please refer to Appendix D.I
4.2 PROTEIN-LIGND BINDING AFFINITY PREDICTION
Setup. For protein-ligand binding affinity, we adopt the setup from Atom3D (Townshend al. 2021), where each sample consists of the bound structure of protein-ligand complex, and our goal is to predict its binding affinity. The dataset includes two splits based on protein sequence similarity: one with max. sequence identity 30% (LBA-30) and another with max. sequence identity 60% (LBA-60). We use the Root Mean Square Error (RMSE), Pearson correlation coefficient, and Spearman's is rank correlation coefficient evaluation metries. We include recent strong baselines, and additionally fine-tune Protenix (Chen et al. 2025), an open-source reproduction of AlphaFold3.
Results. The results in Table 1 demonstrate that all three ADIT variants achieve competitive performance compared to previous strong baselines. Among them, ADiT-L achieves state-of-the-ar results across both data splits and all evaluation metrics. Specifically, on the LBA-30 benchmark, outperforms the previous best method, GET, with 1.9% increase in Pearson correlation and a 0.77% improvement in Spearman correlation. On LBA-60 benchmark, surpasses prior best baselines ProNet and ProFSA, achieving a 7.2% reduction in RMSE, along with 4.2% and 4.3% increases in Pearson and Spearman correlations, respectively. Notably, even the smallest variant, ADiT-S, with only 12M parameters, outperforms most existing baselines. Furthermore, we observe that prior methods exhibit limited generalization capabilities, performing well on either L.BA-30 or LBAA-00, but not both, e.g., GET and ProNet. In contrast, approach demonstrates consistent and robust performance across both dataset splits. Morcover, our results underscore the superiority of ADiT over Protenix across all evaluation metrics, showing that simply fine-tuning generative models for representation learning and downstream tasks is insufficient to unlock their full potential.
4.3 DRUG-TARGET BINDING AFFINITY PREDICTION
Setup. For the drug target binding affinity task, employ the widely used Davis dataset (Davis et al., 2011), consisting of 30,056 data points that capture binding affinities between 72 small-molecule drugs and 442 target protein kinases. Following (He et al. 2023), we perform random split of the
7

---

Published as a conference paper at ICLR 2026
Table Protein-ligand binding affinity prediction results. The best result is highlighted in bold, and the second-best result is underlined. Differences within 0.002 are considered negligible.
Sequence Identity 30%
Sequence Identity 60%
Model
RMSE
Pearson
Spear:
RMSE.
Pearson
Spear:
B B (Bepler & Berger
2019
1.985
0.165
0.152
1.891
0.249
0.275
Dee (Öztirk al.. 2018
1.866
0.472
0.471
1.762
0.666
0.663
DeepAffnity (Karmm et al 2019
1.893
0.415
0.426
MaSIF (Gainza et al. 2020
1.484
0.467
0.455
1.426
0.709
0.701
Atom3I-GGNN Townshend ct al.
2021
1.60)
0,545
0.533
1,408
0.743
0.743
IEConv Herrmsslla et al. 2021
1.554
0,414
0,428
1,473
0.667
0.675
Holoprot Sommath cl al. 2021
1,464
0.509
0.500
1.365
0.749
0.742
ProNet Wang al. 2023
1,463
0.551
0.551
1.343
0.765
0.761
Uni-Mol Zhou
2023
1.520
0.558
0.540
1.619
0.645
0,653
GeoSSL Liu al. 2023
1,451
0.577
0.572
ProFSA
Gao a 20224
1.377
0.628
0.620
1.377
0.764
0.762
GET (Kong al. 2024
1.327
0.620
0.611
1,510
0.675
0.688
GET-PS Kong at.
2024
1.309
0.633
0.642
1,509
0.676
0.680
Protenix CChnn at.
2025
1.354
0.569
0.593
1.781
0.707
0.737
ADiT-S
1.337
0.626
0,618
1,413
0.740
0.740
ADIT-M
1.353
0.622
0.630
1.335
0.764
0.752
ADIT-L
1.308
0.645
0.647
1.246
0.797
0.795
Table Results for drug-target and protein-protein binding affinity prediction. The best result in bold, and the second-best underlined. Differences within 0.002 are considered negligible.
Drug- Target
Protein-Protein Protein
Model
MSE.
Model
Pearson
Spear:
RMSE,
MAE
DeepDTA
0.261
0.630
ESM-IF
0.319
0.280
1.886
1.285
MT-DTI
0.245
0.665
End
0.637
0.488
1.619
1.176
GraphDTA
0.229
0.685
DDGPred
0.658
0.468
1.499
1.082
rZMLP
0.205
0.709
MIF-Network
0.652
0.513
1.593
1.146
EnsembleDLM
0.202
RDE-Network
0.644
0.558
1.579
1.112
DTs
0.208
0.743
Diff'Affiiity
0.669
0.556
1.535
1.093
MgraphDTA
0.207
0.710
Prompt-DDG
0.677
0.591
1.520
1.077
NHGNN-DTA
0.196
0.744
Surface- Surface-VQMAE
0.648
0.561
1.588
1.127
ADIT-S 0.252
0.690
ADiT-S
0.660
0.524
1.597
1.132
ADIT-M
0.216
0.734
ADiT-M
0.683
0.539
1.559
1.098
ADËT-L
0.198
0.751
ADIT-L
0.691
0.560
1.540
1.088
Davis dataset five times and report the average performance across these splits. We employ two evaluation metrics: Mean Squared Error (MSE) and the r metric. Details in Appendix D.3
Results. The results, summarized the left part of Table 2 demonstrate that ADIT delivers competitive performance, with its best-performing variant, ADīT-L, achieving state-of-the-ar results on both the MSE and 22 metrics. Moreover, ADiT exhibits clear scaling effect with respect to model size, suggesting that larger models have the potential to further improve performance.
4.4 PROTEIN-PROTED BINDING AFFINITY PREDICTION
Setup. For the protein-protein binding affinity task, we adopt the SKEMPIv2 dataset (Jankauskaité al. (209)), the most widely used benchmark for AAg prediction. SKEMPIv2 comprehensive, annotated mutation dataset covering 348 protein complexes, containing 7,085 amino acid mutations along with corresponding changes in thermodynamic parameters and kinetic rate constants. Since SKEMPIv2 does not provide structures for the mutated complexes, FoldX Delgado et al. 2019). We evaluate the performance of ADiT using a split-by-complee cross-validation validation on the SKEMPIy2 dataset following (Liu al. 2024). The metrics include: threefold cross
generate them using
(1) Peanson correlation coefficient, (2) Spearman's rank correlation coefficient, (3) root mean squared
error (RMSE) and (4) mean absolute emror (MAE). Details in Appendix D.4
8

---

Published as a conference paper at ICLR 2026
Table 4: Ablation study results on SKEMPIv2 based on ADiT-M.
NO.
Model
Pearsonr
Spearmanrt
RMSE
MAE
0
ADIT-M
0.683
0.539
1.559
1.098
w/o pre-training
0.649
0,511
1.624
1.169
2
w/o all-atom info.
0.658
0.517
1.606
1.153
w/ larger size
0.691
0.560
1.540
1.088
w/ smaller size
0.660
0.524
1.597
1.132
Results. The detailed results are presented Table Among all baselines, ADIT demonstrates
strong and consistent performance. The best-performing variant, ADiT-L, achieves the highest Pearson correlation coefficient and the second-best Spearman's 's rank correlation coefficient. Although this is ranking task, where metrics like RMSE and MAE are less critical, ADiT still performs comparably to the strongest baselines on these metrics, with only marginal differences.
4.5 ANTIBODY-ANTIGEN BINDING AFFINITY PREDICTION
Setup. With the models trained on SKEMPIv2, we test their performance on the HER2 binders test set, following (Cai et al. 2024). The HER2 binders test set is collected from (Shanehsazzadehh et al., 2023) and contains high-qui high-quality binding affinity data, measured by surface plasmon resonance (SPR) 419 HER2 binders with de novo designed CDR loops. The antibodies in the dataset variants of Trastuzumab that have high edit distance (7.6 on average), making them potentially challenging for AAG bind predictors trained on low-edit- low-edit-distance data. Here, we use Pearson and Spearman correlation as evaluation metrics to assess the performance of our ADiT. Details in Appendix D.5
Results. Results are shown in Table Our best-performing variant, ADÎT-L, demonstrates superior performance. outperforming previous methods Bind-ddG (Shan et al.] (022) and GearBind (Cai et al., 2024). These findings underscore the robust generalization capability of ADiT to held-out antibody data.
Table 3: Antibody-antigen binding affinity prediction results. The best result is highlighted in bold, and the second-best result is underlined.
Model
Pearson
Spear.
Bind-ddG
0.387
0.388
GearBind
0.478
0.475
Case Study. We evaluate ADIT's performance on real-world in silico affinity maturation using models trained SKEMPIy2. Two starting antibodies with different formats and target antigens tested (see Appendix for details). Figure [3] shows the rankings of seven wet-lab- validated affinity-enhancing mutations from (Cai ct al. 2024 The results demonstrate that ADiT assigns better ranks to most beneficial mutations (e.g.. S54Y and S57W for Anti-5T4 UdAb, and SHI03W, SH103Y, and IL34W for CR3022), highlighting its potential as an effective tool for antibody optimization.
GearBind +
0.515
0.517
FoldX
0.246
0.329
ADIT-S
0.527
0.558
ADiT-M
0.507
0.523
ADiT-L
0.567
0.576
4.6 ABLATION STUDY
To better understand the key contributors to the performance of ADiT, we conduct an ablation study on the protein-protein binding affinity prediction task using the SKEMPIv2 dataset. We use ADiT-M as the baseline and evaluate the impact of several single-factor modifications: (1) removing pretraining and use randomly initialized parameters, (2) replacing all-atom (AA) structures with backbone-only structures in both the pretraining and finetuning stages, (3) scaling up the model size (35M) ADiT-L (253M), and (4) scaling down ADiT-S (12M).
The results are summarized in Table 4 (1) Comparing No.0 and No.1 highlights the importance of pretraining, which significantly contributes performance gains (5.2% improvement in Pearson, 5.5% improvement in Spearman, 4% improvement in RMSE, 6% improvement in MAE). (2)
Onll limited number of mutations were experimentally validated, suggesting that additional affinity- enhancing candidates likely exist.

---

Published as a conference paper at ICLR 2026
AciT models' ranks of affinity-arhancing mutations
AdiT Pre-training Loss Curves
Amt--544 UsAb
CR3022
200
ADT-S
mutation
mutstion
4.5
ASTT-M
175
354Y
S57W
600
SH100Y
4.0
ADIT
150
(.3WW)
SH100D
Pre-training Loss
2.5
SLB3R
125
D21
100
2.5
500
2.0
75
1.5
50
25
20k
40k
B0k
80k
100k
120k
140K 180K
Training Steps
ADIS
ADTTM
ADTTL
ADTS ADTM ADTL
Figure 3: Case Study on Antibody Affinity Figure 4: Pre-training and Scaling Effects Maturation. ADIT models' ranks of affinity- Across Model Sizes. Pre-training loss curves for the small, medium, and large variants of ADiT. Larger models consistently achieve lower loss
enhancing mutations nanobody Anti-5T4
UdAb (780 mutations) and antibody CR3022 (1420 mutations) are shown. Lower rank better. throughout the pre-training process.
The comparison between No.0 and No.2 demonstrates the value of all-atom structural modeling (3.8% improvement in Pearson, 4.3% improvement in Spearman, 3% improvement in RMSE, 4.8% improvement in MAE), which provides richer structural information and leads to improved representations. (3) Finally, the comparison among No.0, No.3, and No.4 reveals consistent performance improvements with increased model size, reflecting a clear scaling trend in line with observations from other domains. Figure 4 presents the pretraining loss curves for the small, medium, and large model variants. Larger models consistently achieve lower loss during pretraining, underscoring their enhanced ability to capture complex molecular patterns.
5 CONCLUSION
In this work, we take an important step toward developing all-atom structure foundation models for biomolecular interactions. With large-scale pre-training, our unified foundation model achieves state-of-the-an or competitive performance across diverse tasks. Furthermore, we observe scaling effects our model architecture that align with trends seen other domains. Quantitative and qualitative analyses show the practical effectiveness of our approach. We further discuss the LLM usage in Appendix A and limitations in Appendix
Broader Impacts Our work has significant broader impacts for computational biology and drug discovery. Positively, ADiT provides powerful, general-purpose framework for modeling biomolecular interactions, potentially accelerating therapeutic discovery and advancing biological understanding. However, ethical concerns, such as the potential misuse in drug development, also warrant careful consideration. Responsible deployment and transparency are essential to maximize benefits while mitigating risks.
REPRODUCIBILITY STATEMENT
We provide a clear and detailed description of the architecture in Sec. along with the necessary pretraining and finetuning configurations in Sec. D The pretraining dataset constructed from publicly accessible PDB and clustered using publicly available clustering results. All finetuning datasets used in our experiments are also publicly accessible. Our open-source implementation is available at wtttttp::// htttps:///gttub....... httpub/.cco//eettto/////11/
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
The authors acknowledge funding from the Canada CIFAR AI Chair Program and the Intel-Mila partnership program. The computational resources for this project were provided by Mila and the Digital Research Alliance of Canada.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger. Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arX3vv23033007..7 2023.
Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew O'Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, Douglas Renfrew, Vikram Mulligan, Kalli Kappel, et al. The rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation, ((6)):031-3045 2017.
Kyle A Barlow, Shane O Conchuir, Samuel Thompson, Pooja Suresh, James Lucas, Markus Heinonen, and Tanja Kortemme. Flex ddg: Rosetta ensemble-base estimation of changes in protein-proted binding affinity upon mutation. The Journal of Physical Chemistry B, 122(21): 5389-5399, 2018.
Tristan Bepler and Bonnie Berger. Learning protein sequence embeddings using information from structure. In International Conference Learning Representations, 2019. URL https://
openreview ppenrevieewnet//ooeeeeeeeeeeeeee htt.:ee//oormmmll/SsyeeeeeeC
Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic acids research, 28(1):
235-242, 2000.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07255: 2021.
Hannah Busch, Muhammad Yasir Ateeque, Florian Taube, Thomas Wiegand, Björn Corzilius, and Georg Künze. Probing biomolecular interactions with paramagnetic nuclear magnetic resonance spectroscopy. ChemBioChem, pp. e202400903, 2025.
Huiyu Cai, Zuobai Zhang, Mingkai Wang, Bozitao Zhong, Quanxiao Li, Yuxuan Zhong, Yanling Wu, Tianlii Ying, and Jian Tang. Pretrainable geometric graph neural network for antibody affinity maturation. Nature communications, 15(1):7785, 2024.
Chai Discovery. Chai-1: 1: Decoding the molecular interactions of life. bioRxiv, 2024. doi: 10. 1101/2024. 10. 10.61 5955. URL.ttts. //www//wwwww////// /tr.l 10/11/2024. 10. 615955
Can (Sam) Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self-supervised learning. Bioinformatics, 39, 2022.
Xinshi Chen, Yuxuan Zhang, Chan Lu, Wenzhi Ma, Jiaqi Guan, Chengyue Gong, Jincai Yang. Hanyu Zhang, Ke Zhang, Shenghao Wu, et al. Protenix-adve advancing structure prediction through comprehensive alphafold3 reproduction. bioRxiv, pp. 2025-01, 2025.
Feixiong Cheng, Junfei Zhao, Yang Wang, Weiqiang Lu, Zehui Liu, Yadi Zhou, William R Martin, Ruisheng Wang. Jin Huang. Tong Hao, et al. Comprchensive characterization of -pttiin interactions perturbed by disease mutations. Nature genetics, 53(3):42-353, 2021.
11

---

Published as a conference paper at ICLR 2026
Gabriele Corso, Hannes Stärk, Bowen Jing. Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusson steps, twists, and turns for molecular docking. In The Eleventh International Conference Learning Representations, 2023. URL
K-mBbS
Mindy Davis, Jeremy Hunt, Sanna Herrgard, Pietro Ciceri, Lisa M Wodicka, Gabriel Pallares, Michael Hocker, Daniel K Treiber, and Patrick P Zarrinkar. Comprehensive analysis of kinase inhibitor selectivity. Nature biotechnology. 29(11):1046 (011)):04--101122222 2011.
Javier Delgado, Leandro G Radusky, Damiano Cianferoni, and Luis Serrano. Foldx 5.0: working with ma, small molecules and a new graphical interface. Bioinformatics, 35(20)4468-4166, 2019.
Georgy Derevyanko, Sergei Grudinin, Yoshua Bengio, and Guillaume Lamoureux. Deep convolutional networks for quality assessment of protein folds. Bioinformatics, 34(23):4046- 4053, 2018.
Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Wang Yu, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. Prottrans: Towards cracking the language of lifes code through self- supervised deep learning and high performance computing. IEEE Transactions on Puttern Analysis and Machine Intelligence, pp. 1-1, 2021. doi: 10. 109/TPAMI2001 133093333..
Hehe Fan, Zhangyang Wang. Yi Yang. and Mohan Kankanhalli. Continuous discrree convolution for geometry-sequence modeling proteins. The Eleventh International Conference on Learning Representations, 2023.
Luciano Floridi and Massimo Chiriatti. Gpt-3: Its nature, scope, limits, and consequences. Minds and Machines, 30:681-694 2020.
Pablo Gainza, Freyr Sverrisson, Frederico Monti, Emanuele Rodola, D Boscaini, MM Bronstein, and BE Correia. Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods, 17(2):184 192, 2020.
Octavian-Eugen Ganea, Xinyuan Huang, Charlotte Bunne, Yatao Bian, Regina Barzilay, Tommi S. Jaakkola, and Andreas Krause. Independent -equivrqui models for end-to-end rigid protein docking. In International Conference on Learning Representations, 2022. URL https:// openreview net foreet//oormmmmm (od
Bowen Gao, Yinjun Jia, YuanLe Mo, Yuyan Ni, Wei-Ying Ma, Zhi-Ming Ma, and Yanyan Lan. Self-supervised pocket pretraining via protein fragment-surroundiing alignment, In The Twelfth International Conference on Learning Representations, 2023.
Vladimir Gligorijević, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, et al. Structure-based protein function prediction using graph convolutional networks. Nature communications, 12(1):3168, 2021.
Haohuai He, Guanxing Chen, and Calvin Yu-Chian Chen. Nhgnn-dta: node-adaptive hybrid graph neural network for interpretable drug-target binding affinity prediction. Bioinformatics, 39(6): btad355, 2023.
Michael Heinzinger, Konstantin Weissenow, Joaquin Gomez Sanchez, Adrian Henkel, Milot Mirdita, Martin Steinegger, and Burkhard Rost. Bilingual language model for protein sequence and structure. NAR Genomics and Bioinformatics, 655) 11 2024. ISSN 2631-9268. doi: 10. 10.10993/argaa//lqeeI 150. URL /et.lppppe///oo..ror////////////
Pedro Hermosilla. Marco Schäfer, Matëj Lang. Gloria Fackelmann, Pere Pau Vázquez, Barbora Kozliková, Michael Krone, Tobias Ritschel, and Timo Ropinski. Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. Representations, 2021.
International Conference on Learning
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives. Learning inverse folding from millions of predicted structures. In International conference machine learning, pp. 8946-8970. PMI.R, 2022.
12

---

Published as a conference paper at ICLR 2026
Mingyang Hu, Fajie Yuan, Kevin Yang, Fusong Ju, Jin Su, Hui Wang, Fei Yang. and Qiuyang Ding. Exploring evolution-awee aware &-free protein language models as protein function predictors. Advances in Neural Information Processing Systems, 15:38883-38888 2022.
Justina Jankauskaite, Brian Jiménez-García Justas Dapkūnas, Juan Fernández- Recioo and lain H Moal. Skempi 2.0: an updated benchmark of changes in protein- protlirrotttnnoe binding energy, kinetics and thermodynamies upon mutation. Bioinformatics, 35(3):462 469, 2019.
Rui Jiao, Xiangzhe Kong, Ziyang Yu, Wenbing Huang, and Yang Liu. Equivariant pretrained transformer for unified geometric learning on multi-domain 3d molecules. arXiv preprint arXiv:2402.12714. 2024.
Wengong Jin, Xun Chen, Amrita Vetticaden, Siramush Sarzikova, Raktima Raychowdhury, Caroline Uhler, and Nir Hacohen. Dsmbind: Se (3) denoising score matching for unsupervised binding energy prediction and nanobody design. bioRxiv, pp. 2023 2023a.
Wengong Jin, Siranush Sarkizova, Xun Chen, Nir Hacohen, and Caroline Uhler. Unsupervised protein-ligand binding energy prediction via neural euler's 's rotation equation. Advances Neural Information Processing Systems, 36:33514-33528 2023b.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature. (97(59977777733-58888 2021.
Po-Yu Kao, Shu-Min Kao, Nan-Lan Huang, and Yen-Chu Lin. Toward drug-target interaction prediction via ensemble modeling and transfer learning. In 2021 IEEE international conference on bioinformatics and biomedicine (BIBM), pp. 2384-2391. IEEE, 2021.
Mostafa Karimi, Di Wu, Zhangyang Wang, and Yang Shen. Deepaffinity: interpretable deep learning of compound-protein affinity through unified recurrent and convolutional neural networks. Bioinformatics, (11211))3322-333333 2019.
Jacob Devlin Ming-Way Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of naacL-HLT, volume 1. Minneapolis, Minnesota, 2019.
Mohamed Amine Ketata, Cedrik Laue, Ruslan Mammadov, Hannes Stark, Menghua Wu, Gabriele Corso, Céline Marquet, Regina Barzilay, and Tommi S. Jaakkola. Diffdock-PP. Rigid protein- protein docking with diffusion models. In ICLR 2023 Machine Learning for Drug Discovery workshop, 2023. URL
net forum Satrrummid=AMMMMMMXUIII
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 4015-4024 2023.
Xiangzhe Kong, Wenbing Huang, and Yang Liu. Generalist equivariant transformer towards 3d molecular interaction learning. In JCML. OpenReview.net, 2024.
Rohith Krishna, Jue Wang, Woody Ahern, Pascal Sturmfels, Preetham Venkatesh, Indrek Kalvet, Gyu Rie Lee, Felix S Morey-Burrows Ivan Anishchenko. Ian R Humphreys, et al. Generalized biomolecular modeling and design with rosettafold all-atom. Science, 384((6993)ead11252,, 2024.
Alexander Kroll, Sahasra Ranjan, Martin KM Engqvist, and Martin J Lercher. A general model to predict small molecule substrates of enzymes based machine and deep learning. Nature communications, 14(1 14(1):2777, 2023a.
Alexander Kroll, Sahasra Ranjan, and Martin J. Lercher. Drug-targer interaction prediction using a multi-modal transformer network demonstrates high generalizability to unseen proteins. bioRxiv, 2023b. doi: 10. 10.11000/223.08..1111 5551177. URL htt.htt::: http://wwwwboolwloo/o//////////t ear y /023 0100002/11100 000002//11112202222
08 21. 554147
Tianhong Li, Huiwen Chang, Shlok Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan. Mage: Masked generative encoder to unify representation learning and image synthesis. In Proceedings of the IEFF/CVF Conference Computer Vision and Pattern Recognition, pp. 2142-2152, 2023.
13

---

Published as a conference paper at ICLR 2026
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scate prediction of atomic-level protein structure with a language model. Science, 379(6637): 11t0::11111111111 2023.
Shengchao Liu, Weitao Du, Zhi-Ming Ma, Hongyu Guo, and Jian Tang. A group symmetric stochastic differential equation model for molecule multi-modal pretraining. In International Conference on Machine Learning, pp. 21497-21526. PMLR, 2023.
Shiwei Liu, Tian Zhu, Milong Ren, Chungong Yu, Dongbo Bu, and Haicang Zhang. Predicting mutational effects on protein-protein binding via a side-chain diffusion probabilistic model. Advances in Neural Information Processing Systems, 36, 2024.
Shitong Luo, Yufeng Su, Zuofan Wu, Chenpeng Su, Jian Peng, and Jianzhu Ma. Rotamer density estimator is an unsupervised leamer of the effect of mutations on protein-protein interaction. In The Eleventh International Conference on Learning Representations, 2023. URL https //openreview.o net mettfooummiieeeee.ooooo.m.mmmmmm X9Y11K2mD
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Yuanle Mo, Xin Hong, Bowen Gao, Yinjun Jia, and Yanyan Lan. Multi-level interaction modeling for protein mutational effect prediction. arXiv preprint arXiv:2405.17800 2024.
Thin Nguyen, Hang Le, Thomas P Quinn, Tri Nguyen, Thuc Duy Le, and Svetha Venkatesh. Graphdta: predicting drug -aaget binding affinity with graph neural networks. Bioinformatics, 37 (8):1140-1147, 2021.
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan N Gomez, Debora Marks, and Yarin Gal. Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval. In International Conference on Machine Learning, pp. 16990-17017. PMLR, 2022.
Hakime Öztürk, Arzucan Özgür, and Elif Ozkirimli. Deepdta: deep drug target binding affinity prediction. Bioinformatics, 34(17):1821-i22, 2018.
Kamila J Pacholarz, Rachel A Garlish, Richard J Taalor, and Perdita E Barran. Mass spectrometry based tools to investigate protein-ligand interactions for drug discovery. Chemical Society Reviews, 41(11):4335-433 4355, 2012.
Saro Passaro, Gabriele Corso, Jeremy Wohlwend, Mateo Reveiz, Stephan Thaler, Vignesh Ram Somnath, Noah Getz, Tally Portnoi, Julien Roy, Hannes Stark, et al. Boltz-2: Towards accurate and efficient binding affinity prediction. BioRxiv, 2025.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
Zongzhao Qiu, Qihong Jiao, Yuxiao Wang, Cheng Chen, Daming Zhu, and Xuefeng Cui. rzmlp- dta: gmlp network with rezero for sequence-based drug-target affinity prediction. In 2021 IEEE international conference on bioinformatics and biomedicine (BIBM), pp. 308-313. IEEE, 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748 8763. PMLR, 2021.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer lcarning with tape. Advances in neural information processing systems, 32, 2019a.
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with tape. In Advances in Neural Information Processing Systems, 2019b.
14

---

Published as a conference paper at ICLR 2026
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8844-8856. PMLR, 18-24 Jul 2021.
Nikhila Ravi, Valentin Gabeur, Yuan-Time Ting Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, et al. Sam 2: Segment anything images and videos. arXiv preprint arXiv:2408.007711 2024.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, Lawrence Zitmick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning 250 million protein sequences. Proceedings of the National Academy of Sciences, 1 118(15), 2021.
Sisi Shan, Shitong Luo, Ziqing Yang. Junxian Hong, Yufeng Su, Fan Ding, Lili Fu, Chenyu Li, Peng Chen, Jianzhu Ma, al. Deep learning guided optimization of human antibody against sars-cov-2 variants with broad neutralization. Proceedings of the National Academy of Sciences, 119(11): e2122954119, 2022.
Amir Shanehsazzadeh, Sharrol Bachas, Matt McPartlon, George Kasun, John Sutton, Andrea Steiger, Richard Shuai, Christa Kohnert, Goran Rakocevic, Jahir M Gutierrez, et al. Unlocking de novo antibody design with generative artificial intelligence. bioRxiv, pp. 2023-01, 2023.
Noam M. Shazeer. Glu variants improve transformer. arxiv, abs/2022.05002, 2020. URL https:
7/api www.seetiiccchhll....... ottpp//cc.....ooooooooooooaaaaao
Bonggun Shin, Sungsoo Park, Keunsoo Kang, and Joyce Ho. Self-attention based molecule
representation for predicting drug-target interaction. In Machine learning for healthcare conference, pp. 230-248. PMLR, 2019.
Vignesh Ram Somnath, Charlotte Bunne, and Andreas Krause. Multi-scale representation learning proteins. Advances in Neural Information Processing Systems, 34, 2021.
Jin Su, Chenchen Han, Yuyang Zhou, Junjie Shan, Xibin Zhou, and Fajie Yuan. Saprot: Protein language modeling with structure-aware vocabulary. In International Conference on Learning Representations, 2024.
Tanlin Sun, Zhou, Luhua Lai, and Jianfeng Pei. Sequence-based prediction of protein protein interaction using deep-learning algorithm. BMC bioinformatics, 18:1-8, 2017.
Raphael John Lamarre Townshend, Martin Vögele, Patricia Adriana Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis. Sidhika Balachandar, Bowen Jing, Brandon M. Anderson, Stephan Eismann, Risi Kondor, Russ Altman, and Ron O. Dror. ATOM3d: Tasks molecules in three dimensions. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https: htttp::///penreview... wwwpeeewwwwwwwweeeeeeeeeeeeeeeee fformm ? FkDZLpK1M12
I Limei Wang, Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning hierarchical protein representations via complete 3d graph networks. The Eleventh International Conference Learning Representations, 2023. URL https:/ /ttpss:///ppreevee.. net http://wwww.eeeeoutttteaam ffrrmorrrmmi
9X- "X-hgLDLYkQ
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Diffusion language models are versatile protein learners. In International Conference on Machine Learning, 2024a.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Dplm-2: A multimodal diffusion protein language model. arXiv preprint arXiv:2410. 13782, 2024b.
Jeremy Wohlwend, Gabriele Corso, Saro Passaro, Mateo Reveiz, Ken Leidal, Wojtek Swiderski, Tally Portnoi, Itamar Chinn, Jacob Silterra, Tommi Jaakkola, and Regina Barzilay. Boltz-1: Democratizing biomolecular interaction modeling. bioRxiv, 2024. doi: 10. 10.1101/202..1.19 624167.
15

---

Published as a conference paper at ICLR 2026
Fang Wu and Stan Z. Li. Surface- Surface-VQMAE: Vector-quantized masked auto-encoders molecular
surfaces. In Forty-first International Conference on Machine Learning. 2024. URL https
//openreview net/ htt.://ww.....trrorrrooessss/saa
Fang Wu, Siyuan Li, Lirong Wu, Dragomir Radev, Qiang Zhang, and Stan Z Discovering the representation bottleneck of graph neural networks from multi-order interactions. 2022.
Lirong Wu, Yijun Tian, Haitao Lin, Yufei Huang, Siyuan Li, Nitesh V Chawla, and Stan Z. Li.
Learning to predict mutational effects of protein-protein interactions by microenvironmenttavaaa hierarchical prompt learning. Forty-first International Conference on Machine Learning, 2024.
URL httpps::/pppeeeviiiiit.e net Eorum? 1or.mmmdddgg1
Minghao Xu, Xinyu Yuan, Santiago Miret, and Jian Tang. Protst: Multi-modality learning of protein sequences and biomedical texts. In International Conference Machine Learning, pp- 38749-38767. PMLR, 2023.
Fang Yang. Kunjie Fan, Dandan Song, and Huakang Lin. Graph-based prediction of pootein-proteinn interactions with attributed signed graph embedding. BMC bioinformatics, 21:1-16, 2020.
Ziduo Yang, Weihe Zhong, Lu Zhao, and Calvin Yu-Chian Chen. Mgraphdta: deep multiscale graph neural network for explainable drug-target binding affinity prediction. Chemical science, 13(3): 816-833 116-833,20 2022.
Weining Yuan, Guanxing Chen, and Calvin Yu-Chian Chen. Fusiondta: attention-based feature polymerizer and knowledge distillation for drug-target binding affinity prediction. Briefings in Bioinformatics, 23(1):bbab506, 2022.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Genzalzz Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. The Eleventh International Conference on Learning Representations, 2023.
Yangtian Zhang, Huiyu Cai, Chence Shi, and Jian Tang. E3bind: An end-to-end equivariant network for protein-ligand docking. In The Eleventh International Conference on Learning Representations, 2023a.
Zuobai Zhang, Chuanrui Wang, Minghao Xu, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. systematic study of joint representation learning on protein sequences and structures. arXiv preprint 2.XXx:233300677. 2023b.
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. In The Eleventh International Conference on Learning Representations, 2023c.
Zuobai Zhang, Minghao Xu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, and Jian Tang. Pre- Pre-iraining protein encoder via siamese sequence-structuue diffusion trajectory prediction. In Thirth-seventh Conference on Neural Information Processing Systems, 2023d.
Jiale Zhao, Wanru Zhuang, Jia Song, Yaqi Li, and Shuqi Lu. Pre-training protein bi-level representation through span mask strategy on 3d protein chains. In International Conference on Machine Learning, pp. 21791-21 PMLR, 2024.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. URL https //openreview.o net 2eee:/ooommiddd2222
16

---

Published as a conference paper at ICLR 2026
A LLM USAGE STATEMENT
LLMs were used solely for language editing and refinement of the manuscript text. The LLM did contribute research ideation, experimental design, data analysis, content generation beyond linguistic improvements. All scientific concepts, results, and interpretations remain entirely the work of the authors.
The authors reviewed all LLM -Lu-sgggetted changes to ensure accuracy and alignment with the intended meaning. and no text was accepted without human verification.
B LIMITATIONS.
First, the current model is limited protein-protein and protein-ligane interactions, excluding other biomolecules such as DNA and RNA. Given AlphaFold 3's demonstrated ability to model a broader range of biomolecular interaction types, extending our approach to encompass all biomolecular interactions is promising for future research. Second, the computational resource constraints have limited our ability to further scale up the model size. We believe that larger models, combined with advanced pre-training objectives, could lead to stronger foundation models. Additionally, while our study provides initial insights into scaling laws for structure foundation models, further exploration with various model sizes is necessary to better understand these trends. Finally, our framework requires fine-tuning for downstream tasks, which poses challenges scenarios with extremely limited data. Future work should explore zero-shot inference capabilities to broaden the applicability of structure foundation models.
C ALGORITHMS
The pseudo codes for "Hierarchical Representation Learning" and "Diffusion Transformer" are shown in [ and Alg
Algorithm 1 Hierarchical Representation Learning
Input: token condition cluken, pair representation tuken and atom condition da single representation galem and pair representation z
2: Output: updated atom single representation atom
3: # Step 1: Update atom representations
4: gatom <- TreiussionTannsoommerrr g atom eatom atom
5:
6: Step Update token representations
7: glaken &- Atom2Token g gllm
8: g gloenn + Linear( cloken
loken
9: token <- DiffusionTTaassommen stokeen cloken gooeen)
10:
# Step Update atom representations again
12: galom <- Token2Atone (
gokkon) + galom
allm +- DiffusionTrann Transformer( aalom cooo alom
14: Return g atom
Algorithm 2 Diffusion Transformer
1: Input: single representation s, pair representation condition number of DiT blocks Mblock
2: Output: updated representation s
3: for 1 to Notoca do
4: $ <- AdaLN( S, c)
5: 8 AuttiheadAAttn/
6: 8 + Transition 88 c) +
7: end for
8: Return
17

---

Published as a conference paper at ICLR 2026
D EXPERIMENTAL DETAILS
Table 5: Model Hyperparameters.
Model
ADiT-S
ADIT-M
ADiT-L.
Model Configuration
Atom Repr Dim
128
128
128
Atom Pair Dim
16
16
16
Token Repr Dim
384
384
768
Token Pair Dim
32
64
128
Atom Enc Layers
3
3
3
Atom Enc Heads
4
4
4
Token Attn Layers
3
12
24
Token Attn Heads
8
8
Atom Dec Layers
3
3
3
Atom Dec Heads
4
4
4
Pre- Pre-Trainig Setup
Optimizer
Adam
Adam le-4
Adam
Learning Rate
le-4
le-4
Batch Size per GPU
Grad. Acc. Steps
1
2
2
Number of GPUs
4
4
32
Effective Batch Size
24
24
64
Training Steps
180k
180k
180k
Dataset Setup
Truncation Length
250
250
200
Random Rotation
yes
yes
yes
Table Dataset statistics.
Dataset
Split
Train
Validation
Test
Data Type
PDB
Random
142,509 clusters
7,500 clusters
Unlabeled Structure
Identity 30% Identity 60%
3.507
466
LBA
Protein-Ligand Affinity
3,678
460
460
Davis
Random
24,044
3,005
3,005
Drug- Target Affinity
Fold Fold
4.765
1,929
SKEMPIv2
4,282
2,412
Protein-Protein Affinity
Fold 2
4,341
2,353
HER2
419
Anibody-Antigen Affinity
All experiments in this paper performed NVIDIA Tesla A100 GPUs with 40G memory. The dataset statistics are presented in Table
D.1 PRE- TREINININ CONFIGURATION
We implement and pre-train three variants of ADIT: ADIT-S (12M parameters). ADIT-M (35M parameters), and ADiT-L (253M parameters), using the PDB dataset. Detailed training hyperparameters are listed in Table Notably, ADiT-L adopts the same number of atom- and token- level transformer layers as AlphaFold 3. To accommodate GPU memory constraints, we apply random employ data augmentation
sequential spatial cropping to the protein inputs. Additionally, strategy that centers the protein at the origin (based on its center of mass) and applies random rotations. These augmentations help the model learn more robust and generalizable representations.
18

---

Published as a conference paper at ICLR 2026
D.2 PROTEIN-LIGAND BINDING AFFINITY PREDICTION
Baselines We employ diverse set of baselines. including B&B (Bepler & Berger
2019).
De (Özüürk et 2018). DeepAffinity (Karimi et al. 2 MaSIF Gainza al. 2020).
Atom3D-Gii GNN (Townshend et al.
2021), IRConv (Hermosilla et al. 2021). Holoprot (Somnath
al. 202T), ProtNet (Wang et al. 2023). Uni-Mol (Zhou et al. 2023), GeoSSI. (Liu al. (023).
ProFSA (Gao et al., 2023), and GET (Kong al. 2024). Most of the reported results for the baseline
methods are taken from (Goo et al. 2023), while we independently evaluate GET (Kong et al. 2024)
on the LBA-60 dataset using the same configuration in the LBA-30 setting reported by the original
authors.
To fine-tune Protenix, we modified its Diffusion Module by replacing the noisy coordinates with ground-truth coordinates. We then extracted the output from the last Transformer layer of the AtomAAtientionDDecadee This output was passed through LayerNorm, followed by mean pooling to obtain per-token representations, and then sum pooling to derive a complex-level representation, following the same aggregation strategy used in ADiT. As for the hyperparameters, we kept the architectural settings of Protenix unchanged, since its architecture is fixed. However, we carefully tuned the training hyperparameters avoid underestimating the baseline performance. Specifically, we used a learning rate of le-4, an effective batch size of 16, dropout set to 0.0, and no weight decay.
Training For both the 30% and 60% sequence identity threshold splits, we use a batch size of 16, fine-tuning learning rate of le-4, and apply random structural rotation as data augmentation. For 60% sequence identity threshold split, we additionally apply a dropout rate of 0.5. The total epoches on LBA-30 and LBA-60 are 3 and 99. The training process can be efficiently handled using one or two A100 GPU.
D.3 TARGET BINDING AFFINITY PREDICTION
Baselines The baselines are: DeepDTA Öztürk al. 2018 MT-DTI Shhin et al. 2019).
GraphDTA (Nguyen et al. 2021), 1ZMLP-DTA (Qiu et al. (021). EnsembleDD.M (Kao et al.
2021). MGraphDTTA (Yang al. [022]. FusionDTA (Yuan et al. 2022) and NHGNN-DTÄ (He et al.
2023
The reported results are taken from (He et al. 2023).
Training We generate unbound protein structures using AlphaFold 2 and molecular structures using RDKit. Since these structures do not initially reside in the same spatial framework, we remove the edges between proteins and ligands to ensure proper alignment. For all five random splits, an effective batch size of 8, a learning rate of le-4 for ADiT-S and 3e-5 for ADiT-M and ADiT-L, a dropout rate of 0.2 and train for 120 epochs without rotation augmentation. To handle long proteins, we randomly truncate them to a max length of 150. The training process can be efficiently handled using a 4 A100 GPU (for ADiT-L).
D.4 PROTEIN-PROTEIN BINDING AFFINITY PREDICTION
We employ a diverse set of baselines for comparison, including ESM-IF (Hsu et al. 2022). self-attention-bbseed End-to-End network (End-to-End) End) (Jumper et al. 2021). DDGPred (Shan
Baselines
et al. 2022), MIF-Network (Yang et al. 2020). RDE-Network (Luo et al. 2023 DiffAffinity (Li
al. 2024), Prompt-DDG (Wu et al. (2024). and Surface- Surface-VQMAE (Wu & 2024) The reported
results are sourced from (Wu et al. 2024 Wu & Li 2024).
Training We truncate sequences based on the distance between the current residue and mutation- related residues, limiting the remaining sequence length to 50 residues. Training conducted with an effective batch size of 8, a learning rate of 3e-5, a dropout rate of 0.0, and a total of 60 epochs. The training process can be efficiently handled using a single A100 GPU.
D.5 ANTIBODY-ANTIEEL BINDING AFFINITY PREDICTION
Baselines We compare ADiT against several strong baselines for antibody affinity prediction, including Bind-ddG (Shan et al. 2022), and GearBind (Cai et al. 2024). We also include the physics-
19

---

Published as a conference paper at ICLR 2026
based tool FoldX (Delgado al. 2019) additional baseline. The performance of most baselines is reported from (Cai et al. 2024).
Case study Affinity maturation aims to enhance the binding affinity of an antibody to its target antigen. As the search space for antibodies is too vast for wet-lab enumeration, in silico methods increasingly being used to increase success rates. We can formulate the in silico affinity maturation task virtual screening problem, where we use model to rank mutants of the starting antibody by their binding affinity to the target antigen, and validate the top-ranked mutants in wet lab. The two antibodies Anti-5T4 UdAb and CR3022 were chosen in (Cai al., 2024) because (1) they cover two common antibody formats, nanobody and antibody; (2) they cover two antigens: UdAb targets an oncofetal antigen 5T4, while CR3022 targets SARS-CoV- and multiple variants of SARS-CoV-2. The 780 mutations for Anti-5T4 UdAb are obtained from saturation single-point mutagenesis of residues I, 3, 25, 27-33, 39-45, 52-57, 59, 91-93, 95, 99-103, 105, 110, 112. 115-117 (i.e. mutating one of these sites to every possible amino acid type: each site yields 19 mutants), which are residues close to the antibody interface to 5T4. The 1420 mutations for CR3022 are obtained from saturation single-point mutagenesis of residues H25-35, H50-66, H99-108, 1.24-40, 155662, 195-103.
E ENZYME COMMISSION (EC) NUMBER ANNOTATION TASK
To show ADIT's generalization and applicability to non-binding tasks, we run additional experiments on the Enzyme Commission (EC) number annotation task (Glioonijevié al., 202I), which involves predicting whether a protein can catalyze chemical reaction. This task is evaluated using the Fmax metric. The results, presented in Table show that our method consistently outperforms the baseline, even under challenging data splits, demonstrating ADiT's applicability beyond hinding-related functions.
Table 7: Enzyme Commission (EC) Number Annotation Task.
METHOD
30%
40%
50%
70%
DEEPFRI
0.470
0.505
0.545
0.600
GEARNET-EDGE (MULTIVIEW CONTRAST)
0.744
0.769
0.808
0.848
ADIT-S
0.751
0.773
0.806
0.843
F INTERPRETABILITY
Interpretabilityy is crucial next step toward understanding the mechanisms underlying our model's effectiveness. In this work, due to project scope, we primarily focus on establishing performance and scalability. comprehensive and systematic interpretability investigation is left for future research. As preliminary exploration, we conduct a simple case study examining the attention patterns learned by our trained model. We randomly select the 5f63 pocket-ligand complex from the LBA60 test set, which consists of 18 residues in the binding pocket and 41 atoms in the ligand. In this analysis, we extract attention scores from the token-level transformer, since the atom encoder and decoder employ only local block attention.
Several observations arise from this initial investigation:
First, we observe that the average attention weight for intra-chain pairs is 0.0253, compared to only 0.0056 for inter-chain pairs. This suggests that, although the model captures cross-chain interactions, it places greater emphasis on intra-chain relationships, which may play more critical role in binding context representation.
Second, the model's attention mechanism appears to effectively "focus" on key interfacing residues and atoms. We visualize the cross-chain attention scores from the second layer (layer index=1, zero- zero-based) in Figure5] To identify potentially important residues and atoms, highlight the closest one-third of residue tokens to the ligand and the closest one- fifth of ligand atom tokens to the protein. As shown in the results for Head O and Head 4, these highlighted residues exhibit higher attention
20

---

Published as a conference paper at ICLR 2026
Figure 5: Attention score visualization for the 5f63 pocket-ligand ligand complex, comprising 18 residue tokens and 41 ligand atom tokens. The plot highlights residue-aton interaction patterns. Here, on the x- or y-axis denotes the closest one-third of residue tokens to the ligand, and on the x- or y-axis denotes the closest one-fifth of ligand atom tokens to the protein.
weights (columns appear brighter). Similarly, the highlighted ligand atoms receive higher attention in Head 1, Head 3, and Head 5.
21


Reviews:
[{"review_id": "KhmHiNsBc6", "reviewer": "Reviewer_r7bx", "summary": "This paper repurposes AlphaFold3 to a representation learning model and then achieves protein binding affinity prediction. The paper applies two-stage learning algorithm with first pretraining a general model and then finetuning the pretrained model on specific tasks. The pretrained model is finetuned and evaluated on four tasks: protein-ligand, drug-target, protein-protein and antibody-antigen binding predictions.", "questions": "1. In 4.1, the paper mentioned they curated 427,947 protein-ligand interaction examples from PDB. However, generally there are just around 20k protein-ligand complexes in PDB as introduced in previous works. Can the authors explain how they curated each kind of complex data from the PDB?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea of repurposing AlphaFold3 to a binding prediction model is very interesting and useful in real world.", "weaknesses": "The weaknesses of this paper are listed as follows:\n\n1. I appreciate the idea of the building upon AlphaFold3 achieves protein binding prediction and the proposed method is evaluated across several important tasks. However, the paper finetuned four tasks and there are four task-specific models, which is inconvenient for practical use in reality. If there are even more tasks like enzyme-cofactor, enzyme-substrate, target protein-binder affinity prediction,  that means there are will be more models. If different model size is considered, there will be ever more models. That would really reduce the utility of the proposed method. If the paper could merge all downstream tasks together and finetune the pretrained model on all merged tasks, I believe that might both increase the  performance due to multitask learning and also improve the convenience for practical use.\n\n2. The paper lacks some baseline models for affinity prediction, like Rosseta, NERE [1].\n\n[1] Unsupervised Protein-Ligand Binding Energy Prediction via Neural Euler’s Rotation Equation.\n\n3. AlphaFold3 itself provides a metric ipTM to evaluate the binding interface geometry of the binding molecules, which can tell the binding affinity to some extent. On one hand, the author doesn't explain why they need to construct a different metric based on AlphaFold3 instead of directly using ipTM. On the other hand, the proposed method should compare with ipTM.\n\n4. Since the proposed method heavily relies on AlphaFold3, the paper should clearly explain their difference from AF3, like what modules are reused and what modules are newly proposed in this paper.\n\n5. The performance on some tasks are worse than the baseline methods like protein-protein interaction.\n\n6. Some small issues like grammar error in line 214, \"We employs ...\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iXg75XqELh", "reviewer": "Reviewer_Gvv4", "summary": "The paper introduces ADiT, an atom-level diffusion transformer that repurposes AlphaFold 3-style machinery from structure generation to representation learning for binding affinity prediction, jointly encoding sequence and 3D geometry without MSAs or templates. Pretrained on PDB and evaluated across protein–ligand, drug–target, PPI, and antibody–antigen tasks, ADiT achieves state-of-the-art or competitive performance, scales with model size, and identifies wet-lab validated affinity-enhancing antibody mutations.", "questions": "1.  Do you train ADiT entirely from scratch or initialize any modules with AlphaFold 3 parameters?\n\n2.  How is Token2Atom implemented: are atom features a direct copy from their parent token, or do you use learned weights/attention or distance-aware mixing? Please provide formulas and ablations.\n\n3.  What empirical or theoretical reasons justify a non-SE(3)-equivariant design for affinity prediction?\n\n4.  What is the performance impact of rotation augmentation (and centroid-centering) versus no augmentation? \n\n5.  Why rely solely on structure-denoising? Have you compared against masked atom/residue prediction, contrastive objectives, or multi-task pretraining?\n\n6.  What factors drive ADiT’s underperformance on protein-protein affinity benchmarks relative to BA-DDG/Light-DDG?\n\n7.  What are the training compute (FLOPs/GPU hours), wall-clock times, parameter counts, peak memory, and throughput for each model size?\n\n8.  How do ADiT’s training and inference speeds compare to strong baselines (including all-atom and residue-level models) on standardized batch/complex sizes? Please include scaling curves with atom count.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.  The work adopts an all-atom representation and enables interaction between atom-level and token-level features. The all-atom representation unifies the encoding of various biomolecular types, facilitating knowledge transfer across different interaction prediction tasks.\n\n2.  The work  leverages the atom- and token-level representation modules of AF3 while removing dependencies on MSAs and templates, and simplifying the Pairformer blocks to enhance efficiency.", "weaknesses": "1.  When using AlphaFold 3, is the model trained from scratch or initialized with AF3 parameters? The Atom2Token average pooling is simple and reasonable, but the Token2Atom operation is insufficiently explained, does it simply assign token features to all constituent atoms, or are weights/attention mechanisms used?\n\n2.  The performance is relatively weak, at comparable model sizes, ADiT shows no clear advantage, especially given the substantial pretraining cost. On protein–ligand tasks ADiT-S/M trails lighter, task-specific baselines such as CheapNet and BindNet, and on protein–protein binding affinity even the largest variant (ADiT-L) underperforms recent state-of-the-art models like BA-DDG and Light-DDG. This raises concerns about the practical benefit of the proposed pretraining strategy.\n\n3.  The authors emphasize that the model is non-equivariant but do not explain why abandoning SE(3)-invariance is advantageous for binding prediction. Appendix E.1 mentions random rotation augmentation (centered on protein centroids) to encourage robustness and generalization, yet no ablation study is provided to support.\n\n4.  The model uses only a structure-denoising task, which may be insufficient to capture complex features relevant to binding affinity. It would be beneficial to incorporate additional self-supervised tasks such as masked atom/residue prediction.\n\n5.  The paper lacks efficiency analysis and training time reporting. All-atom representation models are typically computationally expensive. It would be helpful to include comparisons of ADiT’s training and inference speeds with other baselines", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cQXBO2uQla", "reviewer": "Reviewer_nSSR", "summary": "The model proposes an all-atom foundation model for the representation learning of biomolecular interactions. The interaction structure is tokenized into atom level and token (protein residue or ligand atom) level representation, in both single and pair-wise manners. The backbone of the model is a diffusion transformer that involves hierarchical atom-token feature aggregation and unpooling. The model is pre-trained on general protein-protein interaction data from PDB and then fine-tuned for multiple downstream prediction tasks: protein-ligand, protein-protein and antibody-antigen interactions, achieving competitive results.", "questions": "- How is Token2Atom implemented?\n\n- Some interpretability is needed. For example, does the attention capture atom or residue level patterns that facilitate the interaction?\n\n- Appendix E3: if the protein-ligand edges are removed, how is interaction information represented and captured? In this way, the problem decomposes into one that separate embeds protein and drug (dual encoders) then combines them for prediction. How would atom level and pairwise information contribute in this scenario?\n\n\n\nMinor:\n- As the model didn't seem to rely on AlphaFold 3's tokenization or trained parameter, only some of the architecture, it is not really \"repurposing AlphaFold 3\".\n\n- Protein-ligand interaction and drug-target interaction largely overlap, so it would be clearer if the two experiments are defined as different datasets instead of different tasks.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The hierarchical atom and token level representation is a good solution that addresses the balance between expressiveness and computational efficiency. The methods are clearly formulated and presented.\n\n- The experimental results show clear improvements on protein-drug interaction prediction tasks and also competitive performance on protein-protein interactions. Also, the comparison between small, medium and large models demonstrates the scalability of the model.\n\n- Additional case studies such as in silico antibody affinity maturation show good generalizability and applicability of the learned patterns.", "weaknesses": "See Questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "86lrxyCp3Z", "reviewer": "Reviewer_kRrt", "summary": "This paper proposes the Atom‑level Diffusion Transformer (ADiT), a unified foundation model that aims to learn transferable representations for biomolecular complexes and predict binding affinity across diverse interaction types. The authors repurpose the AlphaFold 3 architecture for representation learning by:\n\n- **Simplifying inputs.** they remove multiple‑sequence alignments and template conditioning, instead using a unified tokenization scheme where each protein residue and each heavy atom in a small molecule is treated as a token and enriched with features from the ESM‑2 protein language model.\n\n- **Using diffusion transformers.** the model alternates between atom‑level and token‑level DiT blocks, performs atom‑to‑token pooling and token‑to‑atom unpooling, and applies a denoising pre‑training objective where Gaussian noise is added to atom coordinates.\n\n- **Training at scale.** three model sizes (12M, 35M and 253M parameters) are pre‑trained on ~480k protein‑protein, ~433k protein‑ligand and ~427k protein–ligand structures from the PDB without labels. Fine‑tuning is performed on downstream tasks.\n\n- **Evaluating on four benchmarks.** the model is tested on protein‑ligand, drug‑target, protein‑protein and antibody-antigen tasks. According to the authors, the largest variant ADiT‑L achieves the best or comparable results across tasks, shows performance improvements with larger model size, and assigns better ranks to affinity‑enhancing antibody mutations in a case study.", "questions": "- How did you ensure that there is no overlap between pre‑training complexes and those used in downstream evaluation? Will you evaluate ADiT on leak‑free datasets (e.g., CleanSplit) or the synthetic–data‑augmented GatorAffinity benchmark to assess generalization?\n\n[GatorAffinity]: GatorAffinity: Boosting Protein-Ligand Binding Affinity Prediction with Large-Scale Synthetic Structural Data\n\n- Given evidence that equivariant networks like ELGN improve binding affinity prediction by capturing rotational symmetries, why did you choose a non‑equivariant architecture and not include geometric priors? \n\n[ELGN]: Equivariant Line Graph Neural Network for Protein-Ligand Binding Affinity Prediction\n\n- Can ADiT be extended to DNA/RNA–protein interactions or multi‑chain assemblies? Does the unified tokenization scheme generalize to nucleic acids or small peptides? \n\n- Boltz‑2 (2025) jointly predicts structures and binding affinities, matches FEP accuracy, and includes dynamic ensemble data. Did you benchmark ADiT against Boltz‑2? If not, why?\n\nIf the authors address the weaknesses outlined above, or provide a right justification for the points raised in my questions, or clarify points I may have misunderstood, I would be willing to raise my score.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- ADiT attempts to unify representation learning across protein–protein, protein–ligand and antibody–antigen interactions. The use of a single architecture trained on all‑atom inputs is novel and aligns with the trend of foundation models.\n\n- The authors detail a simplified featurization scheme that avoids heavy reliance on MSAs and templates while leveraging pre‑trained ESM embeddings. They clearly describe the hierarchical DiT architecture (Fig. 2) with atom‑to‑token and token‑to‑atom operations, diffusion transformer blocks and gating mechanisms, and provide algorithmic descriptions in the appendices.\n\n- Results are presented on four different tasks. On the protein–ligand benchmark, ADiT‑L outperforms previous models such as ProNet and ProFSA. The authors provide an ablation table and pre‑training loss curves (Fig. 4) demonstrating scaling benefits.\n\n- The antibody affinity maturation case illustrates that ADiT assigns lower (better) ranks to most experimentally validated affinity‑enhancing mutations compared with the smaller variants. Figure 3 in the paper visualizes the ranking of mutations for two antibodies, and the authors claim this suggests utility for antibody design.\n\n- Figures 1–4 are crisp, readable, and self-contained (architecture overview, module layout, case study, scaling/ablation), and each experiment section cleanly separates Setup and Results. This structure makes assumptions, metrics, and splits explicit, significantly improving readability.", "weaknesses": "- **Claim of being the first all‑atom “foundation model” might be overstated.** The paper positions ADiT as a “first step toward developing all‑atom structure foundation models”. However, Boltz‑2 is a biomolecular all-atom foundation model, I believe. Because Boltz‑2 pre‑dates ICLR 2026, the authors should compare to it and clarify what is novel beyond prior work. \n\n- **Limited attention to data leakage and generalization.** Deep models trained on PDBbind and similar benchmarks are known to memorize training complexes due to train–test contamination. A 2025 study in Nature Machine Intelligence introduced the CleanSplit dataset to remove overlaps and found that performance of many models drops significantly when leakage is eliminated, suggesting poor generalization. The authors curate PDB clusters but do not quantify overlap between pre‑training and fine‑tuning sets or evaluate ADiT on leak‑proof benchmarks (e.g., CleanSplit). Without such evaluation, it is unclear if the model genuinely learns transferable representations rather than memorizing structural motifs. Currently, this is the main concern in the field.\n\n[CleanSplit]: Resolving data bias improves generalization in binding affinity prediction\n\n- **Potential over‑claim about scaling.** While larger ADiT variants exhibit modest improvements (e.g., Pearson increases from 0.626 to 0.645 on LBA‑30 and from 0.690 to 0.751 on the Davis dataset), the gains are sometimes small relative to parameter increases. The ablation study shows a 3.8% improvement in Pearson when using all‑atom information, but no comparison is made against equally large equivariant models or models trained on bigger synthetic datasets. It remains unclear whether scaling alone, rather than improved data or architecture, drives gains.\n\n- **Unclear justification for ESM‑2 features and neglect of small‑molecule semantics.** ADiT uses the ESM‑2 (650 M) language model to embed protein residues; however, this model is trained only on protein sequences and may not capture ligand or nucleic‑acid features. For small molecules, the token type embedding is set to zero, meaning no evolutionary or contextual information is added.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)