An attention_mask is a Boolean tensor that marks usable sequence positions and prevents attention from using padding tokens, which contain no information. In Transformer encoders, a source mask marking non-padding tokens is expanded for attention-score broadcasting. section 15.3.2
Use masks for two related but different purposes:
- Pass a non-padding source mask to encoder self-attention so tokens cannot attend to padded positions. section 15.3.2
- Pass the source padding mask to decoder cross-attention when the decoder reads encoder outputs. section 15.3.3
- Use a causal lower-triangular mask for decoder self-attention when each position must see only itself and earlier positions. section 15.3.3
Padding masks during instruction fine-tuning
A preprocessing layer pads token sequences to a fixed length and returns token_ids, a Boolean padding_mask, labels, and sample_weight. The padding mask tracks padded zero inputs, while sample_weight restricts loss computation to response tokens rather than prompts or padding. section 16.3.2
Listing 15.14 A Transformer encoder block
layer_transformer_encoder <- new_layer_class(
"TransformerEncoder",
initialize = function(hidden_dim, intermediate_dim, num_heads) {
super$initialize()
key_dim <- hidden_dim %/% num_heads
self$self_attention <- layer_multi_head_attention( #1
num_heads = num_heads,
key_dim = key_dim
)
self$self_attention_layernorm <- layer_layer_normalization() #1
self$feed_forward_1 <- layer_dense(, intermediate_dim, #2
activation = "relu") #2
self$feed_forward_2 <- layer_dense(, hidden_dim) #2
self$feed_forward_layernorm <- layer_layer_normalization() #2
},
call = function(source, source_mask) {
residual <- x <- source #3
mask <- source_mask@r[, newaxis, ] #3
x <- self$self_attention( #3
query = x, #3
key = x, #3
value = x, #3
attention_mask = mask #3
) #3
x <- x + residual #3
x <- x |> self$self_attention_layernorm() #3
residual <- x #4
x <- x |> #4
self$feed_forward_1() |> #4
self$feed_forward_2() #4
x <- x + residual #4
x <- x |> self$feed_forward_layernorm() #4
x
}
)
Inspect the Transformer encoder implementation to see source_mask reshaped to (batch_size, 1, source_length) before being passed to self-attention. section 15.3.2
preprocessor <- gemma_lm$preprocessor preprocessor$sequence_length <- 512L batch <- iter_next(as_iterator(train_ds)) .[x, y, sample_weight] <- preprocessor(batch) str(x) List of 2 $ token_ids : <jax.Array shape(2, 512), dtype = int32> $ padding_mask: <jax.Array shape(2, 512), dtype = bool> str(y) <jax.Array shape(2, 512), dtype = int32> str(sample_weight) <jax.Array shape(2, 512), dtype = bool>
Inspect the preprocessing output showing a 2-by-512 Boolean padding_mask alongside token IDs, labels, and sample weights. section 16.3.2
