diff --git a/mtenn/combination.py b/mtenn/combination.py index e2cb8c2..3774cf8 100644 --- a/mtenn/combination.py +++ b/mtenn/combination.py @@ -542,3 +542,236 @@ def backward(ctx, comb_grad, pose_grads): # don't get gradients so we just return None return_vals = [None] * 5 + [final_grads[n] for n in ctx.param_names] return tuple(return_vals) + + +class BoltzmannCombination(Combination): + """ + Combine a list of :math:`\mathrm{\Delta G}` predictions according to their + Boltzmann weight. Treat energy in implicit kT units. See the docs for + :py:class:`BoltzmannCombinationFunc ` + for more details. + """ + + def __repr__(self): + return "BoltzmannCombination()" + + def __str__(self): + return repr(self) + + def forward(self, pred_list, grad_dict, param_names, *model_params): + return BoltzmannCombinationFunc.apply( + pred_list, grad_dict, param_names, *model_params + ) + + +class BoltzmannCombinationFunc(torch.autograd.Function): + """ + Custom autograd function that will handle the gradient math for us for combining + :math:`\mathrm{\Delta G}` predictions by Boltzmann weighting. + + .. math:: + + \Delta G &= \sum_{n=1}^{N} w_n \Delta G_n + + w_n &= \mathrm{exp} \\left[ + -\Delta G_n - \mathrm{ln} \\sum_{i=1}^N \\mathrm{exp} (-\Delta G_i ) \\right] + + See :ref:`boltzmann-comb-imp` for more details on the math. + """ + + @staticmethod + def forward(pred_list, grad_dict, param_names, *model_params): + """ + Find the max/min of all input :math:`\mathrm{\Delta G}` predictions. + + Parameters + ---------- + pred_list: List[torch.Tensor] + List of :math:`\mathrm{\Delta G}` predictions to be combined, shape of + ``(n_predictions,)`` + grad_dict: dict[str, List[torch.Tensor]] + Dict mapping from parameter name to list of gradients. Should contain + ``n_model_parameters`` entries, with each entry mapping to a list of + ``n_predictions`` tensors. Each of these tensors is a ``detach`` ed gradient + so the shape of each tensor will depend on the model parameter it + corresponds to, but the shapes of each tensor in any given entry should be + identical + param_names: List[str] + List of parameter names. Should contain ``n_model_parameters`` entries, + corresponding 1:1 with the keys in ``grad_dict`` + model_params: List[torch.Tensor] + Actual parameters that we'll return the gradients for. Each param + should be passed directly for the backward pass to + work right. These tensors should correspond 1:1 with and should be in the + same order as the entries in ``param_names`` (ie the ``i`` th entry in + ``param_names`` should be the name of the ``i`` th model parameter in + ``model_params``) + + Returns + ------- + torch.Tensor + Scalar-value tensor giving the Boltzmann mean of the input + :math:`\mathrm{\Delta G}` predictions + torch.Tensor + Tensor of shape ``(n_predictions,)`` giving the input per-pose predictions + """ + # Save for later so we don't have to keep redoing this + adj_preds = -torch.stack(pred_list).flatten().detach() + + # First calculate the normalization factor + Q = torch.logsumexp(adj_preds, dim=0) + + # Calculate w + w = (adj_preds - Q).exp() + + # Calculate final pred + final_pred = torch.dot(w, -adj_preds) + + return final_pred, -adj_preds + + @staticmethod + def setup_context(ctx, inputs, output): + """ + Store data for backward pass. + + Parameters + ---------- + ctx + Pytorch context manager + inputs : List + List containing all the parameters that will get passed to ``forward`` + output : torch.Tensor + Values returned from ``forward`` + """ + # Unpack the inputs + (pred_list, grad_dict, param_names, *model_params) = inputs + + # Break the grad dict up into lists of keys and corresponding lists of gradients + grad_dict_keys, grad_dict_tensors = Combination.split_grad_dict(grad_dict) + + # Save non-Tensors for backward + ctx.grad_dict_keys = grad_dict_keys + ctx.param_names = param_names + + # Save Tensors for backward + # Saving: + # * Predictions (1 tensor of shape (n_predictions,)) + # * Grad tensors (N params * M poses tensors, where all gradients corresponding + # to a given model parameter are adjacent, ie first M tensors are the + # per-pose gradients for the first model parameter, etc) + # * Model param tensors (N params tensors) + ctx.save_for_backward( + torch.stack(pred_list).flatten(), *grad_dict_tensors, *model_params + ) + + @staticmethod + def backward(ctx, comb_grad, pose_grads): + """ + Compute and return gradients for each parameter. + + Parameters + ---------- + ctx + Pytorch context manager + comb_grad : torch.Tensor + Scalar-value tensor giving the + :math:`\\frac{\\partial L}{\\partial \\Delta \\text{G}}` term from + :eq:`comb-grad` + pose_grads : torch.Tensor + Tensor of shape ``(n_predictions,)``, giving the + :math:`\\frac{\\partial L}{\\partial \\Delta \\text{G}_i}` terms from + :eq:`pose-grad` + """ + print("grads", comb_grad, pose_grads, flush=True) + + # Unpack saved tensors + preds, *other_tensors = ctx.saved_tensors + + # First section of these tensors are the flattened lists of gradients from each + # individual pose or each model parameter + grad_dict_tensors = other_tensors[: len(ctx.grad_dict_keys)] + + # Reconstruct dict mapping from model parameter name to list of gradient tensors + # The ith entry in each list gives the gradient of the ith pose prediction wrt + # that model parameter + grad_dict = Combination.join_grad_dict(ctx.grad_dict_keys, grad_dict_tensors) + + # We use adj_preds here to store the adjusted per-pose prediction values. These + # values have been negated (if we are finding the min), and multiplied by our + # scale value, if given + # These values correspond to the values inside the exponential in eqn (5) (and + # subsequent equations) + adj_preds = -preds.detach() + + # Calculate our normalizing constant (eqn (6)) + Q = torch.logsumexp(adj_preds, dim=0) + + # Calc w values + boltzmann_weights = (adj_preds - Q).exp() + + # Calculate dQ/d_theta, should be a single tensor for each model parameter + # grad_list is list of gradients (one for each input pose -> prediction) for the + # given model_param + # Each gradient for the given model_param gets weighted by its pose weight (w_n) + dQ = { + model_param: -torch.stack( + [w_n * grad for w_n, grad in zip(boltzmann_weights, grad_list)], + axis=-1, + ).sum(axis=-1) + for model_param, grad_list in grad_dict.items() + } + + # Calculate dw/d_theta, will be a list of tensors for each model parameter + # As before, grad_list is list of gradients (one for each + # input pose -> prediction) for the given model_param + # Now, we need a list of gradients for each model_param, one for each pose (n) + dw = { + model_param: [ + w_n * (-grad - dQ[model_param]) + for w_n, grad in zip(boltzmann_weights, grad_list) + ] + for model_param, grad_list in grad_dict.items() + } + + # Calculate final gradients for each parameter + final_grads = {} + for model_param, grad_list in grad_dict.items(): + # Compute the gradient contributions from any combined prediction loss, + # according to eqns (1), (9) + # Stack the tensors for each input pose/prediction, then sum across them + # Multiply by the dL/ddG term to get the final dL/dtheta gradients for each + # model parameter + cur_final_grad = comb_grad * ( + torch.stack( + [ + w_grad * -pred + w_val * grad + for w_grad, pred, w_val, grad in zip( + dw[model_param], adj_preds, boltzmann_weights, grad_list + ) + ], + axis=-1, + ) + .detach() + .sum(axis=-1) + ) + + # Make sure lengths match up (should always be true but just in case) + if len(pose_grads) != len(grad_list): + raise RuntimeError("Mismatch in gradient lengths.") + + # Compute the gradient contributions from any per-pose prediction loss, + # according to eqn (2) + # Will be zero if there was no loss penalty directly applied on the + # individual pose predictions + for pose_grad, param_grad in zip(pose_grads, grad_list): + cur_final_grad += pose_grad * param_grad + + # Store total gradient for each parameter + final_grads[model_param] = cur_final_grad.clone() + + # Return gradients for each of the model parameters that were passed in. Also + # need to return values for the other values that were passed to forward + # (negate_preds, pred_scale, pred_list, grad_dict, param_names), but these + # don't get gradients so we just return None + return_vals = [None] * 3 + [final_grads[n] for n in ctx.param_names] + return tuple(return_vals) diff --git a/mtenn/config.py b/mtenn/config.py index cfafa8a..80c43ca 100644 --- a/mtenn/config.py +++ b/mtenn/config.py @@ -165,10 +165,13 @@ class CombinationConfig(StringEnum): * mean: :py:class:`MeanCombination ` * max: :py:class:`MaxCombination ` + + * boltzmann: :py:class:`BoltzmannCombination ` """ mean = "mean" max = "max" + boltzmann = "boltzmann" class ModelConfigBase(BaseModel, abc.ABC): @@ -365,6 +368,8 @@ def build(self) -> mtenn.model.Model: mtenn_combination = mtenn.combination.MaxCombination( negate_preds=self.max_comb_neg, pred_scale=self.max_comb_scale ) + case CombinationConfig.boltzmann: + mtenn_combination = mtenn.combination.BoltzmannCombination() case None: mtenn_combination = None diff --git a/mtenn/tests/test_combination.py b/mtenn/tests/test_combination.py index 222fe9b..f5015cf 100644 --- a/mtenn/tests/test_combination.py +++ b/mtenn/tests/test_combination.py @@ -3,7 +3,7 @@ import pytest import torch -from mtenn.combination import MeanCombination, MaxCombination +from mtenn.combination import BoltzmannCombination, MeanCombination, MaxCombination from mtenn.conversion_utils.schnet import SchNet @@ -25,13 +25,23 @@ def models_and_inputs(): for _ in range(5) ] target = torch.rand(1) - loss_func = torch.nn.MSELoss() + target_pose = torch.randint(0, len(inp_list), ()) + loss_func_pred = torch.nn.MSELoss() + loss_func_pose = torch.nn.CrossEntropyLoss() - return model_test, model_ref, inp_list, target, loss_func + return ( + model_test, + model_ref, + inp_list, + target, + target_pose, + loss_func_pred, + loss_func_pose, + ) def test_mean_combination(models_and_inputs): - model_test, model_ref, inp_list, target, loss_func = models_and_inputs + model_test, model_ref, inp_list, target, _, loss_func, _ = models_and_inputs # Ref calc pred_list = [model_ref(X)[0] for X in inp_list] @@ -60,7 +70,7 @@ def test_mean_combination(models_and_inputs): def test_max_combination(models_and_inputs): - model_test, model_ref, inp_list, target, loss_func = models_and_inputs + model_test, model_ref, inp_list, target, _, loss_func, _ = models_and_inputs # Ref calc pred_list = [model_ref(X)[0] for X in inp_list] @@ -89,3 +99,83 @@ def test_max_combination(models_and_inputs): for n, p in model_test.named_parameters() ] ) + + +def test_boltzmann_combination(models_and_inputs): + model_test, model_ref, inp_list, target, _, loss_func, _ = models_and_inputs + + # Ref calc + pred_list = [model_ref(X)[0] for X in inp_list] + adj_preds = -torch.stack(pred_list).flatten() + Q = torch.logsumexp(adj_preds, dim=0) + w = (adj_preds - Q).exp() + pred = torch.dot(w, -adj_preds) + + loss = loss_func(pred, target) + loss.backward() + + # Finish setting up GroupedModel + model_test = SchNet.get_model( + model_test, grouped=True, strategy="complex", combination=BoltzmannCombination() + ) + + # Test GroupedModel + pred, _ = model_test(inp_list) + loss = loss_func(pred, target) + loss.backward() + + # Compare + ref_param_dict = dict(model_ref.named_parameters()) + assert all( + [ + np.allclose(p.grad, ref_param_dict[n].grad, atol=5e-7) + for n, p in model_test.named_parameters() + ] + ) + + +def test_boltzmann_combination_multi_loss(models_and_inputs): + ( + model_test, + model_ref, + inp_list, + target, + target_pose, + loss_func_pred, + loss_func_pose, + ) = models_and_inputs + + # Ref calc + pred_list = [model_ref(X)[0] for X in inp_list] + adj_preds = -torch.stack(pred_list).flatten() + Q = torch.logsumexp(adj_preds, dim=0) + w = (adj_preds - Q).exp() + pred = torch.dot(w, -adj_preds) + + pred_loss = loss_func_pred(pred, target) + pose_loss = loss_func_pose(-adj_preds, target_pose) + loss = 0.25 * pose_loss + 0.75 * pred_loss + + loss.backward() + + # Finish setting up GroupedModel + model_test = SchNet.get_model( + model_test, grouped=True, strategy="complex", combination=BoltzmannCombination() + ) + + # Test GroupedModel + pred, preds_list = model_test(inp_list) + pred_loss = loss_func_pred(pred, target) + pose_loss = loss_func_pose(preds_list, target_pose) + loss = 0.25 * pose_loss + 0.75 * pred_loss + + loss.backward() + + # Compare + ref_param_dict = dict(model_ref.named_parameters()) + assert all( + [ + np.allclose(p.grad, ref_param_dict[n].grad, atol=5e-7) + for n, p in model_test.named_parameters() + ] + )