EDM coordinate-denoiser preconditioning

Your task

Function signature

edm_preconditioning(
    x_noisy: torch.Tensor,
    network_update: torch.Tensor,
    x_target: torch.Tensor,
    sigma: torch.Tensor,
    sigma_data: torch.Tensor,
    valid_mask: torch.Tensor,
    training_loss: bool = True,
    sigma_mode: str = 'example',
    log_eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

Both denoising paths contribute

import torch

x_noisy = torch.tensor([[[2.0, -4.0, 6.0]]], dtype=torch.float64, requires_grad=True)
network_update = torch.tensor([[[-1.0, 2.0, 0.5]]], dtype=torch.float64, requires_grad=True)
x_target = torch.tensor([[[0.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
sigma = torch.tensor([3.0], dtype=torch.float64, requires_grad=True)
sigma_data = torch.tensor(4.0, dtype=torch.float64, requires_grad=True)
valid_mask = torch.tensor([[True]], dtype=torch.bool)
training_loss = True
sigma_mode = 'example'
log_eps = 1e-08

result = edm_preconditioning(x_noisy=x_noisy, network_update=network_update, x_target=x_target, sigma=sigma, sigma_data=sigma_data, valid_mask=valid_mask, training_loss=training_loss, sigma_mode=sigma_mode, log_eps=log_eps)

Expected return

{
    'preconditioned': torch.tensor([[[0.4, -0.8, 1.2000000000000002]]], dtype=torch.float64),
    'denoised': torch.tensor([[[-1.1200000000000003, 2.2400000000000007, 5.04]]], dtype=torch.float64),
    'c_in': torch.tensor([0.2], dtype=torch.float64),
    'c_skip': torch.tensor([0.64], dtype=torch.float64),
    'c_out': torch.tensor([2.4000000000000004], dtype=torch.float64),
    'noise_condition': torch.tensor([-0.07192051811294523], dtype=torch.float64),
    'weight': torch.tensor([0.1736111111111111], dtype=torch.float64),
    'weighted_loss': torch.tensor([1.8329629629629631], dtype=torch.float64),
    'loss': torch.tensor(1.8329629629629631, dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Given x_noisy, network_update, x_target [B,A,3], sigma_data a positive scalar tensor, valid_mask [B,A], and sigma [B] when sigma_mode='example' (default), or [B,A] for 'atom', compute d=\sigma^{2}+\sigma_{\mathrm{data}}^{2}; c_{\mathrm{in}}=d^{-1/2}, c_{\mathrm{skip}}=\sigma_{\mathrm{data}}^{2}/d, c_{\mathrm{out}}=\sigma\cdot \sigma_{\mathrm{data}}/\sqrt{d}.
• Return \mathrm{preconditioned}=c_{\mathrm{in}}\cdot x_{\mathrm{noisy}} and \mathrm{denoised}=c_{\mathrm{skip}}\cdot x_{\mathrm{noisy}}+c_{\mathrm{out}}\cdot \mathrm{network\_update}, broadcasting over xyz and, for example sigma, atoms.
• The supplied network_update is F(c_{\mathrm{in}}\cdot x, \mathrm{noise\_condition}), not \epsilon or a complete denoised coordinate; do not add another residual connection. Clear invalid atoms in both coordinate outputs.
• \mathrm{noise\_condition}=\operatorname{log}(\operatorname{max}(\sigma,\mathrm{log\_eps})/\sigma_{\mathrm{data}})/4, \mathrm{log\_eps}>0 defaults to 10^{-8}. Coefficient, noise_condition, and weight outputs retain sigma's shape and are not masked.
• With training_loss=True (default), all sigma entries must be > 0 and \mathrm{weight}=(\sigma^{2}+\sigma_{\mathrm{data}}^{2})/(\sigma\cdot \sigma_{\mathrm{data}})^{2}.
• weighted_loss [b] is the sum over valid atoms of \mathrm{weight}\cdot \sum_{\mathrm{xyz}}((\mathrm{denoised}-x_{\mathrm{target}})^{2}), divided by 3\cdot \operatorname{max}(\mathrm{valid\_count},1). loss is the mean of weighted_loss over all B examples, including empty examples as zeros.
• With training_loss=False, \sigma\ge 0 is allowed and weight, weighted_loss and loss are zero; bypass squared-error arithmetic entirely. At \sigma=0 the valid denoised coordinates equal x_noisy exactly. No \epsilon enters coefficients or the training weight; log_eps is only inside the log clamp.
• There is no alignment, centering, extra atom weighting, or learned callback. valid_mask=True means included; finite padding is ignored. All tensors share dtype/device; B,A\ge 1. The only keys are sigma_mode, log_eps and training_loss.

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 edm_preconditioning(
    x_noisy: torch.Tensor,  # (B, A, 3)
    network_update: torch.Tensor,  # (B, A, 3)
    x_target: torch.Tensor,  # (B, A, 3)
    sigma: torch.Tensor,  # (B,) for sigma_mode="example"; (B, A) for "atom".
    sigma_data: torch.Tensor,  # Scalar tensor, shape ().
    valid_mask: torch.Tensor,  # (B, A)
    training_loss: bool = True,
    sigma_mode: str = 'example',
    log_eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    preconditioned: (B, A, 3).
    denoised: (B, A, 3).
    c_in: Exactly sigma.shape: (B,) for sigma_mode="example", or (B, A) for "atom".
    c_skip: Exactly sigma.shape: (B,) for sigma_mode="example", or (B, A) for "atom".
    c_out: Exactly sigma.shape: (B,) for sigma_mode="example", or (B, A) for "atom".
    noise_condition: Exactly sigma.shape: (B,) for sigma_mode="example", or (B, A) for
        "atom".
    weight: Exactly sigma.shape: (B,) for sigma_mode="example", or (B, A) for "atom".
    weighted_loss: (B,).
    loss: Scalar tensor, shape ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement edm_preconditioning")

Sources and further reading

• google-deepmind/alphafold3 · diffusion_head.py
• jwohlwend/boltz · diffusion.py
• jwohlwend/boltz · diffusionv2.py
• RosettaCommons/foundry · RFD3_diffusion_module.py
• RosettaCommons/foundry · losses.py