Technical Summary: ABOPD – Antibody CDR Design via On-Policy Distillation

1. Problem Statement

The computational design of antibody Complementarity-Determining Regions (CDRs) is a critical challenge in protein engineering, as these flexible loops determine antigen recognition. While recent diffusion-based generative models have advanced the field, a fundamental mismatch exists between the training objective and the inference process:

• Training Distribution: Standard denoising training supervises models on "reference-derived states," which are noisy structures generated by perturbing native (ground truth) CDRs.
• Inference Distribution: During generation, the model recursively produces states based on its own previous predictions.
• The Mismatch: As the generation trajectory proceeds, prediction errors accumulate, causing the model to visit states that deviate significantly from the reference-derived distribution used during training. This is particularly detrimental for flexible CDRs like CDR-H3, where small backbone deviations can propagate, leading to steric clashes and incorrect antigen-facing geometries.

Existing post-training strategies, such as Supervised Fine-Tuning (SFT) or offline distillation, continue to operate on the reference-derived distribution, failing to correct errors that arise specifically along the model's own generation trajectory.

2. Methodology: ABOPD Framework

The authors propose ABOPD (Antibody design framework built on On-Policy Distillation), a post-training strategy designed to align structural supervision with the states actually encountered during recursive generation. The framework consists of three main stages:

A. Hybrid Pretraining (H-DiffAb)

To establish a robust initialization, the authors pretrain a single model, H-DiffAb, across diverse masking scenarios: single-CDR, multi-CDR, and all-six-CDR target masks. This model serves as the initialization for both the student and the teacher, ensuring fair comparisons.

B. Backbone-Aware Teacher Adaptation

A "Teacher" model is derived from H-DiffAb by incorporating privileged information available only during training: the native geometry of the target residues.

• Architecture: The teacher uses lightweight residual MLP adapters to encode native $C_\alpha$ coordinates, atom validity masks (N, $C_\alpha$, C, O, $C_\beta$), and target-to-complex distance features.
• Function: These privileged descriptors condition the denoising prediction without altering the diffusion state representation. The teacher is optimized on reference-derived states to minimize coordinate and orientation errors, resulting in a model with superior geometric prediction capabilities compared to the base H-DiffAb.

C. On-Policy Distillation (The Core Innovation)

The "Student" model is post-trained using a mixed-objective approach:

1. Student Rollout: The student generates reverse-denoising trajectories (rollouts) starting from fully noisy states.
2. State Matching: At specific timesteps ($t \in \{80, 70, \dots, 5\}$), the frozen Teacher evaluates the same student-generated state ($R^S_t$) that the student is currently processing.
3. Coordinate-Transition Distillation: The student is trained to match the Teacher's predicted coordinate noise ($\hat{\epsilon}^{pos}$) on these rollout states. This is formulated as a reverse-KL divergence minimization, simplified to an $L_2$ loss on coordinate noise predictions:
   $$\mathcal{L}_{OPD}^{pos} = \mathbb{E} \left[ \sum_{t \in \mathcal{T}} \sum_{i \in \mathcal{G}} \| \hat{\epsilon}^{pos}_{S, t, i} - \hat{\epsilon}^{pos}_{T, t, i} \|^2_2 \right]$$
   Crucially, this supervision is applied only to $C_\alpha$ coordinates, as they directly control backbone geometry and loop placement. Sequence and orientation predictions are not distilled in this step to avoid overfitting to the teacher's potentially imperfect sequence predictions.
4. Offline Denoising Anchor: To preserve the joint sequence-structure-orientation objective learned during pretraining, the original denoising loss ($\mathcal{L}_{anchor}$) is retained on independently sampled reference-derived states.
5. Final Objective: $\mathcal{L}_{ABOPD} = \mathcal{L}_{anchor} + \beta \mathcal{L}_{OPD}^{pos}$ (with $\beta=0.6$).

3. Key Contributions

1. Identification of State-Distribution Mismatch: The paper highlights that the gap between reference-derived states and student rollout states widens during the denoising process, particularly in later refinement stages where geometric corrections are most critical.
2. ABOPD Framework: The proposal of an on-policy distillation method that combines $C_\alpha$ coordinate-transition distillation on student rollout states with an offline denoising anchor.
3. Backbone-Aware Teacher: The demonstration that a teacher model augmented with privileged native geometry provides significantly more accurate geometric targets than standard models, even if its sequence accuracy is not superior.

4. Experimental Results

The method was evaluated on two benchmarks: Simultaneous multi-CDR redesign (19 SAbDab complexes) and RAbD CDR-H3 generation (60 target cases).

• RAbD CDR-H3 Performance: ABOPD reduced the $C_\alpha$ RMSD from 2.37 Å to 1.95 Å (a 0.42 Å improvement) compared to the H-DiffAb baseline. It also outperformed both SFT and offline distillation controls.
• Multi-CDR Redesign: ABOPD improved RMSD across all six CDRs, with the most significant gain on CDR-H3 (2.94 Å $\to$ 2.48 Å). It also reduced intra-CDR and CDR–antigen clash rates significantly (from 0.93% to 0.42% and 1.58% to 0.67%, respectively).
• Geometric Fidelity: Improvements were observed in lDDT, TM-score, and DockQ, indicating better local geometry, global heavy-chain structure, and interface placement.
• Mechanism Analysis: Timestep-level analysis revealed that the "Teacher advantage" (the error reduction provided by the teacher over the base model) becomes positive only in later timesteps ($t \le 50$). This confirms that on-policy supervision is most effective during the later stages of geometric refinement when the student's trajectory has already diverged from the reference distribution.
• Ablation Studies: Distilling only the coordinate branch yielded the best results. Adding sequence or orientation distillation provided no further benefit, and in some cases, degraded performance.

5. Significance and Claims

The paper claims that ABOPD offers a path to higher-fidelity protein design by addressing the specific limitations of standard post-training strategies.

• Complementary Approach: The authors position ABOPD as a complementary direction to larger-scale pretraining and new generative architectures. It demonstrates that effective structural post-training depends not just on the supervision signal, but on applying that signal to the correct states (i.e., those visited during generation).
• Modest Scope: The authors explicitly state that the method improves backbone recovery and geometric accuracy but does not necessarily claim to solve all aspects of antibody design (e.g., sequence recovery gains were modest). The primary contribution is the validation of on-policy distillation for correcting accumulated geometric errors in diffusion-based antibody generation.