Attention¶

Readable, from-scratch attention mechanisms for the LSTM decoder.

Overview¶

The Transformer gets its attention from torch.nn.Transformer, where the mechanism is real but invisible. This module is the version you can actually read — roughly fifteen lines per scorer.

Every attention mechanism here does the same three things:

Score each source position against the current decoder state.
Softmax those scores into weights that sum to 1.
Average the encoder outputs using those weights, producing a context vector.

Only step 1 differs between the two implementations.

Scorer	Score	Paper	Parameters
`"dot"`	`h_t · h_s`	Luong et al., 2015	none
`"additive"`	`v · tanh(W_dec h_t + W_enc h_s)`	Bahdanau et al., 2014	three matrices

Read them in publication order

Bahdanau's additive score learns how to compare decoder and encoder states. Luong's dot score notices that if the two already live in the same space, an inner product will do — no parameters at all. The Transformer then keeps the dot product, adds a 1/sqrt(d_k) scale, and runs it in parallel heads. Reading the first two makes the third one familiar rather than alien.

Quick Start¶

from torchlingo.models import SimpleSeq2SeqLSTM

model = SimpleSeq2SeqLSTM(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    attention=True,
    attn_type="dot",      # or "additive"
)

logits, weights = model(src, tgt, return_attention=True)
# weights: [batch, tgt_len, src_len], each row summing to 1

Attention is off by default. That keeps the classic bottlenecked seq2seq as the baseline and makes the with-versus-without comparison a single visible flag.

Checkpoints are not interchangeable

Enabling attention adds parameters, so a checkpoint trained with attention=False will not load into a model built with attention=True, or the reverse. Build the model the same way you trained it.

Using the mechanisms directly¶

from torchlingo.models.attention import build_attention

attn = build_attention("dot", hidden_dim=512)
context, weights = attn(dec_out, enc_out, src_pad_mask)

Argument	Shape	Meaning
`dec_out`	`[batch, tgt_len, hidden]`	Decoder states — the queries
`enc_out`	`[batch, src_len, hidden]`	Encoder outputs — keys and values
`src_pad_mask`	`[batch, src_len]`	`True` marks padding to ignore
→ `context`	`[batch, tgt_len, hidden]`	Weighted average of `enc_out`
→ `weights`	`[batch, tgt_len, src_len]`	The alignment matrix

Padding¶

Masked positions are pushed to the most negative finite value before the softmax, not to -inf. Both give a weight of essentially zero, but -inf produces NaN if an entire row is padded — a silent corruption that is unpleasant to track down later.

Masking the weights is not the whole story

The encoder LSTM itself must also stop at the real end of each sentence, or the state handed to the decoder describes padding rather than the sentence. SimpleSeq2SeqLSTM.encode_source packs the batch to guarantee this. Masking attention alone does not fix it — the leak is in the recurrence, not the alignment.

Seeing what it learned¶

Attention weights are the most directly inspectable quantity in an NMT model. See Visualization for the renderers, and run:

python examples/attention_alignment.py

That example trains on a task whose correct alignment is known in advance (the target is the reversed, word-substituted source), so "did attention work?" becomes a measurable number instead of a heatmap you squint at:

configuration            val loss   alignment acc
----------------------------------------------------------------
no attention               1.1031             n/a
dot (Luong)                0.6801           98.4%
additive (Bahdanau)        0.6533           93.6%

Chance is about 12.5% on that task.

API Reference¶

attention ¶

Attention mechanisms for sequence-to-sequence models.

This module contains the readable, from-scratch attention implementations that the rest of the library builds on. It exists for a specific reason: the Transformer in :mod:torchlingo.models.transformer_simple gets its attention from :class:torch.nn.Transformer, where the mechanism is real but invisible. Here it is roughly fifteen lines you can read.

The problem attention solves. A plain LSTM encoder-decoder squeezes the entire source sentence through one fixed-size hidden state. Every source token is encoded and then thrown away; only the final state survives to the decoder. That is an information bottleneck, and it is why translation quality falls off sharply on long sentences. Attention removes it by letting the decoder look back at every encoder output, weighted by relevance, at every step.

Two scoring functions. Both compute the same three things -- a score per source position, a softmax over those scores, and a weighted average of encoder outputs -- and differ only in how the score is computed:

============ ==================================== ============================ Scorer Score Notes ============ ==================================== ============================ additive v @ tanh(W_dec h_t + W_enc h_s) Bahdanau et al., 2014 dot h_t @ h_s Luong et al., 2015 ============ ==================================== ============================

The progression is worth following. Bahdanau's additive score learns three matrices to compare decoder and encoder states. Luong's dot score observes that if the two live in the same space you can just take an inner product -- no parameters at all. The Transformer then keeps the dot product, adds a 1 / sqrt(d_k) scale, and applies it in parallel heads. Reading these two in order makes the third one familiar rather than alien.

Note

The list of supported scorers is :data:torchlingo.config.ATTENTION_TYPES, imported here rather than defined here. Config validates lstm_attn_type while it is still being imported, and every model module imports Config, so defining the list in this module would make that import circular.

Typical usage

import torch attn = DotProductAttention() dec_out = torch.randn(2, 3, 16) # (batch, tgt_len, hidden) enc_out = torch.randn(2, 5, 16) # (batch, src_len, hidden) context, weights = attn(dec_out, enc_out) context.shape, weights.shape (torch.Size([2, 3, 16]), torch.Size([2, 3, 5])) torch.allclose(weights.sum(-1), torch.ones(2, 3)) True

DotProductAttention ¶

Bases: Module

Luong dot-product attention (Luong et al., 2015).

Scores each source position by taking the inner product of the decoder state with the encoder state. This has no learned parameters at all -- the comparison is pure geometry, which is what makes it the direct ancestor of Transformer self-attention.

Requires the decoder and encoder hidden sizes to match, since an inner product is only defined between vectors in the same space.

Note

The Transformer divides these scores by sqrt(d_k) before the softmax; this implementation does not, matching Luong's paper. The scale matters at large d_k, where unscaled dot products grow large enough to push the softmax into saturation and flatten the gradient.

Examples:

>>> import torch
>>> attn = DotProductAttention()
>>> context, weights = attn(torch.randn(1, 2, 8), torch.randn(1, 4, 8))
>>> weights.shape
torch.Size([1, 2, 4])

forward ¶

forward(dec_out: Tensor, enc_out: Tensor, src_pad_mask: Tensor | None = None) -> tuple[Tensor, Tensor]

Compute context vectors and attention weights.

Parameters:

Name	Type	Description	Default
`dec_out`	`Tensor`	Decoder states, shape (batch, tgt_len, hidden).	required
`enc_out`	`Tensor`	Encoder outputs, shape (batch, src_len, hidden).	required
`src_pad_mask`	`Tensor`	Boolean mask of shape (batch, src_len); `True` marks padding to be ignored.	`None`

Returns:

Type	Description
`tuple[Tensor, Tensor]`	tuple[torch.Tensor, torch.Tensor]: The context vectors of shape (batch, tgt_len, hidden), and the attention weights of shape (batch, tgt_len, src_len).

Raises:

Type	Description
`ValueError`	If the decoder and encoder hidden sizes differ.

Source code in src/torchlingo/models/attention.py

def forward(
    self,
    dec_out: torch.Tensor,
    enc_out: torch.Tensor,
    src_pad_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute context vectors and attention weights.

    Args:
        dec_out (torch.Tensor): Decoder states, shape (batch, tgt_len, hidden).
        enc_out (torch.Tensor): Encoder outputs, shape (batch, src_len, hidden).
        src_pad_mask (torch.Tensor, optional): Boolean mask of shape
            (batch, src_len); ``True`` marks padding to be ignored.

    Returns:
        tuple[torch.Tensor, torch.Tensor]: The context vectors of shape
            (batch, tgt_len, hidden), and the attention weights of shape
            (batch, tgt_len, src_len).

    Raises:
        ValueError: If the decoder and encoder hidden sizes differ.
    """
    if dec_out.size(-1) != enc_out.size(-1):
        raise ValueError(
            "Dot-product attention requires matching hidden sizes, got "
            f"decoder {dec_out.size(-1)} and encoder {enc_out.size(-1)}. "
            "Use attn_type='additive', which learns a projection instead."
        )
    # "How well does each decoder state match each source state?"
    # (batch, tgt_len, hidden) @ (batch, hidden, src_len) -> (batch, tgt_len, src_len)
    scores = torch.bmm(dec_out, enc_out.transpose(1, 2))
    weights = _masked_softmax(scores, src_pad_mask)
    # Weighted average of the source representations.
    context = torch.bmm(weights, enc_out)
    return context, weights

AdditiveAttention ¶

AdditiveAttention(hidden_dim: int, enc_dim: int | None = None, attn_dim: int | None = None)

Bases: Module

Bahdanau additive attention (Bahdanau et al., 2014).

Scores each source position with a small one-hidden-layer network over the concatenated decoder and encoder states. Because the comparison is learned rather than geometric, the two sides may have different hidden sizes.

This is the mechanism from the paper that introduced attention to NMT. It works well, but note that the tanh scoring network does not reappear in the Transformer -- the dot product does.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Decoder hidden size.	required
`enc_dim`	`int`	Encoder hidden size. Defaults to `hidden_dim`.	`None`
`attn_dim`	`int`	Width of the scoring network's hidden layer. Defaults to `hidden_dim`.	`None`

Attributes:

Name	Type	Description
`W_dec`	`Linear`	Projects the decoder state into the scoring space.
`W_enc`	`Linear`	Projects each encoder state into the scoring space.
`v`	`Linear`	Collapses the scoring space to one scalar per position.

Examples:

>>> import torch
>>> attn = AdditiveAttention(hidden_dim=8, enc_dim=16)
>>> context, weights = attn(torch.randn(1, 2, 8), torch.randn(1, 4, 16))
>>> context.shape, weights.shape
(torch.Size([1, 2, 16]), torch.Size([1, 2, 4]))

Source code in src/torchlingo/models/attention.py

def __init__(
    self,
    hidden_dim: int,
    enc_dim: int | None = None,
    attn_dim: int | None = None,
):
    super().__init__()
    enc_dim = enc_dim if enc_dim is not None else hidden_dim
    attn_dim = attn_dim if attn_dim is not None else hidden_dim
    self.W_dec = nn.Linear(hidden_dim, attn_dim, bias=False)
    self.W_enc = nn.Linear(enc_dim, attn_dim, bias=False)
    self.v = nn.Linear(attn_dim, 1, bias=False)

forward ¶

forward(dec_out: Tensor, enc_out: Tensor, src_pad_mask: Tensor | None = None) -> tuple[Tensor, Tensor]

Compute context vectors and attention weights.

Parameters:

Name	Type	Description	Default
`dec_out`	`Tensor`	Decoder states, shape (batch, tgt_len, hidden).	required
`enc_out`	`Tensor`	Encoder outputs, shape (batch, src_len, enc_dim).	required
`src_pad_mask`	`Tensor`	Boolean mask of shape (batch, src_len); `True` marks padding to be ignored.	`None`

Returns:

Type	Description
`tuple[Tensor, Tensor]`	tuple[torch.Tensor, torch.Tensor]: The context vectors of shape (batch, tgt_len, enc_dim), and the attention weights of shape (batch, tgt_len, src_len).

Source code in src/torchlingo/models/attention.py

def forward(
    self,
    dec_out: torch.Tensor,
    enc_out: torch.Tensor,
    src_pad_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute context vectors and attention weights.

    Args:
        dec_out (torch.Tensor): Decoder states, shape (batch, tgt_len, hidden).
        enc_out (torch.Tensor): Encoder outputs, shape (batch, src_len, enc_dim).
        src_pad_mask (torch.Tensor, optional): Boolean mask of shape
            (batch, src_len); ``True`` marks padding to be ignored.

    Returns:
        tuple[torch.Tensor, torch.Tensor]: The context vectors of shape
            (batch, tgt_len, enc_dim), and the attention weights of shape
            (batch, tgt_len, src_len).
    """
    # Broadcast decoder over source and source over decoder, so that every
    # (target position, source position) pair gets its own score.
    dec_proj = self.W_dec(dec_out).unsqueeze(2)  # (batch, tgt_len, 1, attn_dim)
    enc_proj = self.W_enc(enc_out).unsqueeze(1)  # (batch, 1, src_len, attn_dim)
    scores = self.v(torch.tanh(dec_proj + enc_proj)).squeeze(-1)
    weights = _masked_softmax(scores, src_pad_mask)
    context = torch.bmm(weights, enc_out)
    return context, weights

build_attention ¶

build_attention(attn_type: str, hidden_dim: int, enc_dim: int | None = None) -> Module

Construct an attention module by name.

Parameters:

Name	Type	Description	Default
`attn_type`	`str`	One of `"dot"` (Luong) or `"additive"` (Bahdanau).	required
`hidden_dim`	`int`	Decoder hidden size.	required
`enc_dim`	`int`	Encoder hidden size. Defaults to `hidden_dim`.	`None`

Returns:

Type	Description
`Module`	nn.Module: An attention module whose `forward` takes `(dec_out, enc_out, src_pad_mask)` and returns `(context, weights)`.

Raises:

Type	Description
`ValueError`	If `attn_type` is not a supported scorer.

Examples:

>>> build_attention("dot", hidden_dim=16)
DotProductAttention()

Source code in src/torchlingo/models/attention.py

def build_attention(
    attn_type: str,
    hidden_dim: int,
    enc_dim: int | None = None,
) -> nn.Module:
    """Construct an attention module by name.

    Args:
        attn_type (str): One of ``"dot"`` (Luong) or ``"additive"`` (Bahdanau).
        hidden_dim (int): Decoder hidden size.
        enc_dim (int, optional): Encoder hidden size. Defaults to ``hidden_dim``.

    Returns:
        nn.Module: An attention module whose ``forward`` takes
            ``(dec_out, enc_out, src_pad_mask)`` and returns
            ``(context, weights)``.

    Raises:
        ValueError: If ``attn_type`` is not a supported scorer.

    Examples:
        >>> build_attention("dot", hidden_dim=16)
        DotProductAttention()
    """
    if attn_type == "dot":
        return DotProductAttention()
    if attn_type == "additive":
        return AdditiveAttention(hidden_dim, enc_dim=enc_dim)
    raise ValueError(
        f"Unknown attn_type '{attn_type}'. Expected one of {list(ATTENTION_TYPES)}."
    )