Transformer¶

Simple transformer-based sequence-to-sequence model with sinusoidal positional encoding.

Overview¶

SimpleTransformer implements the encoder-decoder Transformer architecture from "Attention Is All You Need" (Vaswani et al., 2017), using the paper's sinusoidal positional encoding.

Architecture¶

┌─────────────────────────────────────────────────┐
│                    ENCODER                       │
│  ┌───────────────────────────────────────────┐  │
│  │  Token Embedding + Sinusoidal Position    │  │
│  │   src_vocab_size → d_model                │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                           │
│  ┌───────────────────────────────────────────┐  │
│  │         Transformer Encoder × N           │  │
│  │   • Multi-Head Self-Attention             │  │
│  │   • Feed-Forward Network                  │  │
│  │   • LayerNorm + Residual                  │  │
│  └───────────────────────────────────────────┘  │
└─────────────────────────────────────────────────┘
                       ↓
              Encoder Output (Memory)
                       ↓
┌─────────────────────────────────────────────────┐
│                    DECODER                       │
│  ┌───────────────────────────────────────────┐  │
│  │  Token Embedding + Sinusoidal Position    │  │
│  │   tgt_vocab_size → d_model                │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                           │
│  ┌───────────────────────────────────────────┐  │
│  │         Transformer Decoder × N           │  │
│  │   • Masked Multi-Head Self-Attention      │  │
│  │   • Cross-Attention (to encoder)          │  │
│  │   • Feed-Forward Network                  │  │
│  │   • LayerNorm + Residual                  │  │
│  └───────────────────────────────────────────┘  │
│                      ↓                           │
│  ┌───────────────────────────────────────────┐  │
│  │         Linear Generator                  │  │
│  │   d_model → tgt_vocab_size                │  │
│  └───────────────────────────────────────────┘  │
└─────────────────────────────────────────────────┘

Quick Start¶

from torchlingo.models import SimpleTransformer

model = SimpleTransformer(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    d_model=512,
    n_heads=8,
    num_encoder_layers=6,
    num_decoder_layers=6,
)

# Training
logits = model(src_batch, tgt_batch)  # [batch, tgt_len, vocab]

# Inference
memory = model.encode(src_batch)
logits = model.decode(tgt_batch, memory)

API Reference¶

SimpleTransformer ¶

SimpleTransformer(src_vocab_size: int, tgt_vocab_size: int, d_model: int | None = None, n_heads: int | None = None, num_encoder_layers: int | None = None, num_decoder_layers: int | None = None, d_ff: int | None = None, max_seq_length: int | None = None, dropout: float | None = None, pad_idx: int | None = None, config: Config | None = None)

Bases: Module

Transformer encoder-decoder model with sinusoidal positional encoding.

A standard transformer architecture combining an encoder and decoder, each composed of stacked multi-head self-attention and feed-forward layers. Uses the fixed sinusoidal positional encoding from Vaswani et al. (2017). Embeddings are scaled by sqrt(d_model) to prevent vanishing gradients.

Parameters:

Name	Type	Description	Default
`src_vocab_size`	`int`	Size of the source vocabulary.	required
`tgt_vocab_size`	`int`	Size of the target vocabulary.	required
`d_model`	`int`	Model hidden dimension. Defaults to 512.	`None`
`n_heads`	`int`	Number of attention heads. Defaults to 8.	`None`
`num_encoder_layers`	`int`	Number of encoder transformer blocks. Defaults to 6.	`None`
`num_decoder_layers`	`int`	Number of decoder transformer blocks. Defaults to 6.	`None`
`d_ff`	`int`	Feed-forward inner dimension. Defaults to 2048.	`None`
`max_seq_length`	`int`	Maximum sequence length for the positional encoding table. Defaults to 512.	`None`
`dropout`	`float`	Dropout rate throughout the model. Defaults to 0.1.	`None`
`pad_idx`	`int`	Padding token index. Falls back to config.pad_idx.	`None`
`config`	`Config`	Configuration object. Defaults to default config.	`None`

Attributes:

Name	Type	Description
`d_model`	`int`	Model dimension.
`max_seq_length`	`int`	Maximum sequence length.
`pad_idx`	`int`	Resolved padding index.
`src_tok_emb`	`Embedding`	Source token embedding layer.
`tgt_tok_emb`	`Embedding`	Target token embedding layer.
`pos_encoding`	`SinusoidalPositionalEncoding`	Positional encoding module.
`transformer`	`Transformer`	PyTorch transformer with encoder and decoder.
`generator`	`Linear`	Output projection to target vocabulary logits.

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

def __init__(
    self,
    src_vocab_size: int,
    tgt_vocab_size: int,
    d_model: int | None = None,
    n_heads: int | None = None,
    num_encoder_layers: int | None = None,
    num_decoder_layers: int | None = None,
    d_ff: int | None = None,
    max_seq_length: int | None = None,
    dropout: float | None = None,
    pad_idx: int | None = None,
    config: Config | None = None,
) -> None:
    super().__init__()
    cfg = config if config is not None else get_default_config()
    self.pad_idx = pad_idx if pad_idx is not None else cfg.pad_idx
    d_model = d_model if d_model is not None else cfg.d_model
    n_heads = n_heads if n_heads is not None else cfg.n_heads
    num_encoder_layers = (
        num_encoder_layers
        if num_encoder_layers is not None
        else cfg.num_encoder_layers
    )
    num_decoder_layers = (
        num_decoder_layers
        if num_decoder_layers is not None
        else cfg.num_decoder_layers
    )
    d_ff = d_ff if d_ff is not None else cfg.d_ff
    max_seq_length = (
        max_seq_length if max_seq_length is not None else cfg.max_seq_length
    )
    dropout = dropout if dropout is not None else cfg.dropout

    self.d_model = d_model
    self.max_seq_length = max_seq_length
    self.dropout = dropout

    self.src_tok_emb = nn.Embedding(
        src_vocab_size, d_model, padding_idx=self.pad_idx
    )
    self.tgt_tok_emb = nn.Embedding(
        tgt_vocab_size, d_model, padding_idx=self.pad_idx
    )
    self.pos_encoding = SinusoidalPositionalEncoding(
        d_model, max_seq_length, dropout=dropout
    )
    self.transformer = nn.Transformer(
        d_model=d_model,
        nhead=n_heads,
        num_encoder_layers=num_encoder_layers,
        num_decoder_layers=num_decoder_layers,
        dim_feedforward=d_ff,
        dropout=dropout,
        batch_first=True,
    )
    # Disable the encoder's nested-tensor fast path. In eval mode with a
    # padding mask, PyTorch converts the batch to a nested tensor, and the
    # op that does it is not implemented for Apple's MPS backend -- so
    # encode() raises NotImplementedError on any Apple Silicon GPU. The fast
    # path is a padding optimization only: disabling it changes speed, not
    # results, which the decoding tests assert. nn.Transformer does not
    # expose the constructor argument, so it is set on the encoder here.
    self.transformer.encoder.use_nested_tensor = False

    self.generator = nn.Linear(d_model, tgt_vocab_size)
    self._init_parameters()

forward ¶

forward(src: Tensor, tgt: Tensor, src_key_padding_mask: Tensor | None = None, tgt_key_padding_mask: Tensor | None = None, tgt_mask: Tensor | None = None, return_attention: bool = False) -> Tensor | tuple[Tensor, Tensor]

Encode source and decode target in a single forward pass.

Automatically generates padding and causal masks if not provided. Calls encode() to process source sequences, then decode() to generate target sequence predictions.

Parameters:

Name	Type	Description	Default
`src`	`Tensor`	Source token indices of shape (batch_size, src_len).	required
`tgt`	`Tensor`	Target token indices of shape (batch_size, tgt_len).	required
`src_key_padding_mask`	`Tensor`	Boolean mask for source padding. If None, generated automatically from padding index. Defaults to None.	`None`
`tgt_key_padding_mask`	`Tensor`	Boolean mask for target padding. If None, generated automatically from padding index. Defaults to None.	`None`
`tgt_mask`	`Tensor`	Causal mask for target self-attention. If None, generated automatically. Defaults to None.	`None`
`return_attention`	`bool`	Also return the decoder's cross-attention weights, showing which source positions each target position attended to. Off by default because capturing them costs a faster attention kernel. Matches the same argument on :class:`~torchlingo.models.SimpleSeq2SeqLSTM`, so the call is identical on either architecture. Defaults to False.	`False`

Returns:

Type	Description
`Tensor \| tuple[Tensor, Tensor]`	torch.Tensor: Logits of shape (batch_size, tgt_len, tgt_vocab_size).
`Tensor \| tuple[Tensor, Tensor]`	If `return_attention` is True, returns `(logits, weights)`
`Tensor \| tuple[Tensor, Tensor]`	instead, where weights has shape (batch_size, tgt_len, src_len) and
`Tensor \| tuple[Tensor, Tensor]`	comes from the last decoder layer, heads averaged. Use
`Tensor \| tuple[Tensor, Tensor]`	func:`capture_cross_attention` directly for every layer.

Example

The same call shape as the LSTM, so tooling works on both. Call eval() first: in training mode, attention dropout randomly zeroes weights and rescales the rest, so the rows will not sum to 1 and the map you plot is not the one inference uses.

import torch model = SimpleTransformer(src_vocab_size=20, tgt_vocab_size=20, ... d_model=16, n_heads=2, ... num_encoder_layers=1, ... num_decoder_layers=1) _ = model.eval() src, tgt = torch.tensor([[2, 5, 9, 3]]), torch.tensor([[2, 7, 8]]) with torch.no_grad(): ... logits, weights = model(src, tgt, return_attention=True) weights.shape torch.Size([1, 3, 4])

Each row is a distribution over source positions, so it sums to 1:

bool(torch.allclose(weights.sum(-1), torch.ones(1, 3), atol=1e-5)) True

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

def forward(
    self,
    src: torch.Tensor,
    tgt: torch.Tensor,
    src_key_padding_mask: torch.Tensor | None = None,
    tgt_key_padding_mask: torch.Tensor | None = None,
    tgt_mask: torch.Tensor | None = None,
    return_attention: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Encode source and decode target in a single forward pass.

    Automatically generates padding and causal masks if not provided.
    Calls encode() to process source sequences, then decode() to generate
    target sequence predictions.

    Args:
        src (torch.Tensor): Source token indices of shape (batch_size, src_len).
        tgt (torch.Tensor): Target token indices of shape (batch_size, tgt_len).
        src_key_padding_mask (torch.Tensor, optional): Boolean mask for source padding.
            If None, generated automatically from padding index. Defaults to None.
        tgt_key_padding_mask (torch.Tensor, optional): Boolean mask for target padding.
            If None, generated automatically from padding index. Defaults to None.
        tgt_mask (torch.Tensor, optional): Causal mask for target self-attention.
            If None, generated automatically. Defaults to None.
        return_attention (bool, optional): Also return the decoder's
            cross-attention weights, showing which source positions each
            target position attended to. Off by default because capturing
            them costs a faster attention kernel. Matches the same argument
            on :class:`~torchlingo.models.SimpleSeq2SeqLSTM`, so the call is
            identical on either architecture. Defaults to False.

    Returns:
        torch.Tensor: Logits of shape (batch_size, tgt_len, tgt_vocab_size).
        If ``return_attention`` is True, returns ``(logits, weights)``
        instead, where weights has shape (batch_size, tgt_len, src_len) and
        comes from the **last** decoder layer, heads averaged. Use
        :func:`capture_cross_attention` directly for every layer.

    Example:
        The same call shape as the LSTM, so tooling works on both.
        **Call ``eval()`` first**: in training mode, attention dropout
        randomly zeroes weights and rescales the rest, so the rows will not
        sum to 1 and the map you plot is not the one inference uses.

        >>> import torch
        >>> model = SimpleTransformer(src_vocab_size=20, tgt_vocab_size=20,
        ...                           d_model=16, n_heads=2,
        ...                           num_encoder_layers=1,
        ...                           num_decoder_layers=1)
        >>> _ = model.eval()
        >>> src, tgt = torch.tensor([[2, 5, 9, 3]]), torch.tensor([[2, 7, 8]])
        >>> with torch.no_grad():
        ...     logits, weights = model(src, tgt, return_attention=True)
        >>> weights.shape
        torch.Size([1, 3, 4])

        Each row is a distribution over source positions, so it sums to 1:

        >>> bool(torch.allclose(weights.sum(-1), torch.ones(1, 3), atol=1e-5))
        True
    """
    if src_key_padding_mask is None:
        src_key_padding_mask = create_key_padding_mask(src, pad_idx=self.pad_idx)
    if tgt_key_padding_mask is None:
        tgt_key_padding_mask = create_key_padding_mask(tgt, pad_idx=self.pad_idx)
    if tgt_mask is None:
        tgt_mask = create_causal_mask(tgt.size(1), tgt.device)
    memory = self.encode(src, src_key_padding_mask)

    if not return_attention:
        return self.decode(
            tgt,
            memory,
            src_key_padding_mask=src_key_padding_mask,
            tgt_key_padding_mask=tgt_key_padding_mask,
            tgt_mask=tgt_mask,
        )

    with capture_cross_attention(self.transformer.decoder) as weights:
        logits = self.decode(
            tgt,
            memory,
            src_key_padding_mask=src_key_padding_mask,
            tgt_key_padding_mask=tgt_key_padding_mask,
            tgt_mask=tgt_mask,
        )
    # The last layer by convention: it sits closest to the output and is
    # what alignment visualizations normally show. Earlier layers are
    # available through `capture_cross_attention` if you want to compare.
    return logits, weights[-1]

encode ¶

encode(src: Tensor, src_key_padding_mask: Tensor | None = None) -> Tensor

Encode source sequence using the transformer encoder.

Embeds source tokens, adds sinusoidal positional encodings, then passes the result through the transformer encoder stack.

Parameters:

Name	Type	Description	Default
`src`	`Tensor`	Source token indices of shape (batch_size, src_len).	required
`src_key_padding_mask`	`Tensor`	Boolean mask where True indicates padding positions to ignore. Shape (batch_size, src_len). Defaults to None.	`None`

Returns:

Type	Description
`Tensor`	torch.Tensor: Encoded source representation of shape (batch_size, src_len, d_model).

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

def encode(
    self, src: torch.Tensor, src_key_padding_mask: torch.Tensor | None = None
) -> torch.Tensor:
    """Encode source sequence using the transformer encoder.

    Embeds source tokens, adds sinusoidal positional encodings, then passes
    the result through the transformer encoder stack.

    Args:
        src (torch.Tensor): Source token indices of shape (batch_size, src_len).
        src_key_padding_mask (torch.Tensor, optional): Boolean mask where True indicates
            padding positions to ignore. Shape (batch_size, src_len). Defaults to None.

    Returns:
        torch.Tensor: Encoded source representation of shape (batch_size, src_len, d_model).
    """
    src_emb = self._embed(src, is_src=True)
    return self.transformer.encoder(
        src_emb, src_key_padding_mask=src_key_padding_mask
    )

decode ¶

decode(tgt: Tensor, memory: Tensor, src_key_padding_mask: Tensor | None = None, tgt_key_padding_mask: Tensor | None = None, tgt_mask: Tensor | None = None) -> Tensor

Decode target sequence using the transformer decoder and encoder output.

Embeds target tokens, adds sinusoidal positional encodings, then passes them through the transformer decoder with cross-attention to the encoder output. The decoder is typically run with a causal mask to prevent attending to future tokens.

Parameters:

Name	Type	Description	Default
`tgt`	`Tensor`	Target token indices of shape (batch_size, tgt_len).	required
`memory`	`Tensor`	Encoded source from encoder output, shape (batch_size, src_len, d_model).	required
`src_key_padding_mask`	`Tensor`	Mask for encoder output padding. Shape (batch_size, src_len). Defaults to None.	`None`
`tgt_key_padding_mask`	`Tensor`	Mask for target padding. Shape (batch_size, tgt_len). Defaults to None.	`None`
`tgt_mask`	`Tensor`	Causal mask for target self-attention. Shape (tgt_len, tgt_len). Defaults to None.	`None`

Returns:

Type	Description
`Tensor`	torch.Tensor: Decoder output logits of shape (batch_size, tgt_len, tgt_vocab_size).

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

def decode(
    self,
    tgt: torch.Tensor,
    memory: torch.Tensor,
    src_key_padding_mask: torch.Tensor | None = None,
    tgt_key_padding_mask: torch.Tensor | None = None,
    tgt_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Decode target sequence using the transformer decoder and encoder output.

    Embeds target tokens, adds sinusoidal positional encodings, then passes
    them through the transformer decoder with cross-attention to the encoder output.
    The decoder is typically run with a causal mask to prevent attending to future tokens.

    Args:
        tgt (torch.Tensor): Target token indices of shape (batch_size, tgt_len).
        memory (torch.Tensor): Encoded source from encoder output, shape (batch_size, src_len, d_model).
        src_key_padding_mask (torch.Tensor, optional): Mask for encoder output padding.
            Shape (batch_size, src_len). Defaults to None.
        tgt_key_padding_mask (torch.Tensor, optional): Mask for target padding.
            Shape (batch_size, tgt_len). Defaults to None.
        tgt_mask (torch.Tensor, optional): Causal mask for target self-attention.
            Shape (tgt_len, tgt_len). Defaults to None.

    Returns:
        torch.Tensor: Decoder output logits of shape (batch_size, tgt_len, tgt_vocab_size).
    """
    tgt_emb = self._embed(tgt, is_src=False)
    dec = self.transformer.decoder(
        tgt_emb,
        memory,
        tgt_mask=tgt_mask,
        memory_key_padding_mask=src_key_padding_mask,
        tgt_key_padding_mask=tgt_key_padding_mask,
    )
    return self.generator(dec)

Constructor Parameters¶

Parameter	Type	Default	Description
`src_vocab_size`	int	required	Source vocabulary size
`tgt_vocab_size`	int	required	Target vocabulary size
`d_model`	int	512	Model hidden dimension
`n_heads`	int	8	Number of attention heads
`num_encoder_layers`	int	6	Encoder transformer blocks
`num_decoder_layers`	int	6	Decoder transformer blocks
`d_ff`	int	2048	Feed-forward inner dimension
`max_seq_length`	int	512	Maximum sequence length
`dropout`	float	0.1	Dropout rate
`pad_idx`	int	0	Padding token index
`config`	Config	None	Configuration object

Examples¶

Basic Training¶

import torch
from torchlingo.models import SimpleTransformer
from torchlingo.config import Config

config = Config(d_model=256, n_heads=8)
model = SimpleTransformer(
    src_vocab_size=10000,
    tgt_vocab_size=10000,
    config=config,
)

# Dummy data
src = torch.randint(0, 10000, (32, 20))  # [batch, src_len]
tgt = torch.randint(0, 10000, (32, 25))  # [batch, tgt_len]

# Forward pass
logits = model(src, tgt[:, :-1])  # [32, 24, 10000]

# Compute loss
criterion = torch.nn.CrossEntropyLoss(ignore_index=0)
loss = criterion(
    logits.reshape(-1, logits.size(-1)),
    tgt[:, 1:].reshape(-1)
)

Inference (Greedy Decoding)¶

def greedy_decode(model, src, tgt_vocab, max_len=50, device="cpu"):
    model.eval()
    src = src.to(device)

    # Encode source
    memory = model.encode(src)

    # Start with SOS
    ys = torch.tensor([[tgt_vocab.sos_idx]], device=device)

    for _ in range(max_len):
        logits = model.decode(ys, memory)
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        ys = torch.cat([ys, next_token], dim=1)

        if next_token.item() == tgt_vocab.eos_idx:
            break

    return ys[0].tolist()

Custom Masking¶

# Manual padding mask
src_pad_mask = (src == pad_idx)  # True where padding

# Manual causal mask
tgt_len = tgt.size(1)
causal_mask = torch.triu(
    torch.ones(tgt_len, tgt_len, dtype=torch.bool),
    diagonal=1
)

logits = model(
    src, tgt,
    src_key_padding_mask=src_pad_mask,
    tgt_mask=causal_mask,
)

Key Features¶

Sinusoidal Positional Encoding¶

Fixed sine/cosine waves of geometrically increasing wavelengths are added to the token embeddings:

No learned parameters; valid for any position
Encodings for nearby positions are related by simple rotations, which helps the model attend by relative offset
Matches the original "Attention Is All You Need" recipe

Automatic Masking¶

The model automatically generates:

Padding masks: Prevent attention to PAD tokens
Causal masks: Prevent decoder from seeing future tokens

Embedding Scaling¶

Embeddings are scaled by √d_model to maintain variance:

src_emb = self.src_tok_emb(src) * math.sqrt(self.d_model)

Xavier Initialization¶

Weights are initialized using Xavier uniform for stable training:

for m in self.modules():
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)

Model Variants¶

Tiny (Demo/Testing)¶

model = SimpleTransformer(
    src_vocab_size, tgt_vocab_size,
    d_model=64, n_heads=2,
    num_encoder_layers=1, num_decoder_layers=1,
)
# ~500K params, trains in seconds

Base¶

model = SimpleTransformer(
    src_vocab_size, tgt_vocab_size,
    d_model=512, n_heads=8,
    num_encoder_layers=6, num_decoder_layers=6,
)
# ~65M params, good quality

Large¶

model = SimpleTransformer(
    src_vocab_size, tgt_vocab_size,
    d_model=1024, n_heads=16,
    num_encoder_layers=6, num_decoder_layers=6,
    d_ff=4096,
)
# ~200M params, high quality

Computational Considerations¶

Aspect	Complexity
Self-attention	O(n² × d)
Memory	O(n² + n × d)
Parameters	O(L × d²)

Where n = sequence length, d = d_model, L = num layers.

Long Sequences

Attention has O(n²) complexity. For very long sequences (>1000 tokens), consider using flash attention or other efficient attention implementations.