Aprendizaje de Representaciones Matryoshka Explicado: El Método Detrás de los Embeddings de Texto Eficientes de OpenAI
Desarrollar modelos de aprendizaje automático de la vida real siempre implica un equilibrio entre costo y rendimiento. Por ejemplo, cuanto más grande sea el modelo y el conjunto de datos que usemos durante el entrenamiento, más capaz será finalmente nuestro modelo entrenado. Sin embargo, esta mayor capacidad conlleva un tiempo más prolongado y mayores costos computacionales, ya que el entrenamiento requiere duraciones significativamente más largas. El mismo principio se aplica durante la inferencia, donde los modelos más grandes tienden a generar representaciones de características más grandes que requieren más memoria para su almacenamiento.
Dado que las necesidades durante el entrenamiento y la inferencia de aprendizaje automático pueden variar significativamente de un caso de uso a otro, es esencial contar con un método que nos permita intercambiar una pequeña parte del rendimiento de un modelo a cambio de costos reducidos. Aquí es donde entra en juego un método como Matryoshka Representation Learning (MRL) . Por ejemplo, el modelo text-embedding-3-small de OpenAI utiliza MRL para permitir a los desarrolladores acortar embeddings mientras mantienen sus propiedades principales de representación de conceptos. Al permitir el ajuste de las dimensiones de los embeddings, MRL ayuda a lograr el equilibrio perfecto entre eficiencia de costos y rendimiento del modelo.
Este artículo explorará cómo funciona MRL, su implementación y cómo permite modelos de aprendizaje automático escalables y eficientes. Comencemos con la motivación detrás de MRL.
La motivación detrás de Matryoshka Representation Learning (MRL)
El equilibrio entre costo y rendimiento siempre es una consideración antes de desarrollar o usar cualquier modelo de aprendizaje automático.
Durante el entrenamiento, los parámetros de un modelo y la cantidad de datos de entrenamiento influyen directamente en su rendimiento final. Los modelos más grandes y los conjuntos de datos de entrenamiento tienden a dar como resultado modelos más capaces. Sin embargo, el costo computacional para entrenar estos modelos también es mayor.
Mientras tanto, durante la inferencia, los modelos más grandes producen representaciones de características más grandes. Este aumento en las representaciones de características significa que necesitamos más memoria para almacenarlas.
Usemos BERT como ejemplo en el contexto de la recuperación de información. BERT es uno de los primeros modelos de aprendizaje profundo basados en Transformer que logró un rendimiento de vanguardia en varios conjuntos de datos de referencia en clasificación de texto, reconocimiento de entidades nombradas (NER), respuesta a preguntas (QnA) y más.
Figura: Flujo de trabajo de generación de embeddings usando el modelo base de BERT.
En pocas palabras, BERT toma una oración o una palabra como entrada y la transforma en un embedding de tamaño fijo. Este embedding captura el significado semántico de la entrada original, y su tamaño depende de la variante específica del modelo utilizada.
BERT tiene dos variantes principales: el modelo base y el modelo grande. El modelo base produce un embedding de 768 dimensiones, mientras que el modelo grande genera un embedding de 1024 dimensiones.
Supongamos que queremos almacenar 10 millones de embeddings en una base de datos vectorial para un caso de uso de recuperación de información. Si usamos el modelo base de BERT en formato FP32, necesitaremos aproximadamente 768×10M×4=30.72 GB de memoria para almacenar estos embeddings. En cambio, usar el modelo grande de BERT en el mismo escenario requeriría alrededor de 40.9 GB. Además, un tamaño de embedding más grande puede ralentizar el proceso de recuperación al aumentar la complejidad computacional de las búsquedas de similitud. Sin embargo, la principal ventaja de usar embeddings más grandes es su mayor capacidad para recuperar información altamente relevante en comparación con embeddings más pequeños.
Una solución ideal para estos problemas en la recuperación de información es dividir todo el proceso en dos partes: preselección y reranking. La preselección se refiere a recuperar un conjunto inicial de documentos candidatos de una vasta colección en nuestra base de datos vectorial como Milvus. El reranking luego toma estos candidatos preseleccionados y los reordena para maximizar la relevancia de los resultados finales.
Figura- Flujo de trabajo de preselección y reranking
Figura: Flujo de trabajo de preselección y reranking.
Para realizar la preselección, el algoritmo calcula la similitud entre los embeddings de consulta y una gran colección de embeddings almacenados en la base de datos vectorial. Por lo tanto, usar tamaños de embedding más pequeños puede hacer que este cálculo sea más eficiente y rápido. Por el contrario, el reranking enfatiza la precisión, lo que significa que nos beneficiaríamos de tamaños de embedding más grandes.
El problema con este método es que, una vez que seleccionamos un modelo en particular, el tamaño de los embeddings o las representaciones de características que produce queda fijo. Esta falta de flexibilidad limita nuestra capacidad para usar embeddings de diferentes tamaños desde un solo modelo.
Figura- Flujo de trabajo de un modelo que genera un embedding de tamaño fijo..png
Figura: Flujo de trabajo de un modelo que genera un embedding de tamaño fijo.
El enfoque de Matryoshka Representation Learning (MRL) ofrece una solución intrigante a este problema: los modelos entrenados con este método pueden producir embeddings de varios tamaños. Exploraremos el mecanismo de este enfoque en detalle en la siguiente sección.
¿Qué es Matryoshka Representation Learning (MRL)?
Matryoshka Representation Learning (MRL) es un método para entrenar redes neuronales con el fin de producir representaciones multiescala dentro de un solo modelo. Inspirado en las muñecas matrioshka, donde las muñecas más pequeñas encajan dentro de las más grandes, MRL permite que un modelo genere representaciones de tamaños variables (de gruesas a finas) usando una sola pasada hacia adelante. Este enfoque permite que los modelos de aprendizaje profundo comprendan mejor las relaciones complejas y los matices en los datos. MRL es particularmente eficaz para tareas como búsqueda semántica, recuperación de información, procesamiento multilingüe y cualquier aplicación que requiera representaciones matizadas de datos en diferentes niveles de abstracción.
Visualización de embeddings Matryoshka con múltiples capas de detalle
Figura: Visualización de embeddings Matryoshka con múltiples capas de detalle
Los modelos de embedding populares que han adoptado el enfoque de Matryoshka Representation Learning (MRL) incluyen text-embedding-3-large de OpenAI, nomic-embed-text-v1 de Nomic y gte-multilingual-base de Alibaba.
Cómo funciona el enfoque MRL
El enfoque MRL nos permite extraer representaciones de características de cualquier modelo de aprendizaje automático en varios tamaños. Por ejemplo, en lugar de usar las 1024 dimensiones originales de un embedding, podemos utilizar las primeras 16, 32, 64, 128 o 256 dimensiones (o cualquier dimensión que elijamos). La clave de esta capacidad reside en la forma en que se entrena el modelo al implementar el enfoque MRL.
Figura- Entrenamiento de la función de pérdida MRL y su caso de uso durante la inferencia
Figura: Entrenamiento de la función de pérdida MRL y su caso de uso durante la inferencia. Fuente.
Durante el entrenamiento de modelos con MRL, optimizamos no solo una función de pérdida, como es típico en el entrenamiento estándar de modelos, sino varias. Cada función de pérdida tiene como objetivo optimizar la representación de características para una dimensión específica que definimos de antemano. Si establecemos cinco dimensiones de características diferentes, necesitaremos optimizar cinco funciones de pérdida distintas durante el entrenamiento.
Supongamos que queremos entrenar un modelo BERT large usando el enfoque MRL. Como se mencionó en la sección anterior, este modelo genera un embedding de tamaño 1024. Si establecemos las dimensiones de características en 32, 64, 128, 256 y 1024, optimizaremos cinco funciones de pérdida diferentes durante el entrenamiento para que el modelo aprenda a generar embeddings optimizados en esas dimensiones.
La optimización a través de estas diversas dimensiones es sencilla: MRL descompone la función de pérdida general en la suma de las pérdidas para cada dimensión individual. Usando nuestro ejemplo anterior, esto se puede expresar como:
El proceso de optimización para un modelo entrenado con MRL se puede representar de manera más formal de la siguiente forma:
Figura: Ecuación de optimización del enfoque MRL. Fuente.
Como puedes ver, la definición de la función de pérdida anterior es bastante general. Por lo tanto, el enfoque MRL se puede aplicar a casi cualquier modelo y es completamente independiente de la arquitectura del modelo. También podemos ajustar modelos preentrenados como BERT o cualquier otro modelo basado en transformers para generar embeddings de varios tamaños. Continuando con nuestro escenario anterior, ahora podemos generar embeddings de 1024 dimensiones a partir de nuestro modelo BERT large entrenado y embeddings de 32, 64, 128 y 256 dimensiones.
En general, después de entrenar un modelo con MRL, las dimensiones iniciales de las características generadas por el modelo entrenado contienen información más significativa que las dimensiones posteriores. Las primeras dimensiones contienen detalles de alto nivel, mientras que las dimensiones posteriores se centran en información más granular, imitando cómo está estructurada una muñeca Matryoshka.
Sin embargo, esto no significa que los embeddings más cortos sean siempre simplemente versiones truncadas de los más largos. Los valores de cada elemento en los embeddings más cortos podrían diferir de los de los embeddings más largos debido al factor de escala aplicado a cada dimensión de característica durante el entrenamiento. Sin embargo, si establecemos que el factor de escala en cada dimensión sea igual, los valores de cada elemento entre los embeddings más cortos y más largos podrían parecer idénticos.
Resultados experimentales de MRL
El enfoque MRL se ha evaluado en modelos de aprendizaje automático a través de diferentes modalidades, incluidas texto, visión y visión-texto. Los modelos ResNet50 y ViT representan modelos basados en visión, BERT representa modelos basados en texto, y ALIGN representa la combinación de visión y texto. Estos modelos se han evaluado principalmente para dos casos de uso comunes: clasificación y recuperación. Hablemos primero de la clasificación.
Al comparar el rendimiento del modelo ResNet50 entrenado en el conjunto de datos ImageNet-1K con un ResNet50 estándar entrenado de forma independiente, el modelo MRL logra una precisión top-1 comparable en varios tamaños de representación de características.
Para evaluar aún más la utilidad de las representaciones de características para tareas posteriores, también se midió la precisión de 1-vecino más cercano (1-NN) para cada tamaño de representación de características. La configuración para calcular 1-NN es la siguiente: dadas 1.3K muestras de imágenes en la base de datos, la tarea consiste en encontrar el vecino más cercano para cada una de las 50K consultas. Con esta configuración, el ResNet50 entrenado con MRL es hasta un 2% más preciso en cada tamaño de representación de características en comparación con su contraparte de características fijas.
Figura- Precisión Top-1 y precisión 1-NN de ResNet50 en ImageNet-1K
Figura: Precisión Top-1 y precisión 1-NN de ResNet50 en ImageNet-1K. Fuente.
Mientras tanto, el rendimiento del modelo ViT entrenado con MRL en el conjunto de datos JFT-300M también es muy competitivo en todos los tamaños de representación. Su precisión 1-NN es comparable a la de ViT entrenado con representaciones de características de tamaño fijo. Como se muestra en la figura de abajo, el rendimiento del modelo MRL también es mejor que el de su contraparte de tamaño fijo en representaciones de características más bajas, en parte porque se seleccionan características aleatorias del modelo de tamaño fijo para representar dimensiones más bajas. Se observa una tendencia similar con el modelo ALIGN entrenado mediante el enfoque MRL; su rendimiento coincide con el del modelo ALIGN entrenado con representaciones de tamaño fijo.
Figura- Precisión 1-NN de ViT y ALIGN en ImageNet-1K
Figura: Precisión 1-NN de ViT y ALIGN en ImageNet-1K. Fuente.
La principal ventaja de usar el enfoque MRL es su flexibilidad para cambiar entre diferentes tamaños de representaciones de características desde el mismo modelo. Para aprovechar plenamente esta fortaleza, también se realizaron pruebas de clasificación adaptativa usando ResNet50 en ImageNet-1K.
Esta configuración implica aprender la probabilidad softmax máxima para determinar transiciones de representaciones de características más pequeñas a otras más grandes. Los resultados de las pruebas muestran que el modelo MRL con una representación de características de 37 dimensiones es tan preciso como el modelo ResNet50 entrenado con una representación de características fija de 512 dimensiones, siendo solo un 0.8% menos preciso que el modelo de 2048 dimensiones.
Figura- Precisión Top 1 del modelo ResNet50 en clasificación adaptativa (izquierda) y mAP de ResNet50 en caso de uso de recuperación de imágenes (derecha)
Figura: Precisión Top 1 del modelo ResNet50 en clasificación adaptativa (izquierda) y mAP de ResNet50 en caso de uso de recuperación de imágenes (derecha). Fuente.
A continuación, analicemos la recuperación. Como quizá sepas, el objetivo de la recuperación es encontrar un conjunto de contextos similares para una consulta determinada. Por lo tanto, se probó la calidad de recuperación entre ResNet50 entrenado con MRL y representaciones de características de tamaño fijo usando la precisión media promedio (mAP). El modelo MRL muestra una mejora en mAP de hasta un 3% en comparación con su contraparte de tamaño fijo en cada tamaño de representación de características, como se ilustra en la visualización anterior.
Dado que podemos utilizar varios tamaños de representación de características con MRL, también resulta particularmente interesante explorar cómo podemos acelerar el proceso de recuperación manteniendo la precisión con MRL. Podemos probarlo en un escenario llamado recuperación adaptativa.
Como se mencionó anteriormente, hay dos etapas principales en un proceso de recuperación: preselección y reordenamiento. La configuración de la recuperación adaptativa implicó usar una representación de características de 16 dimensiones para preseleccionar un conjunto de 200 candidatos, mientras que se usa una representación de características de 2048 dimensiones para el reordenamiento. Este enfoque se compara con otro método que usa una representación de características de 2048 dimensiones tanto para la preselección como para el reordenamiento.
En ImageNet-1K, usar el modelo MRL puede acelerar teóricamente el proceso de recuperación hasta 128 veces. Nos referimos a esto como teórico porque, en aplicaciones del mundo real, rara vez usamos un algoritmo ingenuo de vecinos más cercanos para tareas de recuperación. En su lugar, se suelen implementar algoritmos de vecinos más cercanos aproximados (ANN) como FAISS, ANNOY o HNSW para mejorar la velocidad de recuperación. En este caso, la configuración que usa el modelo MRL logra una aceleración de 14 veces en comparación con los procesos de recuperación que usan el algoritmo HNSW en hardware idéntico.
Se observó un resultado similar con el conjunto de datos ImageNet-4K, donde los autores del artículo usaron un modelo MRL con una representación de características de 64 dimensiones para la preselección y una representación de 2048 dimensiones para la reclasificación. Los resultados indican una aceleración teórica de 32 veces y una aceleración en el mundo real de seis veces.
La compensación entre mAP@10 vs MFLOPs/Query para Recuperación Adaptativa (AR) en ImageNet-1K (izquierda) e ImageNet-4K (derecha).
La compensación entre mAP@10 vs MFLOPs/Query para Recuperación Adaptativa (AR) en ImageNet-1K (izquierda) e ImageNet-4K (derecha). Fuente.
Implementación del enfoque MRL
En esta sección, veremos la implementación simple del enfoque MRL. En particular, usaremos un modelo base MPNet entrenado con el enfoque MRL en el conjunto de datos NLI para generar embeddings en diferentes tamaños. Como quizá ya sepas, el modelo base MPNet original solo puede generar un embedding de 768 dimensiones. Sin embargo, este modelo MPNet entrenado puede generar embeddings en 768, 512, 256, 128 y 64 dimensiones, donde el peso de cada dimensión durante el entrenamiento se establece en 1.
Podemos usar este modelo con la ayuda de la biblioteca SentenceTransformers con el siguiente código.
from sentence_transformers import SentenceTransformer
matryoshka_dim_short = 64
matryoshka_dim_long = 768
text = ["The weather is so nice!"]
short_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_short).encode(text)
long_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_long).encode(text)
print(f"Shape: {short_embedding.shape, long_embedding.shape}")
print(short_embedding[0][0:10])
print(long_embedding[0][0:10])
"""
Output:
Shape: ((1, 64), (1, 768))
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
"""
En el código anterior, usamos el modelo base MPNet entrenado para generar embeddings en dos tamaños diferentes: 64 y 768. Dado que el peso o factor de escala de cada dimensión se establece en 1, los dos embeddings tienen elementos idénticos, como puedes ver en los primeros 10 elementos de los dos embeddings anteriores.
Como el valor de cada elemento es idéntico, la similitud entre los dos embeddings sería 1.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
Si deseas entrenar tu propio modelo con el enfoque MRL, puedes consultar el repositorio oficial de GitHub de este método.
Conclusión
El enfoque MRL introduce una solución para equilibrar el compromiso entre costo y rendimiento en el aprendizaje automático. Al permitir que cualquier modelo de aprendizaje automático produzca representaciones de características de distintos tamaños, MRL proporciona la flexibilidad para optimizar ya sea la velocidad o la precisión según nuestro caso de uso y recursos. Hemos visto que esta adaptabilidad es valiosa en aplicaciones como la clasificación y la recuperación, donde cambiar entre representaciones más pequeñas y más grandes puede mejorar significativamente la eficiencia sin sacrificar enormemente el rendimiento.
Los resultados experimentales demuestran que MRL iguala y a menudo supera la precisión de los modelos tradicionales de tamaño fijo en múltiples dominios, incluidos texto, visión y tareas multimodales. Además, la compatibilidad del enfoque con arquitecturas existentes como ResNet50, ViT y BERT destaca su aplicabilidad en diferentes modalidades de modelos. Con el potencial de aceleraciones sustanciales en escenarios del mundo real y una calidad de recuperación mejorada, MRL representa un avance prometedor para soluciones de aprendizaje automático más eficientes y versátiles.
Recursos relacionados
Sigue leyendo
Introducing Customer-Managed Encryption Keys (CMEK) on Zilliz Cloud
We're announcing the general availability of Customer-Managed Encryption Keys (CMEK) on Zilliz Cloud.
Why Teams Are Migrating from Weaviate to Zilliz Cloud — and How to Do It Seamlessly
Explore how Milvus scales for large datasets and complex queries with advanced features, and discover how to migrate from Weaviate to Zilliz Cloud.
Will Amazon S3 Vectors Kill Vector Databases—or Save Them?
AWS S3 Vectors aims for 90% cost savings for vector storage. But will it kill vectordbs like Milvus? A deep dive into costs, limits, and the future of tiered storage.