Masked translational and rotational flow-matching loss

Your task

Function signature

flow_matching_loss(
    x_t: torch.Tensor,
    x1_true: torch.Tensor,
    x1_pred: torch.Tensor,
    R_t: torch.Tensor,
    R1_true: torch.Tensor,
    R1_pred: torch.Tensor,
    t: torch.Tensor,
    masks: dict,
    translation_weight: float = 1.0,
    translation_clip: float | None = 5.0,
    rotation_weight: float = 1.0,
    t_clip: float = 0.9,
    position_scale: float = 1.0,
) -> dict[str, torch.Tensor]

Worked example

Different translation and rotation weights

import torch

x_t = torch.tensor([[[0.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
x1_true = torch.tensor([[[1.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
x1_pred = torch.tensor([[[2.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
R_t = torch.tensor([[[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
R1_true = torch.tensor([[[[0.9800665778412415, -0.19866933079506122, 0.0], [0.19866933079506122, 0.9800665778412415, 0.0], [0.0, 0.0, 1.0000000000000002]]]], dtype=torch.float64, requires_grad=True)
R1_pred = torch.tensor([[[[0.8775825618903728, -0.47942553860420295, 0.0], [0.47942553860420295, 0.8775825618903728, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
t = torch.tensor([0.5], dtype=torch.float64, requires_grad=True)
masks = {
    'valid_mask': torch.tensor([[True]], dtype=torch.bool),
    'diffuse_mask': torch.tensor([[True]], dtype=torch.bool),
}
translation_weight = 0.5
translation_clip = None
rotation_weight = 3.0
t_clip = 0.9
position_scale = 2.0

result = flow_matching_loss(x_t=x_t, x1_true=x1_true, x1_pred=x1_pred, R_t=R_t, R1_true=R1_true, R1_pred=R1_pred, t=t, masks=masks, translation_weight=translation_weight, translation_clip=translation_clip, rotation_weight=rotation_weight, t_clip=t_clip, position_scale=position_scale)

Expected return

{
    'translation_loss': torch.tensor([2.6666666666666665], dtype=torch.float64),
    'rotation_loss': torch.tensor([0.3599999999999998], dtype=torch.float64),
    'per_example': torch.tensor([3.0266666666666664], dtype=torch.float64),
    'loss': torch.tensor(3.0266666666666664, dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs x_t,x1_true,x1_pred [B,N,3], R_t,R1_true,R1_pred [B,N,3,3], t [B] in [0,1], and two bool masks. Let d=1-\operatorname{min}(t,t_{\mathrm{clip}}), where 0<t_{\mathrm{clip}}<1 defaults to .9.
• Translation velocities are (\mathrm{x1}-x_{t})/d; compare their difference after multiplying by \mathrm{position\_scale}>0 (default 1). The shared x_t cancels: form \mathrm{x1\_true}-\mathrm{x1\_pred} directly to preserve this cancellation numerically.
• Evaluate d as \operatorname{max}(1-t,1-t_{\mathrm{clip}}), forming 1-t_{\mathrm{clip}} before casting it to the input dtype, so a clip close to one remains safe.
• Rotational velocities are \operatorname{Log}(R_{t}^{\top} \mathrm{R1})/d for each endpoint; compare these BODY vectors, not matrices and not \operatorname{Log}(\mathrm{R1\_pred}^{\top} \mathrm{R1\_true}). Use only valid_mask & diffuse_mask.
• Sum squared xyz errors and divide each example by 3\cdot \operatorname{max}(\mathrm{movable\_valid\_count},1); multiply by \mathrm{translation\_weight}\ge 0 and \mathrm{rotation\_weight}\ge 0 (both default 1). Bypass a zero-weight term before squaring.
• Then optionally cap the weighted translation loss at \mathrm{translation\_clip}>0 (default 5); translation_clip=None disables this cap. rotation_loss is uncapped. per_example is their sum, and loss is the unweighted mean over ALL B examples, including empty zeros.
• t=1 is safe because time is clipped; there is no denominator \epsilon. Invalid/fixed rotations are replaced by identity before logs. The only keys are t_clip, position_scale, translation_weight, rotation_weight, translation_clip.
• All masks are bool: valid_mask=True denotes a real residue/atom; diffuse_mask=True denotes a movable one. Their conjunction selects updates. B,N (or A) are positive; an all-false mask, not a zero-length tensor, represents an empty example.
• All supplied numbers, including padding, are finite and share dtype/device; floating tensors are float32 or float64. Rotations act on column vectors and map local coordinates to the global frame.
• For v=(x,y,z), \mathrm{hat}(v)=[[0,-z,y],[z,0,-x],[-y,x,0]]. \operatorname{Exp}(v)=I+A \mathrm{hat}(v)+B \mathrm{hat}(v)^{2}, where q=\mathrm{v.v}, A=\operatorname{sin}(\sqrt{q})/\sqrt{q}, B=(1-\operatorname{cos}(\sqrt{q}))/q.
• For q<10^{-4} use A=1-q/6+q^{2}/120-q^{3}/5040 and B=1/2-q/24+q^{2}/720-q^{3}/40320; at zero \operatorname{Exp}(0)=I.
• For \operatorname{Log}(R), set v=\mathrm{vee}(R-R^{\top})/2, q=\mathrm{v.v}, s=\sqrt{q}, c=(\mathrm{trace}(R)-1)/2, \theta=\operatorname{atan2}(s,c). Ordinarily \operatorname{Log}(R)=\theta\cdot v/s. If q<10^{-8} and c>0 use v\cdot (1+q/6+3q^{2}/40+5q^{3}/112).
• If q<10^{-6} and c<0, recover an axis from P=(R+R^{\top}-2\mathrm{cI})/(2(1-c)): choose the largest diagonal (\mathrm{ties} x, \mathrm{then} y, \mathrm{then} z), take its positive square root, and divide its row by that component; normalize the axis.
• Orient it to have nonnegative dot with v when q>10^{-20}; otherwise keep the chosen component positive. Return \theta\cdot \mathrm{axis}.
• For dormant branches clamp 1-c below by 10^{-6} and diagonal square-root arguments below by 10^{-12}; these floors do not alter the selected near-\pi axis. Evaluate every inactive square root/division safely before torch.where. R must be a proper \operatorname{SO}(3) matrix; no projection/repair is requested.
• At \pi axis sign is inherently ambiguous: \operatorname{exp}(\operatorname{log}(R)) is the geometric invariant. No global continuous logarithm exists at the \pi cut locus.

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 flow_matching_loss(
    x_t: torch.Tensor,  # (B, N, 3)
    x1_true: torch.Tensor,  # (B, N, 3)
    x1_pred: torch.Tensor,  # (B, N, 3)
    R_t: torch.Tensor,  # (B, N, 3, 3)
    R1_true: torch.Tensor,  # (B, N, 3, 3)
    R1_pred: torch.Tensor,  # (B, N, 3, 3)
    t: torch.Tensor,  # (B,)
    masks: dict,
    translation_weight: float = 1.0,
    translation_clip: float | None = 5.0,
    rotation_weight: float = 1.0,
    t_clip: float = 0.9,
    position_scale: float = 1.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    translation_loss: (B,).
    rotation_loss: (B,).
    per_example: (B,).
    loss: Scalar tensor, shape ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement flow_matching_loss")

Sources and further reading

• microsoft/protein-frame-flow · flow_module.py
• RosettaCommons/RFdiffusion2 · noisers.py