Attention Mechanism

1import tensorflow as tf
2from tensorflow.keras import layers
3
4# Self-Attention trong Keras
5class SimpleSelfAttention(layers.Layer):
6    def __init__(self, d_model=64):
7        super().__init__()
8        self.d_model = d_model
9        # Q, K, V projections
10        self.wq = layers.Dense(d_model)
11        self.wk = layers.Dense(d_model)
12        self.wv = layers.Dense(d_model)
13    
14    def call(self, x):
15        # x shape: (batch, seq_len, d_model)
16        Q = self.wq(x)
17        K = self.wk(x)
18        V = self.wv(x)
19        
20        # Attention scores
21        scores = tf.matmul(Q, K, transpose_b=True)  # Q × K^T
22        scores = scores / tf.sqrt(tf.cast(self.d_model, tf.float32))  # Scale
23        
24        # Softmax → weights
25        weights = tf.nn.softmax(scores, axis=-1)
26        
27        # Output = weights × V
28        output = tf.matmul(weights, V)
29        
30        return output, weights
31
32# Demo
33self_attn = SimpleSelfAttention(d_model=64)
34x = tf.random.normal((2, 5, 64))  # 2 câu, 5 từ mỗi câu, 64-dim
35
36output, weights = self_attn(x)
37print(f"Input: {x.shape}")
38print(f"Output: {output.shape}")
39print(f"\nAttention weights (câu 1):")
40print(weights[0].numpy().round(2))

1import tensorflow as tf
2from tensorflow.keras import layers
3
4# Multi-Head Attention có sẵn trong Keras!
5mha = layers.MultiHeadAttention(
6    num_heads=4,      # 4 heads
7    key_dim=64,       # Dimension của mỗi head
8)
9
10# Input
11x = tf.random.normal((2, 10, 256))  # 2 câu, 10 từ, 256-dim
12
13# Self-Attention (Q=K=V=x)
14output = mha(query=x, key=x, value=x)
15
16print(f"Input shape: {x.shape}")
17print(f"Output shape: {output.shape}")

1import tensorflow as tf
2import numpy as np
3
4def create_padding_mask(seq):
5    """Mask cho padding tokens (giá trị 0)"""
6    mask = tf.cast(tf.equal(seq, 0), tf.float32)
7    return mask[:, tf.newaxis, tf.newaxis, :]
8
9def create_look_ahead_mask(size):
10    """Mask không cho nhìn tương lai"""
11    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
12    return mask
13
14# Demo
15seq = tf.constant([[1, 2, 3, 0, 0]])  # Câu có 3 từ + 2 padding
16
17padding_mask = create_padding_mask(seq)
18print("Padding mask shape:", padding_mask.shape)
19print("Padding mask:\n", padding_mask[0, 0, 0].numpy())
20
21look_ahead = create_look_ahead_mask(5)
22print("\nLook-ahead mask:\n", look_ahead.numpy())