Pooling Strategies: De Múltiples Tokens a Un Vector

Descripción de la cápsula

Después de que el Transformer procesa los tokens, tienes múltiples hidden states (un vector por token). Para obtener un embedding de frase completa, necesitas "pooling": reducir múltiples vectores a uno solo.

En esta cápsula aprenderás las 3 estrategias principales de pooling (mean, CLS, max), implementarás cada una con numpy, compararás sus trade-offs, y entenderás por qué mean pooling domina en sentence embeddings. También verás código práctico que simula el proceso de pooling.

Al final, podrás tomar decisiones informadas sobre qué pooling strategy usar según tu caso de uso.

El problema del pooling

Situación:

# Después del Transformer encoder:
text = "Python es un lenguaje"

# Hidden states (un vector por token):
hidden_states = [
    [0.8, 0.1, 0.3, 0.5, 0.2],  # "Python"
    [0.9, 0.3, 0.2, 0.4, 0.1],  # "es"
    [0.7, 0.2, 0.4, 0.6, 0.3],  # "un"
    [0.6, 0.4, 0.1, 0.3, 0.4]   # "lenguaje"
]

# Necesitas 1 solo embedding para toda la frase:
sentence_embedding = ???  # ¿Cómo obtenerlo?

Solución: Pooling (agregación).

Strategy #1: Mean Pooling

Qué es:

Promedio de todos los hidden states (vectores de tokens).

Fórmula:

embedding = (v₁ + v₂ + v₃ + ... + vₙ) / n

Donde:
- v₁, v₂, ... = Hidden states de cada token
- n = Cantidad de tokens

Implementación con numpy:

import numpy as np

def mean_pooling(hidden_states):
    """
    Mean pooling: Promedio de todos los hidden states
    
    Args:
        hidden_states: Array de shape (n_tokens, hidden_size)
    
    Returns:
        Embedding de shape (hidden_size,)
    """
    return np.mean(hidden_states, axis=0)

# Ejemplo
hidden_states = np.array([
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Token 1
    [0.9, 0.3, 0.2, 0.4, 0.1],  # Token 2
    [0.7, 0.2, 0.4, 0.6, 0.3],  # Token 3
    [0.6, 0.4, 0.1, 0.3, 0.4]   # Token 4
])

embedding = mean_pooling(hidden_states)
print(f"Embedding: {embedding}")
print(f"Shape: {embedding.shape}")

Output:

Embedding: [0.75 0.25 0.25 0.45 0.25]
Shape: (5,)

Explicación:

• Dim 0: (0.8 + 0.9 + 0.7 + 0.6) / 4 = 0.75
• Dim 1: (0.1 + 0.3 + 0.2 + 0.4) / 4 = 0.25
• ...

Ventajas de mean pooling:

✅ 1. Simple y efectivo

# Una línea de código:
embedding = np.mean(hidden_states, axis=0)

✅ 2. Considera todos los tokens

# Cada token contribuye al embedding final
# No descartas información

✅ 3. Robusto a longitud variable

# Funciona igual para 5 tokens o 500 tokens
embedding_short = mean_pooling(hidden_states_5)    # Shape: (768,)
embedding_long = mean_pooling(hidden_states_500)   # Shape: (768,)

Desventajas de mean pooling:

❌ 1. Tokens irrelevantes diluyen

text = "Python, um, es, like, un lenguaje"
# "um", "like" contribuyen al promedio aunque son ruido

# Mean pooling: Promedia TODO (incluido ruido)

❌ 2. No pondera importancia

# "Python" y "es" contribuyen igual al promedio
# Pero "Python" es más importante semánticamente

Strategy #2: CLS Pooling

Qué es:

Usar solo el hidden state del token [CLS] (token especial agregado al inicio).

BERT agrega [CLS] al inicio:

text = "Python es un lenguaje"

# Tokenización BERT:
tokens = ["[CLS]", "Python", "es", "un", "lenguaje"]

# Hidden states:
hidden_states = [
    [0.5, 0.6, 0.7, 0.8, 0.9],  # [CLS] ← Usamos SOLO este
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Python
    [0.9, 0.3, 0.2, 0.4, 0.1],  # es
    [0.7, 0.2, 0.4, 0.6, 0.3],  # un
    [0.6, 0.4, 0.1, 0.3, 0.4]   # lenguaje
]

# CLS pooling: Tomar solo hidden_states[0]
embedding = hidden_states[0]  # [0.5, 0.6, 0.7, 0.8, 0.9]

Implementación con numpy:

def cls_pooling(hidden_states):
    """
    CLS pooling: Usar solo el primer hidden state ([CLS] token)
    
    Args:
        hidden_states: Array de shape (n_tokens, hidden_size)
    
    Returns:
        Embedding de shape (hidden_size,)
    """
    return hidden_states[0]

# Ejemplo (asumiendo [CLS] es primer token)
hidden_states = np.array([
    [0.5, 0.6, 0.7, 0.8, 0.9],  # [CLS]
    [0.8, 0.1, 0.3, 0.5, 0.2],  # Token 1
    [0.9, 0.3, 0.2, 0.4, 0.1],  # Token 2
])

embedding = cls_pooling(hidden_states)
print(f"Embedding: {embedding}")

Output:

Embedding: [0.5 0.6 0.7 0.8 0.9]

Ventajas de CLS pooling:

✅ 1. Token entrenado específicamente

# [CLS] se entrena para "resumir" toda la frase
# Durante training de BERT, [CLS] aprende a capturar significado global

✅ 2. No diluye con tokens irrelevantes

# Tokens de ruido no afectan (no son promediados)

Desventajas de CLS pooling:

❌ 1. Ignora otros tokens

# Solo usa [CLS], descarta información de tokens 2-N
# Potencial pérdida de matices

❌ 2. Requiere token especial

# OpenAI embeddings NO usan [CLS]
# Solo funciona con modelos BERT-style

Strategy #3: Max Pooling

Qué es:

Tomar el valor máximo de cada dimensión a través de todos los tokens.

def max_pooling(hidden_states):
    """
    Max pooling: Máximo por dimensión
    
    Args:
        hidden_states: Array de shape (n_tokens, hidden_size)
    
    Returns:
        Embedding de shape (hidden_size,)
    """
    return np.max(hidden_states, axis=0)

# Ejemplo
hidden_states = np.array([
    [0.8, 0.1, 0.3, 0.5, 0.2],
    [0.9, 0.3, 0.2, 0.4, 0.1],
    [0.7, 0.2, 0.4, 0.6, 0.3],
    [0.6, 0.4, 0.1, 0.3, 0.4]
])

embedding = max_pooling(hidden_states)
print(f"Embedding: {embedding}")

Output:

Embedding: [0.9 0.4 0.4 0.6 0.4]

Explicación:

• Dim 0: max(0.8, 0.9, 0.7, 0.6) = 0.9
• Dim 1: max(0.1, 0.3, 0.2, 0.4) = 0.4
• ...

Ventajas de max pooling:

✅ Captura features más "salientes"

# Si un token tiene valor muy alto en una dimensión,
# ese valor se preserva en el embedding final

Desventajas de max pooling:

❌ 1. Pierde información promedio

# Solo captura valores máximos, no el "sentido general"

❌ 2. Menos usado en práctica

# Mean pooling y CLS pooling dominan
# Max pooling es raro en sentence embeddings

Comparación práctica: 3 pooling strategies

Código completo de comparación:

import numpy as np
from openai import OpenAI
import os
from dotenv import load_dotenv

# Setup
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Simular hidden states (en producción vienen del Transformer)
# Para demostración, generamos aleatorios
np.random.seed(42)
hidden_states = np.random.randn(10, 5)  # 10 tokens, 5 dims

print("Hidden States (10 tokens × 5 dims):")
print(hidden_states)
print()

# Mean pooling
mean_emb = np.mean(hidden_states, axis=0)
print(f"Mean Pooling: {mean_emb}")

# CLS pooling (tomar primer token)
cls_emb = hidden_states[0]
print(f"CLS Pooling:  {cls_emb}")

# Max pooling
max_emb = np.max(hidden_states, axis=0)
print(f"Max Pooling:  {max_emb}")

# Comparar con cosine similarity
def cosine_similarity(vec_a, vec_b):
    return np.dot(vec_a, vec_b) / (np.linalg.norm(vec_a) * np.linalg.norm(vec_b))

print(f"\nComparación:")
print(f"Mean vs CLS: {cosine_similarity(mean_emb, cls_emb):.4f}")
print(f"Mean vs Max: {cosine_similarity(mean_emb, max_emb):.4f}")
print(f"CLS vs Max:  {cosine_similarity(cls_emb, max_emb):.4f}")

Output (ejemplo):

Mean Pooling: [ 0.12 -0.05  0.18 -0.08  0.15]
CLS Pooling:  [ 0.50  0.60  0.70  0.80  0.90]
Max Pooling:  [ 1.45  0.85  1.20  0.95  1.10]

Comparación:
Mean vs CLS: 0.65  ← Algo similares
Mean vs Max: 0.78  ← Más similares (mean suaviza max)
CLS vs Max:  0.55  ← Menos similares

Qué pooling usa cada modelo

Tabla de referencia:

Tendencia: Mean pooling domina en sentence embeddings (SBERT lo popularizó).

Attention Mask: Pooling correcto con padding

Problema: Secuencias de diferente longitud

# Batch de 3 frases:
texts = [
    "Python",                    # 1 token
    "Python es un lenguaje",     # 4 tokens
    "JavaScript"                 # 1 token
]

# Para procesar en batch, necesitas MISMO tamaño:
# Padding (rellenar con tokens especiales):
# [Python, PAD, PAD, PAD]
# [Python, es, un, lenguaje]
# [JavaScript, PAD, PAD, PAD]

Mean pooling con attention mask:

def mean_pooling_with_mask(hidden_states, attention_mask):
    """
    Mean pooling ignorando tokens de padding
    
    Args:
        hidden_states: (batch_size, seq_len, hidden_size)
        attention_mask: (batch_size, seq_len) - 1 = token real, 0 = padding
    
    Returns:
        Embeddings: (batch_size, hidden_size)
    """
    # Expandir attention mask para multiplicación
    attention_mask_expanded = np.expand_dims(attention_mask, axis=-1)
    
    # Multiplicar hidden states por mask (anula padding)
    masked_hidden_states = hidden_states * attention_mask_expanded
    
    # Sumar (solo tokens reales)
    sum_hidden = np.sum(masked_hidden_states, axis=1)
    
    # Dividir por cantidad de tokens reales
    sum_mask = np.sum(attention_mask_expanded, axis=1)
    sum_mask = np.clip(sum_mask, a_min=1e-9, a_max=None)  # Evitar división por 0
    
    # Mean
    embeddings = sum_hidden / sum_mask
    
    return embeddings

# Ejemplo
hidden_states = np.array([
    [[0.8, 0.1, 0.3], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],  # Batch 1: 1 token real
    [[0.9, 0.3, 0.2], [0.7, 0.2, 0.4], [0.6, 0.4, 0.1]]   # Batch 2: 3 tokens reales
])

attention_mask = np.array([
    [1, 0, 0],  # Batch 1: Solo primer token es real
    [1, 1, 1]   # Batch 2: Todos los tokens son reales
])

embeddings = mean_pooling_with_mask(hidden_states, attention_mask)
print("Embeddings:")
print(embeddings)

Output:

Embeddings:
[[0.8  0.1  0.3 ]  ← Batch 1: Solo usó token 1
 [0.73 0.3  0.23]]  ← Batch 2: Promedio de 3 tokens

Comparación empírica: Mean vs CLS

Experimento: ¿Cuál es mejor?

Setup: Evaluar ambos pooling strategies en tarea de similaridad.

from sentence_transformers import SentenceTransformer

# Cargar modelo SBERT (usa mean pooling por defecto)
model = SentenceTransformer('all-MiniLM-L6-v2')

# Pares de frases similares
pairs = [
    ("El gato duerme", "El felino descansa"),
    ("Python es popular", "Python es muy usado"),
    ("Me gusta la pizza", "Adoro la pizza")
]

# Generar embeddings con mean pooling (default)
for text_a, text_b in pairs:
    emb_a = model.encode(text_a)
    emb_b = model.encode(text_b)
    
    sim = np.dot(emb_a, emb_b) / (np.linalg.norm(emb_a) * np.linalg.norm(emb_b))
    print(f"'{text_a}' vs '{text_b}': {sim:.4f}")

Output con mean pooling:

'El gato duerme' vs 'El felino descansa': 0.78
'Python es popular' vs 'Python es muy usado': 0.85
'Me gusta la pizza' vs 'Adoro la pizza': 0.82

Mean pooling captura paráfrasis efectivamente.

Ejercicios

Ejercicio 1: Implementar mean pooling

Implementa mean pooling para estos hidden states:

import numpy as np

hidden_states = np.array([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0]
])

# Calcula mean pooling

def mean_pooling(hidden_states):
    return np.mean(hidden_states, axis=0)

embedding = mean_pooling(hidden_states)
print(f"Embedding: {embedding}")

Embedding: [4. 5. 6.]

Ejercicio 2: CLS vs Mean pooling

Calcula CLS y mean pooling para los mismos hidden states. ¿Son similares?

hidden_states = np.array([
    [0.5, 0.6, 0.7],  # [CLS]
    [0.8, 0.1, 0.3],
    [0.9, 0.3, 0.2]
])

# Calcula ambos y compara

# CLS pooling
cls_emb = hidden_states[0]

# Mean pooling
mean_emb = np.mean(hidden_states, axis=0)

print(f"CLS:  {cls_emb}")
print(f"Mean: {mean_emb}")

# Comparar
similarity = np.dot(cls_emb, mean_emb) / (np.linalg.norm(cls_emb) * np.linalg.norm(mean_emb))
print(f"\nSimilarity: {similarity:.4f}")

CLS:  [0.5 0.6 0.7]
Mean: [0.73 0.33 0.4]

Similarity: 0.9850

Ejercicio 3: Max pooling

Implementa max pooling:

hidden_states = np.array([
    [0.8, 0.1, 0.3],
    [0.9, 0.3, 0.2],
    [0.7, 0.2, 0.4]
])

# Implementa max pooling

def max_pooling(hidden_states):
    return np.max(hidden_states, axis=0)

embedding = max_pooling(hidden_states)
print(f"Embedding: {embedding}")

Embedding: [0.9 0.3 0.4]

Resumen

Qué aprendiste:

• ✅ Pooling: Reducir múltiples tokens → 1 embedding
• ✅ Mean pooling: Promedio (todos los tokens importan)
• ✅ CLS pooling: Token especial [CLS] (BERT-style)
• ✅ Max pooling: Máximo por dimensión (raro)
• ✅ Attention mask: Ignorar padding en batch processing
• ✅ Best practice: Mean pooling para sentence embeddings

Conceptos clave:

1. Mean pooling domina en sentence embeddings modernos
2. CLS pooling requiere token especial (BERT)
3. Attention mask es crítico para batch processing

Recursos adicionales

1. Sentence-BERT Paper - Justifica mean pooling
2. Pooling Strategies Comparison - Experimentos
3. BERT Pooling - Tutorial CLS
4. Attention Mask Explained - HuggingFace
5. Mean Pooling Implementation - SBERT código

En la siguiente cápsula

Cápsula 06: Normalization

Aprenderás:

• L2 normalization (código numpy)
• Por qué normalizar embeddings
• Dot product vs cosine similarity
• Cuándo normalizar (depende del modelo)
• Implementación práctica

De pooling a normalización.

Módulo 2 - Embeddings Deep Dive Guide Agregando tokens: de muchos vectores a uno