DiffDock three-space score-matching loss

Your task

Function signature

diffusion_score_matching_loss(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    apply_mean: bool = True,
    no_torsion: bool = False,
) -> dict[str, torch.Tensor]

Worked example

Three spaces have three different normalization rules

import torch

activations = {
    'torsion_prediction': torch.tensor([3.0, -1.0, 4.0], dtype=torch.float64, requires_grad=True),
    'torsion_target': torch.tensor([1.0, 1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'torsion_score_norm_sq': torch.tensor([2.0, 8.0, 3.0], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'component_weights': torch.tensor([0.5, 2.0, 3.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'graph_mask': torch.tensor([True, True], dtype=torch.bool),
    'torsion_mask': torch.tensor([True, True, True], dtype=torch.bool),
}
structure = {
    'translation_prediction': torch.tensor([[2.0, 0.0, -1.0], [1.0, 3.0, 2.0]], dtype=torch.float64, requires_grad=True),
    'translation_target': torch.tensor([[1.0, -1.0, 0.0], [0.0, 1.0, 2.0]], dtype=torch.float64, requires_grad=True),
    'rotation_prediction': torch.tensor([[1.0, 2.0, 0.0], [-1.0, 0.0, 2.0]], dtype=torch.float64, requires_grad=True),
    'rotation_target': torch.tensor([[0.0, 0.0, 0.0], [1.0, -1.0, 0.0]], dtype=torch.float64, requires_grad=True),
    'translation_sigma': torch.tensor([2.0, 0.5], dtype=torch.float64, requires_grad=True),
    'rotation_score_norm': torch.tensor([0.5, 2.0], dtype=torch.float64, requires_grad=True),
    'torsion_graph': torch.tensor([0, 0, 1], dtype=torch.int64),
}
apply_mean = True
no_torsion = False

result = diffusion_score_matching_loss(activations=activations, parameters=parameters, masks=masks, structure=structure, apply_mean=apply_mean, no_torsion=no_torsion)

Expected return

{
    'loss': torch.tensor([14.020833333333334], dtype=torch.float64),
    'translation_loss': torch.tensor(2.2083333333333335, dtype=torch.float64),
    'rotation_loss': torch.tensor(3.7083333333333335, dtype=torch.float64),
    'torsion_loss': torch.tensor([1.8333333333333333], dtype=torch.float64),
    'translation_base_loss': torch.tensor(1.5416666666666665, dtype=torch.float64),
    'rotation_base_loss': torch.tensor(0.08333333333333333, dtype=torch.float64),
    'torsion_base_loss': torch.tensor([0.3194444444444444], dtype=torch.float64),
    'torsion_counts': torch.tensor([2, 1], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement the \mathrm{original}-\mathrm{v1} three-space score-matching objective from supplied predictions, physical score targets and normalization values. B\ge 1 graphs have translation_prediction/translation_target and rotation_prediction/rotation_target [B,3].
• K\ge 0 flattened torsions have torsion_prediction/torsion_target [K] and int64 torsion_graph [K]. graph_mask [B] and torsion_mask [K] are bool with True meaning valid. All real tensors use one common \mathrm{float32}/64 CPU dtype; inputs are never mutated.
• component_weights [3] contains finite nonnegative translation/rotation/torsion coefficients, applied as a plain weighted SUM without normalization. The source function defaults to unit coefficients, while its CLI used 0.33 each; here every weight is explicitly supplied. Zero weight does not disable a space or its input checks.
• translation_sigma [B], rotation_score_norm [B] and torsion_score_norm_sq [K] are strictly positive and finite on active entries. rotation_score_norm means original DiffDock \sqrt{\mathbb E[\lVert s_{\mathrm{rot}}\rVert^2]/\pi}, not mathematical RMS.
• torsion_score_norm_sq means per-BOND \mathbb E[s_{\mathrm{tor}}^2], the output meaning of original torus.score_norm, not its square root and not a quantity to square again. The source derives these torsion moments from per-edge noise scales; do not substitute a graph's rotation normalizer or a graph-level torsion width.
• No density lookup or random draw occurs here. Targets are physical scores, not noise predictions; for example a supplied translation Gaussian \epsilon would first give the target -\epsilon/\mathrm{translation\_sigma} outside this loss.
• For graph b define e_{\mathrm{tr}}[b]=(\sigma_{\mathrm{tr}}[b]^{2}/3)\cdot \sum_{c}(\mathrm{pred\_tr}-\mathrm{target\_tr})^{2} and e_{\mathrm{rot}}[b]=(1/(3\cdot n_{\mathrm{rot}}[b]^{2}))\cdot \sum_{c}(\mathrm{pred\_rot}-\mathrm{target\_rot})^{2}.
• For bond k define e_{\mathrm{tor}}[k]=(\mathrm{pred\_tor}-\mathrm{target\_tor})^{2}/n_{\mathrm{tor\_sq}}[k]. Define each base error by replacing its prediction with zero, not by subtracting a baseline or changing the component weight. Config apply_mean defaults to True and no_torsion defaults to False; both must be bool.
• When torsions are enabled, torsion_mask=True entries require a graph index in [0,B); the bond is active only if that graph's graph_mask is also True. Masked torsion indices may be arbitrary. no_torsion=True disables all bonds, ignores their values/normalizers/indices, and returns zero bond counts and torsion losses.
• Inactive graph/bond scores and scales may contain NaN/Inf or zero normalizers: replace them before subtraction, squares or division.
• For apply_mean=True, translation and rotation losses are means over valid graphs (already means over three components). Torsion loss is the mean over ALL active BONDS, not the mean of per-graph torsion means. Use the same reductions for base losses.
• Return scalar [] translation/rotation losses and bases, [1] torsion loss/base, and [1] loss, preserving the source's mixed scalar/vector shapes. An empty active graph or bond mean is defined as zero with zero gradients. This finite-zero empty global torsion convention deliberately replaces the upstream enabled-empty NaN.
• For apply_mean=False, return every loss/base as [B].
• Keep e_tr/e_rot per graph, and divide each graph's active torsion-error SUM by (\text{its bond count}+0.0001), exactly retaining the original count \epsilon even when the count is positive; do not use \operatorname{max}(\mathrm{count},1), an unbiased mean, or add \epsilon to the physical score normalizers.
• Graphs with no active torsions have zero torsion loss/base. Invalid graphs return zero components.
• In either mode: \mathrm{loss}=w_{\mathrm{tr}}\cdot \mathrm{translation\_loss}+w_{\mathrm{rot}}\cdot \mathrm{rotation\_loss}+w_{\mathrm{tor}}\cdot \mathrm{torsion\_loss}. Return torsion_counts [B] as int64 active bond counts. Base/component losses are unweighted diagnostics, never subtracted from loss.
• All returned floating quantities remain differentiable in this teaching contract; upstream detaches logging metrics/baselines. Supplied targets and normalizers receive mathematical derivatives here, without implying differentiation through source cached-table indexing or training-data preparation.
• All active values and intermediate results must be finite and representable. Invalid listed shapes/dtypes/configuration types, active graph indices, nonfinite active inputs, nonpositive active normalizers, or invalid component weights raise ValueError.
• Complete training execution, score-network inference and target-distribution construction are separate.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def diffusion_score_matching_loss(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    apply_mean: bool = True,
    no_torsion: bool = False,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    loss: (1) or (B,).
    translation_loss: Scalar tensor, shape (). or (B,).
    rotation_loss: Scalar tensor, shape (). or (B,).
    torsion_loss: (1) or (B,).
    translation_base_loss: Scalar tensor, shape (). or (B,).
    rotation_base_loss: Scalar tensor, shape (). or (B,).
    torsion_base_loss: (1) or (B,).
    torsion_counts: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement diffusion_score_matching_loss")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · training.py
• gcorso/DiffDock · torus.py
• gcorso/DiffDock · so3.py
• gcorso/DiffDock · parsing.py