¿Qué deberíamos aprender de ModernBERT?

Allá por 2018, Google lanzó BERT y fue un cambio revolucionario para NLP, mucho antes de la actual ola de LLM. Incluso ahora, muchos Modelos de Lenguaje Pequeños están construidos sobre BERT. En diciembre de 2024, ModernBERT toma lo que hemos aprendido de los recientes desarrollos de LLM y lo aplica a estos modelos más pequeños. ¿Las claves? Mejor eficiencia de parámetros, comprensión de código y manejo de contextos largos.

En esta publicación, analizaremos cómo se compara ModernBERT con dos modelos que conocemos a fondo: jina-XLM-RoBERTa (la columna vertebral multilingüe detrás de jina-embeddings-v3) y RobERTa-large. Veamos cada modelo:

• ModernBERT (dic. 2024) es un SLM recientemente lanzado, desarrollado en colaboración por Answer.AI, LightOn y HuggingFace. Aprovecha optimizaciones modernas como RoPE para una ventana de contexto de 8,192 tokens y capas GeGLU, mejorando el rendimiento mientras mantiene la eficiencia.
• jina-XLM-RoBERTa (sept. 2024) es un modelo de embeddings de texto multilingüe basado en XLM-RoBERTa de Meta. Mientras que el XLM-RoBERTa original mejora RoBERTa usando el conjunto de datos multilingüe XLM, jina-XLM-RoBERTa va más allá con entrenamiento de contexto extendido, implementación de RoPE y soporte para FlashAttention-2. Este modelo sirve como base para jina-embeddings-v3.
• RoBERTa-large (julio 2019) desarrollado por Meta, es una versión mejorada de BERT con 355 millones de parámetros. A través de entrenamiento extendido, conjuntos de datos más grandes e innovaciones como el enmascaramiento dinámico, ha logrado resultados impresionantes en benchmarks clave incluyendo GLUE, SQuAD y RACE. Esto lo hace adecuado para varias tareas de NLP, desde clasificación de texto hasta respuesta de preguntas.

Al comparar estos modelos en tres aspectos centrales, buscamos destacar las efectivas decisiones de diseño de ModernBERT para otros desarrolladores de modelos e identificar insights clave de desarrollo para futuros modelos tipo BERT. También compartiremos nuestros aprendizajes del desarrollo de jina-embeddings-v3 y discutiremos mejoras planeadas para jina-embeddings-v4 y jina-reranker-v3.

tagEficiencia de Parámetros de ModernBERT

Primero examinemos el enfoque de ModernBERT hacia la eficiencia de parámetros - incorpora varios insights clave de desarrollos recientes de LLM. ModernBERT aprovecha tres estrategias principales: una arquitectura más profunda pero más delgada, tamaño de vocabulario controlado, y escalamiento progresivo del modelo comenzando desde modelos más pequeños.

tagArquitectura Profunda y Delgada

ModernBERT-large se hace más profundo con 28 capas, mientras que jina-XLM-RoBERTa y RoBERTa-large funcionan con 24. Pero aquí está lo interesante - iguala a RoBERTa-large en cantidad de parámetros a pesar de esas capas extra. jina-XLM-RoBERTa necesita más parámetros ya que maneja 89 idiomas, mientras que los otros dos se enfocan solo en inglés.

La mayoría de los parámetros de un transformer provienen de las capas de atención y completamente conectadas. ModernBERT se mantiene competitivo en tamaño haciéndose "más delgado" - ejecuta 2,624 unidades ocultas a través de 28 capas, comparado con las 4,096 unidades de RoBERTa-large a través de 24 capas. Esta configuración más "profunda" pero delgada les permite alcanzar sus objetivos de rendimiento sin inflar el modelo.

Este enfoque se alinea con la investigación MobileLLM de Meta, que encontró que para modelos más pequeños, la profundidad importa más que el ancho cuando se trata de capturar patrones complejos y mejorar el rendimiento. Esencialmente, la capacidad de procesar información a través de más capas transformer resulta más valiosa que tener capas más anchas para procesamiento paralelo.

Veamos los datos sobre cómo se desempeña esta arquitectura profunda y delgada.

Tome jina-XLM-RoBERTa - se basa en la arquitectura superficial-gruesa de RoBERTa-large pero aumenta el vocabulario de 50K a 250K tokens y entrena con más datos. Sin embargo, ModernBERT aún lo supera, sugiriendo que el cambio arquitectónico está marcando una diferencia real en eficiencia.

tagEl Tamaño del Vocabulario Importa

Primero, veamos cómo se cuentan los parámetros de vocabulario en transformers. Para cualquier transformer, parámetros de vocabulario = número de tokens distintos × tamaño oculto. Tome jina-XLM-RoBERTa: con 250K tokens y 1,024 dimensiones, necesita 256M parámetros solo para codificación de vocabulario - ¡antes de manejar cualquier tarea de lenguaje real!

Pero aquí está el asunto: los pesos del vocabulario no contribuyen a los mecanismos de atención - son solo tablas de búsqueda. Para los SLM que trabajan con presupuestos fijos de parámetros, un vocabulario más grande significa menos parámetros disponibles para las capas de atención, que realizan el procesamiento real del lenguaje. Esto explica por qué ModernBERT-large en inglés supera al multilingüe jina-XLM-RoBERTa a pesar de ser más pequeño - jina-XLM-RoBERTa asigna más parámetros (¡47%!) para soportar múltiples idiomas. El vocabulario enfocado de ModernBERT no solo mejora el rendimiento sino que también acelera la inferencia, haciéndolo particularmente efectivo para aplicaciones con recursos limitados.

Así que ahora si miramos solo los parámetros del modelo central (excluyendo los pesos del vocabulario), ModernBERT en realidad tiene más poder computacional que sus pares: ModernBERT dedica 19% más parámetros al modelado de lenguaje real que jina-XLM-RoBERTa y 15% más que RoBERTa-large!

tagEscalado del Modelo mediante "Weight Tiling"

Al construir el backbone de jina-BERT-v2, descubrimos que entrenar SLM desde cero era intensivo en recursos y complejo. ModernBERT aborda esto con un enfoque inteligente de inicialización llamado weight tiling - esencialmente inicializando ModernBERT-large desde los pesos de su versión base más pequeña.

Esta técnica no es completamente nueva - se basa en el trabajo de DeepMind con Gopher y también aparece en los modelos Phi-2 de Microsoft. Pero su aplicación aquí es particularmente efectiva para abordar el cuello de botella del entrenamiento de SLM.

Esta estrategia de inicialización le da a ModernBERT-large una ventaja significativa - en lugar de comenzar desde cero, aprovecha los patrones pre-aprendidos de su contraparte más pequeña. Ha demostrado ser particularmente efectiva para escalar modelos de lenguaje en este rango de tamaño.

Encontramos que un modelo con inicio en caliente se recupera rápidamente de una pérdida inicial alta (debido a los parámetros añadidos) a una pérdida bastante cercana a la del modelo base. Podemos expandir 417M parámetros por más de 3× en tamaño y mantener un rendimiento superior al de un modelo equivalente entrenado desde cero hasta la convergencia, lo que implica que las ganancias no se limitaron al inicio del entrenamiento. Sin embargo, en tamaños más grandes, las ganancias relativas logradas en la convergencia disminuyen, especialmente con expansiones en anchura.

El envolvimiento cíclico de pesos no es solo por conveniencia - se alinea bien con cómo las matrices de atención naturalmente exhiben patrones periódicos. La investigación de Gopher muestra que este enfoque realmente brilla para SLM (menos de 9B parámetros), aunque los beneficios comienzan a disminuir a medida que te mueves hacia territorios de modelos más grandes.

tagModelado de Código de ModernBERT

ModernBERT trae un enfoque especializado para la comprensión de código con su tokenizador optimizado para código y datos de entrenamiento. Esta optimización para el procesamiento de código rinde frutos tanto en tareas de comprensión como de recuperación.

Ejecutamos un benchmark usando el corpus jina-embeddings-v2-code, comparando tres modelos como backbones: ModernBERT, jina-XLM-RoBERTa, y RoBERTa-large. ¿La prueba? CodeSearchNet - emparejar descripciones de texto con fragmentos de código. ModernBERT superó a ambas alternativas en toda la línea.

tagLa ventaja del tokenizador

Veamos por qué ModernBERT maneja tan bien el código - utiliza el tokenizador OLMo, que fue entrenado específicamente en código, en lugar de los tokenizadores estándar de BERT/RoBERTa.

Un tokenizador divide el texto UTF-8 en tokens que se mapean a vectores - estos son los que el modelo realmente procesa. Durante el entrenamiento, aprende a combinar secuencias de caracteres que ocurren frecuentemente en tokens únicos. ¿La diferencia? Un tokenizador estándar podría dividir init en in + it, perdiendo el contexto de programación. Pero el tokenizador de ModernBERT, consciente del código, lo obtiene sin dividirlo.

Aquí es donde se pone interesante con el manejo de espacios: ModernBERT preserva los espacios iniciales de Python como tokens únicos y diferencia entre 4 y 8 espacios - crucial para la estructura del código. Mientras tanto, jina-XLM-RoBERTa colapsa todos los espacios continuos en un solo _, y RoBERTa-large trata cada espacio como su propio token. Esto significa que el codificador de ModernBERT recibe una entrada más limpia y significativa al procesar código, mientras que los otros trabajan con tokens fracturados y menos coherentes.

tagManejo de contexto largo de ModernBERT

ModernBERT ha logrado avances significativos en el procesamiento de texto largo, gracias a su extenso corpus de entrenamiento (300B tokens con muestras de 8,192 tokens) y técnicas avanzadas como la atención combinada global y local.

Para evaluar las capacidades de manejo de documentos largos, usamos el conjunto de datos MLDR - un benchmark exhaustivo de texto largo que abarca 13 idiomas. Dado que ModernBERT actualmente solo admite inglés, nos centramos en el subconjunto en inglés de MLDR para comparar ModernBERT con jina-XLM-RoBERTa. Si bien ambos modelos pueden manejar entradas de 8K tokens, RoBERTa-large fue excluido de este benchmark debido a su límite de 512 tokens, que es insuficiente para el análisis de texto largo.

El rendimiento superior de ModernBERT no se debe solo a su extenso entrenamiento en texto largo - se debe en gran parte a su innovadora combinación de mecanismos de atención global y local. A diferencia de jina-XLM-RoBERTa, que aplica una atención global computacionalmente costosa a cada capa, ModernBERT adopta un enfoque más eficiente. Alterna entre atención global (usada cada tercera capa con un theta de 160,000) y atención local (usando una ventana deslizante de 128 tokens con un theta de 100,000). Esta estrategia híbrida mantiene un alto rendimiento mientras reduce dramáticamente el tiempo de entrenamiento.

En ModernBERT, cada tercera capa emplea atención global con un theta RoPE de 160,000 y las capas restantes usan una ventana deslizante de atención local de 128 tokens con un theta RoPE de 10,000. —— ModernBERT

tag¿La lección amarga?

La ley de escalado y la lección amarga sugieren que las mejoras significativas de rendimiento provienen principalmente del aumento en el número de parámetros y datos de entrenamiento. Este principio guió nuestro enfoque de expandir el corpus y usar LoRA para adaptaciones específicas de tareas.

Sin embargo, el éxito de ModernBERT ha revelado que subestimamos el poder de la optimización arquitectónica. Demuestra que los SLMs pueden lograr resultados excepcionales a través de una mejor eficiencia datos-modelo, sin necesariamente aumentar los parámetros. Un reciente informe técnico de Stella Embeddings refuerza este hallazgo, indicando que los métodos actuales de entrenamiento de modelos de embedding pueden mejorarse sin aumentar el corpus o el tamaño del modelo.

Avanzando, anticipamos costos computacionales más bajos y tamaños de modelo más pequeños a medida que obtenemos conocimientos más profundos sobre la utilización de datos e implementamos las técnicas de ModernBERT. A corto plazo, podemos implementar mejoras directas descritas en el paper de ModernBERT - específicamente integrando más datos relacionados con código y adoptando un tokenizador amigable con el código. Cambios más complejos, como cambiar a una arquitectura profunda y delgada o arrancar modelos grandes desde otros más pequeños, requerirán construir modelos base desde cero - una iniciativa a mediano plazo.

Si bien la eficiencia de ModernBERT es notable, su limitación a solo texto apunta a desafíos futuros. A medida que los modelos de embedding multimodales ganan popularidad, nuestro próximo desafío es desarrollar modelos base de búsqueda más inteligentes, rápidos y capaces que puedan manejar entradas para aplicaciones multimodales. Estas aplicaciones demandan ventanas de contexto aún más largas - un desafío de eficiencia que queda por resolver.

tagConclusión

A lo largo de esta publicación, hemos explorado cómo ModernBERT hace avanzar los modelos de la familia BERT a través de tres innovaciones clave: su arquitectura profunda y delgada, tokenizador optimizado, y escalado eficiente usando mosaico de pesos. Estas mejoras permiten a ModernBERT ofrecer un rendimiento sobresaliente en un tamaño relativamente compacto, superando tanto a RoBERTa-large como a jina-XLM-RoBERTa en varias tareas. ModernBERT demuestra que las mejoras arquitectónicas pueden importar más que el tamaño de los parámetros, abriendo puertas para modelos más eficientes. Su uso exitoso del mosaico de pesos muestra cómo el escalado progresivo puede reducir los costos de entrenamiento mientras preserva o incluso mejora el rendimiento. Además, su vocabulario compacto y optimizaciones dirigidas sugieren oportunidades crecientes para SLMs especializados en entornos con recursos limitados.