File size: 6,628 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from typing import Literal, Protocol
import torch
from torch_scatter import scatter
from ...diffusion.corruption.corruption import Corruption
from ...diffusion.corruption.sde_lib import maybe_expand
from ...diffusion.d3pm.d3pm import compute_kl_reverse_process
from ...diffusion.data.batched_data import BatchedData
from ...diffusion.discrete_time import to_discrete_time
from ...diffusion.model_target import ModelTarget
def compute_noise_given_sample_and_corruption(
x: torch.Tensor,
x_noisy: torch.Tensor,
corruption: Corruption,
t: torch.Tensor,
batch_idx: torch.LongTensor | None,
batch: BatchedData,
) -> torch.Tensor:
"""
Recover the (unit-Gaussian-distributed) raw noise that was used to corrupt a batch of samples.
We first obtain the mean and std of the noisy samples from the corruption via `t` and the clean batch.
Then we solve:
x_noisy = x_mean + noise * std w.r.t. `noise`:
noise = (x_noisy - x_mean) / std
"""
x_mean, std = corruption.marginal_prob(
x,
t=t,
batch_idx=batch_idx,
batch=batch,
)
return (x_noisy - x_mean) / std
class FieldLoss(Protocol):
"""Loss function for a single field. Because loss functions are defined different ways in different papers,
we pass loads of keyword arguments. Each loss function will only use a subset of these arguments.
"""
def __call__(
self,
*,
corruption: Corruption,
score_model_output: torch.Tensor,
t: torch.Tensor,
batch_idx: torch.LongTensor | None,
batch_size: int,
x: torch.Tensor,
noisy_x: torch.Tensor,
reduce: Literal["sum", "mean"],
batch: BatchedData,
) -> torch.Tensor:
"""Calculate loss per sample for a single field. Returns a loss tensor of shape (batch_size,)."""
pass
def denoising_score_matching(
*,
corruption: Corruption,
score_model_output: torch.Tensor,
t: torch.Tensor,
batch_idx: torch.LongTensor | None,
batch_size: int,
x: torch.Tensor,
noisy_x: torch.Tensor,
reduce: Literal["sum", "mean"],
batch: BatchedData,
model_target: ModelTarget,
node_is_unmasked: torch.LongTensor | None = None,
**_,
) -> torch.Tensor:
"""Mean square error in predicting raw noise, optionally reweighted."""
assert score_model_output.ndim >= 2
model_target = ModelTarget(model_target) # in case str was passed
losses = get_losses(
corruption=corruption,
score_model_output=score_model_output,
t=t,
batch_idx=batch_idx,
x=x,
noisy_x=noisy_x,
batch=batch,
model_target=model_target,
)
if node_is_unmasked is not None:
losses = node_is_unmasked.unsqueeze(-1) * losses # Apply masking.
original_reduce = reduce
reduce = "sum" # We sum first and handle the division by nodes_per_sample for the mean manually later.
loss_per_sample = aggregate_per_sample(losses, batch_idx, reduce=reduce, batch_size=batch_size)
if (node_is_unmasked is not None) and (original_reduce == "mean"):
nodes_per_sample = scatter(node_is_unmasked, batch_idx, dim=0, reduce="sum")
loss_per_sample /= nodes_per_sample
return loss_per_sample
def get_losses(
corruption: Corruption,
score_model_output: torch.Tensor,
t: torch.Tensor,
batch_idx: torch.LongTensor | None,
x: torch.Tensor,
noisy_x: torch.Tensor,
batch: BatchedData,
model_target: ModelTarget,
) -> torch.Tensor:
if model_target == ModelTarget.score_times_std:
raw_noise = compute_noise_given_sample_and_corruption(
x=x, x_noisy=noisy_x, corruption=corruption, t=t, batch_idx=batch_idx, batch=batch
)
target = -raw_noise
losses = (score_model_output - target).square()
else:
raise ValueError(f"Unknown model_target {model_target}")
return losses
def aggregate_per_sample(
loss_per_row: torch.Tensor,
batch_idx: torch.Tensor | None,
reduce: Literal["sum", "mean"],
batch_size: int,
):
"""
Aggregate (potentially) batched input tensor to get a scalar for each sample in the batch.
E.g., (num_atoms, d1, d2, ..., dn) -> (batch_size, d1, d2, ..., dn) -> (batch_size,),
where the first aggregation only happens when batch_idx is provided.
Args:
loss_per_row: shape (num_nodes, any_more_dims). May contain multiple nodes per sample.
batch_idx: shape (num_nodes,). Indicates which sample each row belongs to. If not provided,
then we assume the first dimension is the batch dimension.
reduce: determines how to aggregate over nodes within each sample. (Aggregation over samples
and within dims for one node is always mean.)
batch_size: number of samples in the batch.
Returns:
Scalar for each sample, shape (batch_size,).
"""
# Sum over all but 0th dimension.
loss_per_row = torch.mean(loss_per_row.reshape(loss_per_row.shape[0], -1), dim=1)
if batch_idx is None:
# First dimension is batch dimension. In this case 'reduce' is ignored.
loss_per_sample = loss_per_row
else:
# Aggregate over nodes within each sample.
loss_per_sample = scatter(
src=loss_per_row,
index=batch_idx,
dim_size=batch_size,
reduce=reduce,
)
return loss_per_sample
def d3pm_loss(
*,
corruption: Corruption,
score_model_output: torch.Tensor,
t: torch.Tensor,
batch_idx: torch.LongTensor | None,
batch_size: int,
x: torch.Tensor,
noisy_x: torch.Tensor,
reduce: Literal["sum", "mean"],
d3pm_hybrid_lambda: float = 0.0,
**_,
) -> torch.Tensor:
assert hasattr(corruption, "N") # mypy
assert hasattr(corruption, "_to_zero_based") # mypy
assert hasattr(corruption, "d3pm") # mypy
t = maybe_expand(to_discrete_time(t, N=corruption.N, T=corruption.T), batch_idx)
metrics_dict = compute_kl_reverse_process(
corruption._to_zero_based(x.long()),
t,
diffusion=corruption.d3pm,
log_space=True,
denoise_fn=lambda targets, timestep: score_model_output,
hybrid_lambda=d3pm_hybrid_lambda,
x_t_plus_1=corruption._to_zero_based(noisy_x.long()),
)
loss = metrics_dict.pop("loss")
loss_per_structure = aggregate_per_sample(
loss, batch_idx=batch_idx, reduce=reduce, batch_size=batch_size
)
return loss_per_structure
|