Bare Vaswani encoder-decoder Transformer trunk.
Implements the canonical seq2seq architecture of Vaswani, Shazeer,
Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin, 2017. A
stack of encoder layers (multi-head self-attention + FFN)
processes the source input_ids; a stack of decoder
layers (masked self-attention + cross-attention + FFN) consumes the
right-shifted decoder_input_ids and the encoder memory. Token
embeddings are scaled by and combined
with sinusoidal positional encodings before the trunk.
Forward inputs / outputs follow the Seq2SeqLMOutput contract.
For task-specific outputs use the TransformerForXxx wrappers; this
bare trunk is exposed so callers needing raw decoder hidden states
(e.g. for custom heads) can skip the default LM projection.
Parameters
configTransformerConfigshare_embeddings /
tie_word_embeddings toggles.Attributes
src_tok_embnn.Embedding(vocab_size, d_model).tgt_tok_embnn.Embeddingsrc_tok_emb
when config.share_embeddings is True.positional_encodingnn.SinusoidalEmbedding(max_position_embeddings, d_model) (Vaswani 2017 §3.5).dropoutnn.Dropouttransformernn.Transformernn.TransformerEncoder + nn.TransformerDecoder.Notes
Reference: Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin, "Attention Is All You Need", NeurIPS, 2017 (arXiv:1706.03762).
Scaled dot-product attention:
with . The decoder applies an additional lower-triangular causal mask on its self-attention so position cannot attend to positions .
Examples
>>> import lucid
>>> from lucid.models.text.transformer import TransformerConfig, TransformerModel
>>> cfg = TransformerConfig(num_hidden_layers=2, num_decoder_layers=2,
... hidden_size=64, num_attention_heads=2,
... intermediate_size=128, vocab_size=1000)
>>> model = TransformerModel(cfg).eval()
>>> src = lucid.tensor([[1, 234, 567, 2]])
>>> tgt = lucid.tensor([[1, 100, 200]])
>>> out = model(src, decoder_input_ids=tgt)
>>> out.logits.shape # decoder hidden (B=1, T_tgt=3, d_model=64)
(1, 3, 64)
>>> out.encoder_last_hidden_state.shape
(1, 4, 64)Used by 2
Constructors
1Instance methods
6decode(tgt_ids: Tensor, memory: Tensor, tgt_attention_mask: Tensor | None = None, memory_attention_mask: Tensor | None = None, past_key_value: EncoderDecoderCache | None = None, use_cache: bool = False)Run only the decoder against precomputed encoder memory.
When use_cache is set, past_key_value accumulates the decoder
self-attention keys/values (and caches the cross-attention projection
of memory once), so each step only needs the freshly appended
tgt_ids token(s).
Parameters
tgt_idsTensor(B, T) target token ids to decode this step.memoryTensor(B, S, d_model) encoder output to cross-attend.(B, T) target padding mask.(B, S) source padding mask.EncoderDecoderCache accumulating
decoder self-attention and the once-computed
cross-attention key/value.use_cachebool= Falsepast_key_value so each step only
decodes the freshly appended token(s).Returns
Tensor(B, T, d_model) decoder hidden states (the ForSeq2SeqLM
Run only the encoder; returns memory (B, S, d_model).
forward(input_ids: Tensor, decoder_input_ids: Tensor, attention_mask: Tensor | None = None, decoder_attention_mask: Tensor | None = None)Full encode + decode pass.
Parameters
Returns
Seq2SeqLMOutputclass:Seq2SeqLMOutput carrying decoder hidden states as
Target-side embedding — used to tie an LM head against.