esm 3.0.0__py3-none-any.whl

esm/layers/attention.py

@@ -0,0 +1,70 @@
+import functools
+
+import einops
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from esm.layers.rotary import RotaryEmbedding
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        bias: bool = False,
+        qk_layernorm: bool = True,
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+
+        self.d_head = self.d_model // self.n_heads
+        self.layernorm_qkv = nn.Sequential(
+            nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=bias)
+        )
+        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
+
+        if qk_layernorm:
+            self.q_ln = nn.LayerNorm(d_model, bias=bias)
+            self.k_ln = nn.LayerNorm(d_model, bias=bias)
+        else:
+            self.q_ln = nn.Identity()
+            self.k_ln = nn.Identity()
+
+        self.rotary = RotaryEmbedding(d_model // n_heads)
+
+    def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor):
+        q = q.unflatten(-1, (self.n_heads, self.d_head))
+        k = k.unflatten(-1, (self.n_heads, self.d_head))
+        q, k = self.rotary(q, k)
+        q = q.flatten(-2, -1)
+        k = k.flatten(-2, -1)
+        return q, k
+
+    def forward(self, x, seq_id):
+        qkv_BLD3 = self.layernorm_qkv(x)
+        query_BLD, key_BLD, value_BLD = torch.chunk(qkv_BLD3, 3, dim=-1)
+        query_BLD, key_BLD = self.q_ln(query_BLD), self.k_ln(key_BLD)
+        query_BLD, key_BLD = self._apply_rotary(query_BLD, key_BLD)
+
+        n_heads = self.n_heads
+        reshaper = functools.partial(
+            einops.rearrange, pattern="b s (h d) -> b h s d", h=n_heads
+        )
+
+        query_BHLD, key_BHLD, value_BHLD = map(
+            reshaper, (query_BLD, key_BLD, value_BLD)
+        )
+
+        # Where True, enable participation in attention.
+        mask_BLL = seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)
+        mask_BHLL = mask_BLL.unsqueeze(1)
+
+        context_BHLD = F.scaled_dot_product_attention(
+            query_BHLD, key_BHLD, value_BHLD, mask_BHLL
+        )
+        context_BLD = einops.rearrange(context_BHLD, "b h s d -> b s (h d)")
+        return self.out_proj(context_BLD)

esm/layers/blocks.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from esm.layers.attention import MultiHeadAttention
+from esm.layers.geom_attention import (
+    GeometricReasoningOriginalImpl,
+)
+from esm.utils.structure.affine3d import Affine3D
+
+
+def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
+    # set hidden dimesion to nearest multiple of 256 after expansion ratio
+    return int(((expansion_ratio * d_model) + 255) // 256 * 256)
+
+
+class SwiGLU(nn.Module):
+    """
+    SwiGLU activation function as an nn.Module, allowing it to be used within nn.Sequential.
+    This module splits the input tensor along the last dimension and applies the SiLU (Swish)
+    activation function to the first half, then multiplies it by the second half.
+    """
+
+    def __init__(self):
+        super(SwiGLU, self).__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+def swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool):
+    return nn.Sequential(
+        nn.LayerNorm(d_model),
+        nn.Linear(
+            d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=bias
+        ),
+        SwiGLU(),
+        nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=bias),
+    )
+
+
+def gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool):
+    hidden_dim = int(expansion_ratio * d_model)
+    return nn.Sequential(
+        nn.LayerNorm(d_model),
+        nn.Linear(d_model, hidden_dim, bias=bias),
+        nn.GELU(),
+        nn.Linear(hidden_dim, d_model, bias=bias),
+    )
+
+
+class UnifiedTransformerBlock(nn.Module):
+    """
+    A unified transformer block that can optionally incorporate geometric attention.
+
+    This class defines a transformer block that can be configured to use geometric attention
+    alongside the standard multi-head attention mechanism. It is designed to be a flexible
+    component of transformer-based models, allowing for the integration of geometric reasoning.
+
+    Parameters
+    ----------
+    d_model : int
+        The dimensionality of the input and output features of the transformer block.
+    n_heads : int
+        The number of attention heads in the multi-head attention mechanism.
+    n_layers : int
+        The number of layers in the transformer block.
+    use_geom_attn : bool, optional
+        Whether to use geometric attention in addition to the standard multi-head attention. Defaults to False.
+    v_heads : int, optional
+        The number of heads to use for the geometric attention mechanism, if enabled. Must be specified if `use_geom_attn` is True.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        use_geom_attn: bool = False,
+        use_plain_attn: bool = True,
+        v_heads: int | None = None,
+        bias: bool = False,
+        expansion_ratio: float = 4.0,
+        residue_scaling_factor: float = 1,
+        mask_and_zero_frameless: bool = False,
+        qk_layernorm: bool = True,
+        ffn_type: str = "swiglu",  # swiglu | gelu
+    ):
+        super().__init__()
+        self.use_plain_attn = use_plain_attn
+        if self.use_plain_attn:
+            self.attn = MultiHeadAttention(
+                d_model, n_heads, bias, qk_layernorm=qk_layernorm
+            )
+        self.use_geom_attn = use_geom_attn
+        if self.use_geom_attn:
+            if v_heads is None:
+                raise ValueError("v_heads must be specified when use_geom_attn is True")
+            self.geom_attn = GeometricReasoningOriginalImpl(
+                c_s=d_model,
+                v_heads=v_heads,
+                bias=bias,
+                mask_and_zero_frameless=mask_and_zero_frameless,
+            )
+        if ffn_type == "swiglu":
+            self.ffn = swiglu_ln_ffn(d_model, expansion_ratio, bias)
+        elif ffn_type == "gelu":
+            self.ffn = gelu_ln_ffn(d_model, expansion_ratio, bias)
+        else:
+            raise ValueError(f"Unknown ffn_type: {ffn_type}")
+        self.scaling_factor = residue_scaling_factor
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sequence_id: torch.Tensor,
+        frames: Affine3D,
+        frames_mask: torch.Tensor,
+        chain_id: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass for the UnifiedTransformerBlock.
+
+        Parameters
+        ----------
+        x : torch.Tensor[float]
+            Input tensor to the transformer block, typically the output from the previous layer.
+        sequence_id : torch.Tensor[int]
+            Tensor containing sequence IDs for each element in the batch, used for attention masking.
+        frames : Affine3D
+            Affine3D containing geometric frame information for geometric attention.
+        frames_mask : torch.Tensor[bool]
+            Boolean mask tensor indicating valid frames for geometric attention.
+        chain_id : torch.Tensor[int]
+            Tensor containing chain IDs for each element, used for attention masking in geometric attention.
+
+        Returns
+        -------
+        torch.Tensor[float]
+            The output tensor after applying the transformer block operations.
+        """
+        if self.use_plain_attn:
+            r1 = self.attn(x, sequence_id)
+            x = x + r1 / self.scaling_factor
+
+        if self.use_geom_attn:
+            r2 = self.geom_attn(x, frames, frames_mask, sequence_id, chain_id)
+            x = x + r2 / self.scaling_factor
+
+        r3 = self.ffn(x) / self.scaling_factor
+        x = x + r3
+
+        return x

esm/layers/codebook.py

@@ -0,0 +1,88 @@
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class EMACodebook(nn.Module):
+    def __init__(
+        self,
+        n_codes,
+        embedding_dim,
+        no_random_restart=True,
+        restart_thres=1.0,
+        ema_decay=0.99,
+    ):
+        super().__init__()
+        self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
+        self.register_buffer("N", torch.zeros(n_codes))
+        self.register_buffer("z_avg", self.embeddings.data.clone())
+
+        self.n_codes = n_codes
+        self.embedding_dim = embedding_dim
+        self._need_init = True
+        self.no_random_restart = no_random_restart
+        self.restart_thres = restart_thres
+        self.freeze_codebook = False
+        self.ema_decay = ema_decay
+
+    def reset_parameters(self):
+        # For meta init
+        pass
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.n_codes:
+            n_repeats = (self.n_codes + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + torch.randn_like(x) * std
+        return x
+
+    def _init_embeddings(self, z):
+        # z: [b, t, c]
+        self._need_init = False
+        flat_inputs = z.view(-1, self.embedding_dim)
+        y = self._tile(flat_inputs)
+
+        y.shape[0]
+        _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
+        if dist.is_initialized():
+            dist.broadcast(_k_rand, 0)
+        self.embeddings.data.copy_(_k_rand)
+        self.z_avg.data.copy_(_k_rand)
+        self.N.data.copy_(torch.ones(self.n_codes))
+
+    def forward(self, z):
+        # z: [b, t, c]
+        if self._need_init and self.training and not self.freeze_codebook:
+            self._init_embeddings(z)
+        # z is of shape [batch_size, sequence length, channels]
+        flat_inputs = z.view(-1, self.embedding_dim)
+        distances = (
+            (flat_inputs**2).sum(dim=1, keepdim=True)
+            - 2 * flat_inputs @ self.embeddings.t()
+            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
+        )  # [bt, c]
+
+        encoding_indices = torch.argmin(distances, dim=1)
+        encoding_indices = encoding_indices.view(*z.shape[:2])  # [b, t, ncode]
+
+        embeddings = F.embedding(encoding_indices, self.embeddings)  # [b, t, c]
+
+        commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
+
+        # EMA codebook update
+        if self.training and not self.freeze_codebook:
+            assert False, "Not implemented"
+        embeddings_st = (embeddings - z).detach() + z
+
+        return embeddings_st, encoding_indices, commitment_loss
+
+    def dictionary_lookup(self, encodings):
+        embeddings = F.embedding(encodings, self.embeddings)
+        return embeddings
+
+    def soft_codebook_lookup(self, weights: torch.Tensor) -> torch.Tensor:
+        return weights @ self.embeddings

esm/layers/ffn.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+# NOT CURRENTLY USED
+
+
+class SwiGLU(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return hidden
+
+
+class FFN(nn.Module):
+    def __init__(self, in_proj, activation, out_proj) -> None:
+        super().__init__()
+        self.in_proj = in_proj
+        self.activation = activation
+        self.out_proj = out_proj
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.in_proj(x)
+        x = self.activation(x)
+        x = self.out_proj(x)
+        return x

esm/layers/geom_attention.py

@@ -0,0 +1,151 @@
+from math import sqrt
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+
+class GeometricReasoningOriginalImpl(nn.Module):
+    def __init__(
+        self,
+        c_s: int,
+        v_heads: int,
+        num_vector_messages: int = 1,
+        mask_and_zero_frameless: bool = True,
+        divide_residual_by_depth: bool = False,
+        bias: bool = False,
+    ):
+        """Approximate implementation:
+
+        ATTN(A, v) := (softmax_j A_ij) v_j
+        make_rot_vectors(x) := R(i->g) Linear(x).reshape(..., 3)
+        make_vectors(x) := T(i->g) Linear(x).reshape(..., 3)
+
+        v <- make_rot_vectors(x)
+        q_dir, k_dir <- make_rot_vectors(x)
+        q_dist, k_dist <- make_vectors(x)
+
+        A_ij       <- dot(q_dir_i, k_dir_j) -||q_dist_i - k_dist_j||^2
+        x          <- x + Linear(T(g->i) ATTN(A, v))
+        """
+        super().__init__()
+        self.c_s = c_s
+        self.v_heads = v_heads
+        self.num_vector_messages = num_vector_messages
+        self.mask_and_zero_frameless = mask_and_zero_frameless
+
+        self.s_norm = nn.LayerNorm(c_s, bias=bias)
+        dim_proj = (
+            4 * self.v_heads * 3 + self.v_heads * 3 * self.num_vector_messages
+        )  # 2 x (q, k) * number of heads * (x, y, z) + number of heads * number of vector messages * (x, y, z)
+        self.proj = nn.Linear(c_s, dim_proj, bias=bias)
+        channels_out = self.v_heads * 3 * self.num_vector_messages
+        self.out_proj = nn.Linear(channels_out, c_s, bias=bias)
+
+        # The basic idea is for some attention heads to pay more or less attention to rotation versus distance,
+        # as well as to control the sharpness of the softmax (i.e., should this head only attend to those residues
+        # very nearby or should there be shallower dropoff in attention weight?)
+        self.distance_scale_per_head = nn.Parameter(torch.zeros((self.v_heads)))
+        self.rotation_scale_per_head = nn.Parameter(torch.zeros((self.v_heads)))
+
+    def forward(self, s, affine, affine_mask, sequence_id, chain_id):
+        attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)
+        attn_bias = attn_bias.unsqueeze(1).float()
+        attn_bias = attn_bias.masked_fill(
+            ~affine_mask[:, None, None, :], torch.finfo(attn_bias.dtype).min
+        )
+        chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2)
+        attn_bias = attn_bias.masked_fill(
+            chain_id_mask.unsqueeze(1), torch.finfo(s.dtype).min
+        )
+
+        ns = self.s_norm(s)
+        vec_rot, vec_dist = self.proj(ns).split(
+            [
+                self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages,
+                self.v_heads * 2 * 3,
+            ],
+            dim=-1,
+        )
+
+        # Rotate the queries and keys for the rotation term. We also rotate the values.
+        # NOTE(zeming, thayes): Values are only rotated, not translated. We may wish to change
+        # this in the future.
+        query_rot, key_rot, value = (
+            affine.rot[..., None]
+            .apply(rearrange(vec_rot, "... (h c) -> ... h c", c=3))
+            .split(
+                [
+                    self.v_heads,
+                    self.v_heads,
+                    self.v_heads * self.num_vector_messages,
+                ],
+                dim=-2,
+            )
+        )
+
+        # Rotate and translate the queries and keys for the distance term
+        # NOTE(thayes): a simple speedup would be to apply all rotations together, then
+        # separately apply the translations.
+        query_dist, key_dist = (
+            affine[..., None]
+            .apply(rearrange(vec_dist, "... (h c) -> ... h c", c=3))
+            .chunk(2, dim=-2)
+        )
+
+        query_dist = rearrange(query_dist, "b s h d -> b h s 1 d")
+        key_dist = rearrange(key_dist, "b s h d -> b h 1 s d")
+        query_rot = rearrange(query_rot, "b s h d -> b h s d")
+        key_rot = rearrange(key_rot, "b s h d -> b h d s")
+        value = rearrange(
+            value, "b s (h m) d -> b h s (m d)", m=self.num_vector_messages
+        )
+
+        distance_term = (query_dist - key_dist).norm(dim=-1) / sqrt(3)
+        rotation_term = query_rot.matmul(key_rot) / sqrt(3)
+        distance_term_weight = rearrange(
+            F.softplus(self.distance_scale_per_head), "h -> h 1 1"
+        )
+        rotation_term_weight = rearrange(
+            F.softplus(self.rotation_scale_per_head), "h -> h 1 1"
+        )
+
+        attn_weight = (
+            rotation_term * rotation_term_weight - distance_term * distance_term_weight
+        )
+
+        if attn_bias is not None:
+            # we can re-use the attention bias from the transformer layers
+            # NOTE(thayes): This attention bias is expected to handle two things:
+            # 1. Masking attention on padding tokens
+            # 2. Masking cross sequence attention in the case of bin packing
+            s_q = attn_weight.size(2)
+            s_k = attn_weight.size(3)
+            _s_q = max(0, attn_bias.size(2) - s_q)
+            _s_k = max(0, attn_bias.size(3) - s_k)
+            attn_bias = attn_bias[:, :, _s_q:, _s_k:]
+            attn_weight = attn_weight + attn_bias
+
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+
+        attn_out = attn_weight.matmul(value)
+
+        attn_out = (
+            affine.rot[..., None]
+            .invert()
+            .apply(
+                rearrange(
+                    attn_out, "b h s (m d) -> b s (h m) d", m=self.num_vector_messages
+                )
+            )
+        )
+
+        attn_out = rearrange(
+            attn_out, "b s (h m) d -> b s (h m d)", m=self.num_vector_messages
+        )
+        if self.mask_and_zero_frameless:
+            attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0)
+        s = self.out_proj(attn_out)
+
+        return s

esm/layers/regression_head.py

@@ -0,0 +1,24 @@
+import torch.nn as nn
+
+
+def RegressionHead(
+    d_model: int,
+    output_dim: int,
+    hidden_dim: int | None = None,
+) -> nn.Module:
+    """Single-hidden layer MLP for supervised output.
+
+    Args:
+        d_model: input dimension
+        output_dim: dimensionality of the output.
+        hidden_dim: optional dimension of hidden layer, defaults to d_model.
+    Returns:
+        output MLP module.
+    """
+    hidden_dim = hidden_dim if hidden_dim is not None else d_model
+    return nn.Sequential(
+        nn.Linear(d_model, hidden_dim),
+        nn.GELU(),
+        nn.LayerNorm(hidden_dim),
+        nn.Linear(hidden_dim, output_dim),
+    )