Possible bug in `smooth_lddt_loss`: averaging `eps` over batch before masking

[Aldrich-ux]

I believe there is a semantic issue in the implementation of smooth_lddt_loss related to batch averaging.
Below is the original code (unchanged):

# Compute epsilon values
eps = (
    (
        (
            F.sigmoid(0.5 - dist_diff)
            + F.sigmoid(1.0 - dist_diff)
            + F.sigmoid(2.0 - dist_diff)
            + F.sigmoid(4.0 - dist_diff)
        )
        / 4.0
    )
    .view(B, N, N)
    .mean(dim=0)
)

# Calculate masked averaging
num = (eps * mask).sum(dim=(-1, -2))
den = mask.sum(dim=(-1, -2)).clamp(min=1)
lddt = num / den

• dist_diff has shape [B, N, N]
• mask also has shape [B, N, N] and is per-sample, since it depends on coords_mask
• However, eps is averaged over the batch dimension using .mean(dim=0), resulting in shape [N, N]

This means each sample uses batch-averaged pairwise scores, even though the mask is sample-specific.
When coords_mask differs across samples or B > 1, this no longer matches the intended lDDT / smooth-lDDT semantics.

Thanks a lot for the great work on this project, and I hope this report is helpful.

[yuyangw]

Hi @Aldrich-ux, thanks for pointing this out. This indeed leads to confusion where eps wasn't intended to be averaged in the beginning. However, this doesn't lead to wrong smooth LDDT loss since in each step, a batch in each GPU keeps copies of the same protein (with different rotation and flow timestep). The mask is the same across all samples within the batch. Having said that, we'll have it fixed soon to avoid confusion.