File size: 8,318 Bytes
8e04e6f | 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 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 | from abc import abstractmethod
from typing import Dict, Optional, Tuple, Union
import torch
from jaxtyping import Bool
from torch import Tensor
ABS_CLASS_ERR_MSG = "Method not implemented in abstract class"
class BaseFlowMatcher:
"""
Base class for flow matchers. All flow matching methods in the `base_flow_matching`
directory should inherit from this class (and implement most -see below- methods).
Some details:
- Attributes:
- `zero_com` specifies whether it should center in that modality
(see `mask_n_zero_com` method below).
- `guidance_enabled` specifies whether the underlying data modality
admits guidance or not.
- Methods:
- `nn_out_to_clean`: [x_t, t, nn_out] -> x_1_pred (+ some stuff).
- `interpolate`: [x_0, x_1, t] -> x_t
- `compute_loss` In principle, we use the v-loss || v_pred - v_true ||^2
for all modalities. We can add a config to change this.
- `mask_n_zero_com`. Should be used carefully, since centering multiple modalities
independently yields incorrect behavior. For now, centering can only be used for a
single modality. For some modalities centering may not make sense, for these the
`mask_n_zero_com` method just masks and does not center.
- `sample_noise`. Samples the reference distribution.
- `simulation_step`. Takes a simulation step.
- `extract_clean_sample_from_batch`. Interfaces with our dataloader's batch.
"""
def __init__(self, guidance_enabled: bool, dim: int):
self.guidance_enabled = guidance_enabled
self.dim = dim
@abstractmethod
def mask_n_zero_com(
self, x: torch.Tensor, mask: torch.Tensor = None
) -> torch.Tensor:
"""
Masks sample and fixes center of mass to zero (if applicable).
Args:
x: sample to mask and zero com
mask: mask, shape [*, n]
Returns:
sample x masked and with zero com
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def sample_noise(
self,
n: int,
device: torch.device,
shape: Tuple = tuple(),
mask: Optional[Bool[Tensor, "* n"]] = None,
) -> torch.Tensor:
"""
Samples reference distribution (possibly centered).
Args:
n: number of residues in a single sample (i.e. protein length), int
mask_n_zero_com
shape: tuple (if empty then single sample)
mask (optional): Binary mask of shape [*, n]
Returns:
Samples from refenrece with shape [*shape, n, ...]
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def interpolate(
self,
x_0: torch.Tensor,
x_1: torch.Tensor,
t: torch.Tensor,
mask: torch.Tensor = None,
) -> torch.Tensor:
"""
Interpolates between x_0 (base) and x_1 (data) using t.
Args:
x_0: Samples from reference with batch shape *
x_1: Sampels from target with batch shape *
t: Interpolation times, shape [*]
mask (optional): Binary mask, shape [*, n]
Returns:
x_t: Interpolated samples, same shape as x_0 and x_1
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def extract_clean_sample_from_batch(self, batch: Dict) -> torch.Tensor:
"""
Extracts clean sample from the batch coming from the dataloader.
Args:
batch: batch from dataloader.
Returns:
Clean sample x_1, a tensor with batch shape [*]
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def nn_out_add_clean_sample_prediction(
self,
x_t: torch.Tensor,
t: torch.Tensor,
mask: torch.Tensor,
nn_out: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
"""
Computes predicted clean sample given nn output, and adds it to the nn output
(if not there due to parameterization used).
Args:
x_0: noise sample, shape [*, n, 3]
x_1: clean sample, shape [*, n, 3]
x_t: interpolated sample, shape [*, n, 3]
t: time sampled, shape [*]
nn_out: output of neural network for this flow matcher, Dict[str, torch.Tensor]
Returns:
The nn_out dictionary updated with clean sample prediction (key "x_1").
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def nn_out_add_simulation_tensor(
self,
x_t: torch.Tensor,
t: torch.Tensor,
mask: torch.Tensor,
nn_out: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
"""
Computes simulation tensor (v or score, depending on base flow matcher) given nn output, and
adds it to the nn output if not there.
Args:
x_0: noise sample, shape [*, n, 3]
x_1: clean sample, shape [*, n, 3]
x_t: interpolated sample, shape [*, n, 3]
t: time sampled, shape [*]
nn_out: output of neural network, Dict[str, torch.Tensor]
Returns:
The nn_out dictionary updated with simulation tensor (key "v" or "score").
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def compute_fm_loss(
self,
x_0: torch.Tensor,
x_1: torch.Tensor,
x_t: torch.Tensor,
mask: torch.Tensor,
t: torch.Tensor,
x_1_pred: torch.Tensor,
) -> torch.Tensor:
"""
Computes flow matching loss per element in the batch.
Args:
x_0: noise sample, shape [b, ...]
x_1: clean sample, shape [b, ...]
x_t: interpolated sample, shape [b, ...]
mask (optional): Binary mask, shape [*, n]
t: time sampled, shape [b]
x_1_pred: predicted clean sample, shape [b, ...]
Returns:
Loss per batch element, shape [b]
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def nn_out_add_guided_simulation_tensor(
self,
nn_out: Dict[str, torch.Tensor],
nn_out_ag: Union[Dict[str, torch.Tensor], None],
nn_out_ucond: Union[Dict[str, torch.Tensor], None],
guidance_w: float,
ag_ratio: float,
) -> Dict[str, torch.Tensor]:
"""
Guidance logic, assumes the nn_out stuff contain the corresponding
simulation tensors. See `R3NFlowMatcher` for an example.
Args:
nn_out: output of neural network from full model, Dict[str, torch.Tensor]
nn_out_ag: output of neural network from autoguidance model, Dict[str, torch.Tensor] or None
nn_out_ucond: output of neural network from unconditional model, Dict[str, torch.Tensor] or None
guidance_w: guidance weight, float
ag_ratio: autoguidance ratio, float
Returns:
The nn_out dictionary updated with guided simulation tensor ("v" or "score", or whatever is needed).
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
@abstractmethod
def simulation_step(
self,
x_t: torch.Tensor,
nn_out: Dict[str, torch.Tensor],
t: torch.Tensor,
dt: float,
gt: float,
mask: torch.Tensor,
simulation_step_params: Dict,
):
"""
Single integration step of ODE \dot{x_t} = v(x_t, t) using Euler integration scheme.
Args:
x_t: Current value, batch shape *
nn_out: Dictionary with all available predictions, should include "v" and possibly guided "v_guided".
May include "x_1", etc as well. All batch shape *
t: Current time, shape [*]
dt: Step-size, float
gt: Noise injection, float
mask: Binary mask of shape [*, n]
simulation_step_params: parameters for the simulation step, depends on
data mode.
Returns:
Updated x_t after integration step, same shape as input
"""
raise NotImplementedError(ABS_CLASS_ERR_MSG)
|