1. INTRODUCCIÓN Y ESTADO DEL ARTE

Para realizar tareas de procesamiento de lenguaje natural (NLP), los grandes modelos
de lenguaje (LLMs) deben convertir la información de entrada, texto, en vectores de alta
dimensión que capturan información sintáctica y semántica. El aprendizaje de estas representaciones o embeddings (representation learning) permite al modelo generalizar su conocimiento a nuevas tareas (transfer learning).

Cada palabra o secuencia de palabras es transformada, en distintas etapas del modelo,
en vectores que resumen su significado en contexto. Entender cómo se construyen estas representaciones y cómo pueden manipularse es clave para interpretar y mejorar el rendimiento de los modelos. [1]

En esta entrada exploraremos paso a paso cómo se procesan los textos en un modelo
preentrenado, desde la tokenización hasta las primeras capas del modelo. A partir de ahí, experimentaremos con transformaciones del espacio latente a otros de menor dimensionalidad mediante autoencoders —como PCA o VAEs— para obtener representaciones más compactas. Finalmente, mediremos si es posible adaptar el modelo a trabajar con esta nueva representación sin perder demasiada precisión. [2]

2. MATERIALES

El objetivo principal de esta experimentación es observar cómo se representan las secuencias de texto en un modelo preentrenado y cómo afectan distintas transformaciones de estas representaciones al rendimiento en tareas de clasificación.

DATASET Y TAREA

Para este estudio utilizamos el conjunto de datos TREC (Question Classification), una
colección clásica en procesamiento del lenguaje natural compuesta por preguntas cortas etiquetadas según su tipo. Las categorías incluyen seis clases generales: ABBR
(abreviaciones), DESC (descripciones), ENTY (entidades), HUM (personas), LOC
(localizaciones) y NUM (números). El objetivo del modelo es predecir correctamente la
categoría de una pregunta a partir de su texto. El conjunto incluye aproximadamente
5.000 ejemplos de entrenamiento y 500 de test [3]. Es una tarea sencilla, que nos facilita probar diversas aproximaciones, y que evaluaremos con la métrica F-score.

MODELO

Se ha empleado el modelo distilbert-base-uncased, una versión ligera de BERT que
mantiene su arquitectura principal pero con menos capas [4]. Esto facilita la inspección
de las representaciones internas y reduce el coste computacional.

3. DEL LENGUAJE NATURAL AL ESPACIO LATENTE: PROCESADO PASO A PASO EN UN MODELO DE CLASIFICACIÓN

El primer paso en el recorrido de una secuencia textual por un modelo como DistilBERT es la tokenización. Esto transforma el texto en fragmentos más pequeños (tokens), que luego son convertidos en números mediante un vocabulario entrenado.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

texto = "Where is the Eiffel Tower?"
tokens = tokenizer.tokenize(texto)

print(tokens)

['where', 'is', 'the', 'e', '##iff', '##el', 'tower', '?']

Cada token es una subpalabra, generada según reglas específicas del vocabulario del modelo (WordPiece). Estos tokens se traducen a identificadores numéricos únicos que corresponden a posiciones concretas en el vocabulario de DistilBERT.

token_ids = tokenizer.convert_tokens_to_ids(tokens)
print(token_ids)

[2073, 2003, 1996, 1041, 13355, 2884, 3578, 1029]

En una aplicación real, el modelo espera una entrada que incluya tokens especiales como [CLS] (al inicio) y [SEP] (al final). También se aplica padding y truncamiento para ajustar las secuencias a una longitud fija.

encoded = tokenizer(
    texto,
    padding="max_length",
    truncation=True,
    max_length=12,
    return_tensors="pt"
)

print(encoded["input_ids"])

tensor([[  101,  2073,  2003,  1996,  1041, 13355,  2884,  3578,  1029,   102,
             0,     0]])

Donde:

Una vez que los tokens están representados como IDs, se pasan por la capa de embeddings del modelo. Cada ID es convertido en un vector denso de dimensión fija (en DistilBERT, 768).

from transformers import AutoModel
import torch

model = AutoModel.from_pretrained("distilbert-base-uncased")
with torch.no_grad():
    embedding_output = model.embeddings(encoded["input_ids"])

print(embedding_output.shape)

torch.Size([1, 12, 768])

Esto indica que hay una secuencia de 12 tokens, cada uno representado como un vector de 768 dimensiones. Estos vectores combinan embeddings de tokens, embeddings de posición y embeddings de tipo de segmento (para tareas específicas, por ejemplo la tarea de clasificación a continuación).

Una vez hemos analizado el proceso de transformación del texto hasta su representación en el espacio latente, vamos a explorar este espacio y ver qué transformaciones podemos hacer sobre el mismo.

4. AJUSTE DEL MODELO A EMBEDDINGS REDUCIDOS

El objetivo es que el modelo pueda trabajar directamente con representaciones comprimidas [5]. Exploramos dos enfoques para lograrlo:

• Reducción de dimensionalidad con PCA
• Compresión mediante un autoencoder variacional (VAE)

Ambos se implementan de dos formas:

• Dentro del modelo, añadiendo una capa de expansión al principio que reconstituye los embeddings comprimidos a su dimensión original (768).
• Fuera del modelo, reconstruyendo los embeddings antes de pasarlos a DistilBERT.

ENFOQUE 1: REDUCCIÓN DE DIMENSIONALIDAD CON PCA

REDUCCIÓN DE DIMENSIONALIDAD

Como primer enfoque, utilizamos Análisis de Componentes Principales (PCA) para reducir los embeddings generados por la capa de entrada de DistilBERT. Estos embeddings tienen una dimensión original de 768. En nuestro experimento, los comprimimos a 512 dimensiones.

pca = PCA(n_components=512)
reduced = pca.fit_transform(flat.numpy())

Primero tokenizamos el texto y extraemos los embeddings de entrada del modelo, que tienen forma [batch_size, seq_len, 768]. Luego aplicamos PCA, reduciendo la dimensión de 768 a 512. Para ello, reordenamos los datos como una matriz de vectores [batch_size × seq_len, 768], aplicamos la transformación, y reconstruimos el tensor original con la nueva dimensión reducida: [batch_size, seq_len, 512].

ADAPTACIÓN DEL MODELO

El siguiente reto fue adaptar el modelo para que pudiera trabajar con estos nuevos embeddings reducidos. Para ello, implementamos una clase que incorpora:

• Una capa lineal de proyección: transforma los vectores de 512 dimensiones de vuelta a los 768 requeridos por DistilBERT.
• El encoder original de DistilBERT, reutilizado tal cual.
• Una capa de clasificación final para predecir la categoría de cada pregunta del dataset TREC.

class FineTunedDistilBERT(nn.Module):
    def __init__(self, model, num_labels):
        super(FineTunedDistilBERT, self).__init__()
        self.num_labels = num_labels
        self.distilbert = model.distilbert  # Get the base model
        self.embedding_projector = nn.Linear(num_components, 768)  # Project custom embeddings
        self.classifier = nn.Linear(768, num_labels)    # Your classifier head
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, inputs_embeds, labels=None, attention_mask=None):
        # Project the reduced embedding to 768
        projected = self.embedding_projector(inputs_embeds)

        # Pass through DistilBERT encoder
        outputs = self.distilbert(inputs_embeds=projected, attention_mask=attention_mask)
        hidden_state = outputs.last_hidden_state  # shape: (batch_size, seq_len, hidden_size)

        # Use the [CLS]-like token (first token) for classification
        cls_representation = hidden_state[:, 0, :]  # (batch_size, hidden_size)
        logits = self.classifier(cls_representation)

        loss = None
        if labels is not None:
            loss = self.loss_fn(logits, labels)

        return {"loss": loss, "logits": logits}

RESULTADOS

Aunque la implementación era funcional, el modelo no aprendía y los resultados fueron muy malos. Especulamos sobre algunas posibles razones:

• El modelo original está optimizado para trabajar con una distribución específica de entradas; al modificar la entrada, rompemos esa alineación.
• La capa de expansión al inicio del modelo no recupera adecuadamente la información textual.

ENFOQUE 2: RECONSTRUCCIÓN CON PCA

En base a estos resultados, decidimos cambiar el enfoque: aplicamos reducción de dimensionalidad con PCA, pero esta vez reconstruyendo los embeddings antes de ingresarlos al modelo. Es decir:

• Obtenemos los embeddings de entrada (768 dim).
• Los comprimimos a un espacio latente de menor dimensión (ej. 512).
• Reconstruimos los embeddings originales con la transformación inversa del PCA.
• Pasamos estos vectores reconstruidos al modelo, sin modificar su arquitectura.

Realmente, al experimentar con esta arquitectura lo que estamos evaluando es cuánta información se pierde al codificar y decodificar la información con un autoencoder, en este caso PCA.

recovered = compressor.decode(compressor.encode(embeddings), embeddings.shape)

A diferencia del enfoque anterior, este método sí ofrece buenos resultados (que presentaremos más adelante), especialmente cuando el espacio latente no es demasiado reducido.

ENFOQUE 3: VAE PARA CODIFICACIÓN Y DECODIFICACIÓN

VAE INTEGRADO CON EL MODELO

En esta aproximación, de nuevo partimos de la arquitectura de DistilBERT, a la que añadimos un módulo de compresión VAE entrenado de forma conjunta. La arquitectura utilizada es simple y simétrica: un codificador con una capa oculta de 512 unidades que proyecta a un espacio latente de dimensión configurable, y un decodificador con la misma estructura en orden inverso. Se usa ReLU como activación intermedia. La pérdida combina MSE (error de reconstrucción) y divergencia KL (forzar distribución normal en el espacio latente).

class VAECompressor(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(VAECompressor, self).__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        
        # Encoder: input -> hidden -> latent
        self.fc1 = nn.Linear(input_dim, 512)
        self.fc2_mu = nn.Linear(512, latent_dim)
        self.fc2_logvar = nn.Linear(512, latent_dim)
        
        # Decoder: latent -> hidden -> output
        self.fc3 = nn.Linear(latent_dim, 512)
        self.fc4 = nn.Linear(512, input_dim)
        
    def encode(self, x):
        h1 = torch.relu(self.fc1(x))
        mu = self.fc2_mu(h1)
        logvar = self.fc2_logvar(h1)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        h3 = torch.relu(self.fc3(z))
        return self.fc4(h3)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        reconstructed = self.decode(z)
        return reconstructed, mu, logvar

    def loss_function(self, recon_x, x, mu, logvar):
        # Binary cross entropy loss for reconstruction
        BCE = nn.functional.mse_loss(recon_x, x, reduction='sum')
        
        # KL divergence loss
        KL = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        # Total loss is reconstruction + KL
        return BCE + KL

El modelo optimiza simultáneamente la pérdida de clasificación y la pérdida del VAE, definiendo como función de pérdida la suma de ambos.

class FineTunedDistilBERTVAE(nn.Module):
    def __init__(self, model, vae_compressor, num_labels):
        super(FineTunedDistilBERTVAE, self).__init__()
        self.num_labels = num_labels
        self.distilbert = model.distilbert
        self.classifier = nn.Linear(768, num_labels)
        self.loss_fn = nn.CrossEntropyLoss()
        self.vae_compressor = vae_compressor

    def forward(self, inputs_embeds, labels=None, attention_mask=None):
        # Compresión y reconstrucción de embeddings
        reconstructed, mu, logvar = self.vae_compressor(inputs_embeds)
        
        # Paso por DistilBERT
        outputs = self.distilbert(inputs_embeds=reconstructed, attention_mask=attention_mask)
        cls_representation = outputs.last_hidden_state[:, 0, :]
        logits = self.classifier(cls_representation)

        # Cálculo de pérdidas
        loss = self.loss_fn(logits, labels) if labels is not None else None
        vae_loss = self.vae_compressor.loss_function(reconstructed, inputs_embeds, mu, logvar)
        total_loss = loss + vae_loss if loss is not None else vae_loss

        return {"loss": total_loss, "logits": logits}

Sin embargo, los resultados obtenidos fueron muy malos. Es posible que la combinación de ambos entrenamientos dificultara la optimización, ya que ambos objetivos (compresión y clasificación) competían por los recursos del modelo. También probamos a ajustar la función de pérdida mediante un parámetro beta, que ponderaba la pérdida de clasificación y del VAE, aunque de nuevo sin éxito.

total_loss = loss + beta * vae_loss

ENTRENAMIENTO SEPARADO

Finalmente decidimos separar los entrenamientos: primero entrenamos el VAE para realizar la compresión de los embeddings, y luego utilizamos los embeddings reconstruidos en el modelo de clasificación. El código completo está disponible en este enlace.

Este punto fue el más interesante, puesto que experimentamos con la arquitectura del VAE de diversas formas. Para medir su rendimiento (calidad de la reconstrucción) implementamos una función que devuelve el error MSE de reconstrucción sobre un subconjunto de datos, y de esta forma ir ajustando el VAE. Con ello finalmente implementamos los siguientes cambios que mejoraban el rendimiento, entre otros:

• Nº y tamaño de capas: añadir una segunda capa y ajustar el tamaño por capa. La arquitectura definitiva fue (con dimensión de entrada 768): 2 capas de 640 y 512 neuronas cada una para codificación, y decodificación simétrica.
• Función de pérdida: es mejor ponderar el componente KL al mínimo (beta = 0.001) para priorizar la reconstrucción.
• Duración del entrenamiento: aumentar el nº de épocas hasta 30.
• Función de activación: GELU (en vez de RELU).
• Estabilidad del entrenamiento:
  • Acotar la magnitud del gradiente para evitar “explosiones del gradiente”.
  • Establecer una condición de parada en caso de empeoramiento (3 épocas de paciencia).
  • Acotar el rango del logaritmo de la varianza (logvar) al computar KL para evitar error.

Esto son sólo algunos de los detalles que tuvimos que considerar, y que darían para una entrada completa de blog 🙂

Con esta configuración, conseguimos un entrenamiento estable, cuya evolución podemos visualizar para las distintas dimensiones del espacio latente. El valor de la función de pérdida es muy alto, ya que utilizar la suma MSE, en vez de la media, mejoraba la estabilidad. Es probable que hubiera margen para mejorar el rendimiento aumentando el número de épocas, aunque con su consiguiente coste computacional.

A continuación, presentamos los resultados obtenidos en cuanto a la calidad de la reconstrucción. Como cabría esperar, a mayor dimensionalidad, menor pérdida de información.

Y, finalmente, podemos evaluar el rendimiento en la tarea de clasificación, comparando los resultados obtenidos con PCA frente a VAE. Previsiblemente, nuestro VAE proyecta de forma más eficiente la información, pudiendo alcanzar una reducción mayor de dimensionalidad con un mejor rendimiento.

5. CONCLUSIONES

A lo largo del trabajo, hemos explorado el espacio latente en los LLMs, desde el proceso
para obtener representaciones numéricas de información textual hasta la exploración y
transformaciones de dichos espacios.

Hemos intentado aplicar diversas estrategias para reducir la dimensionalidad de estas
representaciones textuales sin perder rendimiento. Probamos dos enfoques principales: el uso de PCA y VAEs, tanto integrados directamente en el modelo como aplicados de manera separada, obteniendo las siguientes conclusiones:

• Modificar la arquitectura de una red neuronal sin perder rendimiento no es trivial, así como diseñar un VAE funcional: hay que tener en cuenta y elegir muchos detalles (tamaño y nº de capas, función de pérdida, etc.), y no está garantizado que funcione.
• Para tareas específicas, es posible obtener buenos resultados tras reducir la representación a un espacio de menor dimensionalidad, lo que sugiere que se pueden desarrollar modelos de lenguaje mucho más pequeños y eficientes para aplicaciones concretas.
• La reducción significativa de la dimensionalidad de las representaciones textuales es viable, siempre que no se altere la arquitectura original del modelo.

Realizado por Adrián Maldonado Robles, Paula Martín Fernández y Elena Moyano González

6. REFERENCIAS

• [1] Hugging Face. Natural Language Processing and Large Language Models. https://huggingface.co/learn/llm-course/chapter1/2

• [2] Alammar, J. The Illustrated BERT, ELMo, and co. (How NLP Cracked Transfer
  Learning). https://jalammar.github.io/illustrated-bert/

• [3] Hugging Face. Dataset Card for “trec”. https://huggingface.co/datasets/CogComp/trec

• [4] Hugging Face. DistilBERT base model (uncased). https://huggingface.co/distilbert/distilbert-base-uncased

• [5] EmbedLLM: Learning Compact Representations of Large Language Models. https://openreview.net/forum?id=Fs9EabmQrJ