Transformer

salt.models.transformer_v2.Attention

Bases: torch.nn.Module, abc.ABC

Multihead attention module.

Parameters:

Name	Type	Description	Default
embed_dim	int	Dimension of the input.	required
num_heads	int	Number of attention heads.	required
attn_type	str	Type of backend kernel to use.	'torch-meff'
n_kv_heads	int | None	Number of heads for the keys and values. If None, defaults to num_heads.	None
window_size	int | None	Window size for flash attention kernel. If None, defaults to global attention.	None
dropout	float	Dropout rate.	0.0
bias	bool	Whether to include bias terms.	True
add_zero_attn	bool	Whether to add a dummy token to attend to. This avoids nan when all tokens are padded.	True

Source code in salt/models/transformer_v2.py

def __init__(
    self,
    embed_dim: int,
    num_heads: int,
    attn_type: str = "torch-meff",
    n_kv_heads: int | None = None,
    window_size: int | None = None,
    dropout: float = 0.0,
    bias: bool = True,
    add_zero_attn: bool = True,
):
    """Multihead attention module.

    Parameters
    ----------
    embed_dim : int
        Dimension of the input.
    num_heads : int
        Number of attention heads.
    attn_type : str, optional
        Type of backend kernel to use.
    n_kv_heads : int | None, optional
        Number of heads for the keys and values. If None, defaults to num_heads.
    window_size : int | None, optional
        Window size for flash attention kernel. If None, defaults to global attention.
    dropout : float, optional
        Dropout rate.
    bias : bool, optional
        Whether to include bias terms.
    add_zero_attn : bool, optional
        Whether to add a dummy token to attend to. This avoids nan when all tokens are padded.
    """
    super().__init__()

    self.embed_dim = embed_dim
    self.num_heads = num_heads
    self.head_dim = embed_dim // num_heads

    self.n_kv_heads = num_heads if n_kv_heads is None else n_kv_heads
    assert self.n_kv_heads is not None
    self.repeats = self.num_heads // self.n_kv_heads
    self.scale = self.head_dim**-0.5
    self.dropout = dropout
    self.bias = bias
    self.add_zero_attn = add_zero_attn

    self.attn_type = attn_type
    self.attn_func = ATTN_BACKENDS[self.attn_type]
    self.backend = self._flash_backend if self.attn_type == "flash" else self._torch_backend
    if window_size is None:
        self.window_size = (-1, -1)
    else:
        assert attn_type == "flash"
        assert window_size % 2 == 0
        self.window_size = (window_size // 2, window_size // 2)

    self.wq = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=self.bias)
    self.wk = nn.Linear(self.embed_dim, self.n_kv_heads * self.head_dim, bias=self.bias)
    self.wv = nn.Linear(self.embed_dim, self.n_kv_heads * self.head_dim, bias=self.bias)
    self.wo = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=self.bias)

forward

Attention forward pass.

Parameters:

Name	Type	Description	Default
q	torch.Tensor	Queries of shape (batch, q_len, dim).	required
k	torch.Tensor	Keys of shape (batch, kv_len, dim).	required
v	torch.Tensor	Values of shape (batch, kv_len, dim).	required
q_mask	torch.BoolTensor	Mask for the queries, by default None.	None
kv_mask	torch.BoolTensor	Mask for the keys and values, by default None.	None
attn_mask	torch.BoolTensor	Full attention mask, by default None.	None

Returns:

Type	Description
torch.Tensor	Output of shape (batch, q_len, dim).

Source code in salt/models/transformer_v2.py

def forward(
    self,
    q: Tensor,
    k: Tensor,
    v: Tensor,
    q_mask: BoolTensor | None = None,
    kv_mask: BoolTensor | None = None,
    attn_mask: BoolTensor | None = None,
) -> Tensor:
    """Attention forward pass.

    Parameters
    ----------
    q : Tensor
        Queries of shape (batch, q_len, dim).
    k : Tensor
        Keys of shape (batch, kv_len, dim).
    v : Tensor
        Values of shape (batch, kv_len, dim).
    q_mask : BoolTensor, optional
        Mask for the queries, by default None.
    kv_mask : BoolTensor, optional
        Mask for the keys and values, by default None.
    attn_mask : BoolTensor, optional
        Full attention mask, by default None.

    Returns
    -------
    Tensor
        Output of shape (batch, q_len, dim).
    """
    # combine masks
    attn_mask = merge_masks(q_mask, kv_mask, attn_mask, q.shape, k.shape)

    # input projections
    q, k, v = self.wq(q), self.wk(k), self.wv(v)

    # add a dummy token to attend to - avoids nan when all tokens are padded
    if self.add_zero_attn:
        batch = q.shape[0]
        zero_attn_shape = (batch, 1, self.embed_dim)
        k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
        v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
        if attn_mask is not None:
            attn_mask = nn.functional.pad(attn_mask, (0, 1), value=False)
        if kv_mask is not None:
            kv_mask = nn.functional.pad(kv_mask, (0, 1), value=False)

    # run attention
    output = self.backend(q, k, v, attn_mask)

    # return output projection
    return self.wo(output)

salt.models.transformer_v2.SelfAttention

Bases: torch.nn.Module

Self attention module.

Parameters:

Name	Type	Description	Default
embed_dim	int	Dimension of the input.	required
kwargs	dict	Keyword arguments for salt.models.transformer_v2.Attention.	{}

Source code in salt/models/transformer_v2.py

def __init__(self, embed_dim: int, **kwargs):
    """Self attention module.

    Parameters
    ----------
    embed_dim : int
        Dimension of the input.
    kwargs : dict
        Keyword arguments for
        [salt.models.transformer_v2.Attention][salt.models.transformer_v2.Attention].
    """
    super().__init__()
    self.embed_dim = embed_dim
    self.attention = Attention(embed_dim=embed_dim, **kwargs)

salt.models.transformer_v2.GLU

Bases: torch.nn.Module

Dense update with gated linear unit.

See 2002.05202.

Parameters:

Name	Type	Description	Default
embed_dim	int	Dimension of the input and output.	required
hidden_dim	int | None	Dimension of the hidden layer. If None, defaults to embed_dim * 2.	None
activation	str	Activation function.	'ReLU'
bias	bool	Whether to include bias in the linear layers.	True
gated	bool	Whether to gate the output of the hidden layer.	False

Source code in salt/models/transformer_v2.py

def __init__(
    self,
    embed_dim: int,
    hidden_dim: int | None = None,
    activation: str = "ReLU",
    bias: bool = True,
    gated: bool = False,
):
    """Dense update with gated linear unit.

    See [2002.05202](https://arxiv.org/abs/2002.05202).

    Parameters
    ----------
    embed_dim : int
        Dimension of the input and output.
    hidden_dim : int | None, optional
        Dimension of the hidden layer. If None, defaults to embed_dim * 2.
    activation : str, optional
        Activation function.
    bias : bool, optional
        Whether to include bias in the linear layers.
    gated : bool, optional
        Whether to gate the output of the hidden layer.
    """
    super().__init__()

    if hidden_dim is None:
        hidden_dim = embed_dim * 2

    self.in_proj = nn.Linear(embed_dim, hidden_dim, bias=bias)
    self.out_proj = nn.Linear(hidden_dim, embed_dim, bias=bias)
    self.gate = None
    if gated:
        self.gate = nn.Linear(embed_dim, hidden_dim, bias=bias)
    self.activation = getattr(nn, activation)()

salt.models.transformer_v2.EncoderLayer

Bases: torch.nn.Module

Encoder layer consisting of a self-attention and a feed-forward layer.

Parameters:

Name	Type	Description	Default
embed_dim	int	Dimension of the embeddings at each layer.	required
norm	str	Normalization style, by default "LayerNorm".	'LayerNorm'
dense_kwargs	dict | None	Keyword arguments for salt.models.transformer_v2.GLU.	None
attn_kwargs	dict | None	Keyword arguments for salt.models.transformer_v2.SelfAttention.	None

Source code in salt/models/transformer_v2.py

def __init__(
    self,
    embed_dim: int,
    norm: str = "LayerNorm",
    dense_kwargs: dict | None = None,
    attn_kwargs: dict | None = None,
):
    """Encoder layer consisting of a self-attention and a feed-forward layer.

    Parameters
    ----------
    embed_dim : int
        Dimension of the embeddings at each layer.
    norm : str, optional
        Normalization style, by default "LayerNorm".
    dense_kwargs : dict | None, optional
        Keyword arguments for [salt.models.transformer_v2.GLU][salt.models.transformer_v2.GLU].
    attn_kwargs : dict | None, optional
        Keyword arguments for
        [salt.models.transformer_v2.SelfAttention][salt.models.transformer_v2.SelfAttention].
    """
    super().__init__()
    if attn_kwargs is None:
        attn_kwargs = {}
    if dense_kwargs is None:
        dense_kwargs = {}
    self.embed_dim = embed_dim
    self.attn = SelfAttention(embed_dim=embed_dim, **attn_kwargs)
    self.attn_norm = getattr(layernorms, norm)(embed_dim)
    self.dense = GLU(embed_dim, **dense_kwargs)
    self.dense_norm = getattr(layernorms, norm)(embed_dim)

salt.models.transformer_v2.TransformerV2

Bases: torch.nn.Module

Transformer model consisting of a series of stacked Transformer encoder layers.

Parameters:

Name	Type	Description	Default
num_layers	int	Number of layers.	required
embed_dim	int	Dimension of the embeddings at each layer.	required
out_dim	int | None	Optionally project the output to a different dimension.	None
norm	str	Normalization style, by default "LayerNorm".	'LayerNorm'
kwargs	dict	Keyword arguments for [salt.models.transformer_v2.EncoderLayer].	{}

Source code in salt/models/transformer_v2.py

def __init__(
    self,
    num_layers: int,
    embed_dim: int,
    out_dim: int | None = None,
    norm: str = "LayerNorm",
    **kwargs,
):
    """Transformer model consisting of a series of stacked Transformer encoder layers.

    Parameters
    ----------
    num_layers : int
        Number of layers.
    embed_dim : int
        Dimension of the embeddings at each layer.
    out_dim : int | None, optional
        Optionally project the output to a different dimension.
    norm : str, optional
        Normalization style, by default "LayerNorm".
    kwargs : dict
        Keyword arguments for [salt.models.transformer_v2.EncoderLayer].
    """
    super().__init__()
    self.num_layers = num_layers
    self.embed_dim = embed_dim

    self.layers = torch.nn.ModuleList([
        EncoderLayer(embed_dim=embed_dim, norm=norm, **kwargs) for _ in range(num_layers)
    ])
    self.out_norm = getattr(layernorms, norm)(embed_dim if out_dim is None else out_dim)
    self.out_proj = None
    if out_dim is not None:
        self.out_proj = nn.Linear(self.embed_dim, out_dim)

---

Last update: January 25, 2024
Created: January 25, 2024