Attention and MultiHeadAttention

The mechanism that ate machine learning

Attention lets a model selectively focus on different parts of the input when producing each output. Rather than compressing an entire sequence into a single fixed-size vector (the bottleneck of vanilla seq2seq), attention lets the decoder look back at all encoder states and decide which are most relevant at each step.

Self-attention goes further: every position in a sequence attends to every other position simultaneously. This is the key innovation in Transformer architectures. Multi-Head Attention runs several attention operations in parallel, each learning different relationship types (syntactic, semantic, positional).

TensorFlow provides tf.keras.layers.MultiHeadAttention as a first-class layer. The signature is intuitive: pass query, key, value. For self-attention they're all the same tensor; for cross-attention (decoder attending to encoder), query is decoder, key/value are encoder.

Code

MultiHeadAttention 사용·python

import tensorflow as tf
from tensorflow.keras import layers

mha = layers.MultiHeadAttention(num_heads=8, key_dim=64)

# Self-attention: query = key = value
batch_size, seq_len, d_model = 32, 100, 512
x = tf.random.normal((batch_size, seq_len, d_model))

output, attention_weights = mha(
    query=x, value=x, key=x,
    return_attention_scores=True,
)
# output:            (32, 100, 512)
# attention_weights: (32, num_heads, 100, 100)

# Cross-attention: decoder attends to encoder output
decoder_seq = tf.random.normal((32, 20, 512))
encoder_seq = tf.random.normal((32, 100, 512))
cross_out = mha(query=decoder_seq, value=encoder_seq, key=encoder_seq)

Transformer encoder block — 핵심 빌딩 블록·python

import tensorflow as tf
from tensorflow.keras import layers

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, ff_dim, dropout_rate=0.1):
        super().__init__()
        self.attention = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=d_model // num_heads
        )
        self.ffn = tf.keras.Sequential([
            layers.Dense(ff_dim, activation='relu'),
            layers.Dense(d_model),
        ])
        self.norm1 = layers.LayerNormalization(epsilon=1e-6)
        self.norm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(dropout_rate)
        self.dropout2 = layers.Dropout(dropout_rate)

    def call(self, x, training=False):
        # Self-attention with residual + layer norm
        attn = self.attention(query=x, value=x, key=x, training=training)
        attn = self.dropout1(attn, training=training)
        out1 = self.norm1(x + attn)

        # Feedforward with residual + layer norm
        ffn = self.ffn(out1)
        ffn = self.dropout2(ffn, training=training)
        return self.norm2(out1 + ffn)

Attention and MultiHeadAttention

The mechanism that ate machine learning

Code

External links

Progress

Comments 0