File size: 4,662 Bytes
c65e212 4c8d0e0 c65e212 4c8d0e0 c65e212 4c8d0e0 c65e212 4c8d0e0 c65e212 4c8d0e0 c65e212 4c8d0e0 | 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 | """Precompute pair and atom biases used by the Boltz2 diffusion stack."""
from __future__ import annotations
import torch
from torch import nn
from .vb_modules_encodersv2 import AtomEncoder, PairwiseConditioning
def _bias_projections(
depth: int,
input_dim: int,
num_heads: int,
) -> nn.ModuleList:
"""Build one normalized, bias-free projection per transformer block."""
return nn.ModuleList(
[
nn.Sequential(
nn.LayerNorm(input_dim),
nn.Linear(input_dim, num_heads, bias=False),
)
for _ in range(depth)
]
)
def _concatenate_biases(
projections: nn.ModuleList,
pair_features: torch.Tensor,
) -> torch.Tensor:
# pair_features: (..., d_pair); each projection: (..., h).
return torch.cat(
[projection(pair_features) for projection in projections], dim=-1
) # (..., depth * h)
class DiffusionConditioning(nn.Module):
"""Prepare conditioned atom features and per-layer attention biases."""
def __init__(
self,
token_s: int,
token_z: int,
atom_s: int,
atom_z: int,
atoms_per_window_queries: int = 32,
atoms_per_window_keys: int = 128,
atom_encoder_depth: int = 3,
atom_encoder_heads: int = 4,
token_transformer_depth: int = 24,
token_transformer_heads: int = 8,
atom_decoder_depth: int = 3,
atom_decoder_heads: int = 4,
atom_feature_dim: int = 128,
conditioning_transition_layers: int = 2,
use_no_atom_char: bool = False,
use_atom_backbone_feat: bool = False,
use_residue_feats_atoms: bool = False,
) -> None:
super().__init__()
self.pairwise_conditioner = PairwiseConditioning(
token_z=token_z,
dim_token_rel_pos_feats=token_z,
num_transitions=conditioning_transition_layers,
)
self.atom_encoder = AtomEncoder(
atom_s=atom_s,
atom_z=atom_z,
token_s=token_s,
token_z=token_z,
atoms_per_window_queries=atoms_per_window_queries,
atoms_per_window_keys=atoms_per_window_keys,
atom_feature_dim=atom_feature_dim,
structure_prediction=True,
use_no_atom_char=use_no_atom_char,
use_atom_backbone_feat=use_atom_backbone_feat,
use_residue_feats_atoms=use_residue_feats_atoms,
)
self.atom_enc_proj_z = _bias_projections(
atom_encoder_depth,
atom_z,
atom_encoder_heads,
)
self.atom_dec_proj_z = _bias_projections(
atom_decoder_depth,
atom_z,
atom_decoder_heads,
)
self.token_trans_proj_z = _bias_projections(
token_transformer_depth,
token_z,
token_transformer_heads,
)
def forward(
self,
s_trunk: torch.Tensor,
z_trunk: torch.Tensor,
relative_position_encoding: torch.Tensor,
feats: dict[str, torch.Tensor],
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Return conditioned atom tensors and concatenated layer biases.
``S`` has shape ``(b, n, d_s)`` and each ``Z`` tensor has shape
``(b, n, n, d_z)``. Biases are concatenated in transformer-block
order so downstream code can select them by layer.
"""
# b is batch size, t token count, a atom count, and k the atom-window count.
z_conditioned = self.pairwise_conditioner(
z_trunk,
relative_position_encoding,
) # (b, t, t, d_z)
q, c, p, to_keys = self.atom_encoder(
feats=feats,
s_trunk=s_trunk,
z=z_conditioned,
) # q/c: (b, a, d_a); p: (b, k, w, h_k, d_p); to_keys: callable
atom_encoder_bias = _concatenate_biases(
self.atom_enc_proj_z, p
) # (b, k, w, h_k, depth_enc * heads_enc)
atom_decoder_bias = _concatenate_biases(
self.atom_dec_proj_z, p
) # (b, k, w, h_k, depth_dec * heads_dec)
token_transformer_bias = _concatenate_biases(
self.token_trans_proj_z,
z_conditioned,
) # (b, t, t, depth_token * heads_token)
return (
q,
c,
to_keys,
atom_encoder_bias,
atom_decoder_bias,
token_transformer_bias,
) # tensor shapes are traced above; to_keys is the atom-key gatherer
|