Entrenamiento consciente de la cuantización de jina-embeddings-v4

La cuantización es un método muy utilizado para abordar los problemas de escalamiento en la IA. El nombre hace que suene complicado, pero simplemente consiste en redondear los números para que ocupen menos espacio. Esto significa vectores de modelo (Embeddings) más pequeños que ocupan menos memoria y espacio de almacenamiento, y una recuperación de información más rápida porque se tarda menos tiempo en comparar vectores. La cuantización es una técnica puramente numérica que no se preocupa por el tipo de datos que procesa tu modelo ni por los casos de uso que tienes, por lo que puede aportar mejoras sin necesidad de muchos conocimientos especializados caros.

Uno podría esperar que la cuantización implique las típicas compensaciones y el cliché de que nada es gratis, donde debemos sacrificar algo de precisión. En este artículo, te mostraremos una forma de hacerla sin pérdidas mediante el entrenamiento consciente de la cuantización (QAT). Esta técnica se utiliza en jina-embeddings-v4 para proporcionar modelos más pequeños que se necesitan en aplicaciones con limitaciones de espacio.

tagDescripción general de las técnicas de cuantización

La cuantización de modelos suele significar una de estas cuatro cosas:

• Cuantización posterior al entrenamiento (PTQ)
• Entrenamiento para salidas de modelo cuantizadas (Output QAT)
• Entrenamiento para modelos totalmente cuantizados (Full QAT)
• Destilación de un nuevo modelo cuantizado a partir de uno no cuantizado existente

La cuantización posterior al entrenamiento (PTQ) acepta el modelo de modelo entrenado tal cual y no lo modifica de ninguna manera. Se trata simplemente de descartar los dígitos menos significativos de los valores de coma flotante producidos por el modelo. Simplemente redondeamos los números y, a veces, los escalamos a un rango.

Output QAT significa ajustar el modelo de modelo para producir vectores óptimos de precisión reducida. Esto significa modificar el modelo, pero no cambia la precisión de los pesos del modelo y, por lo tanto, no reduce su tamaño. Solo se reduce el tamaño del vector de salida.

Full QAT comienza con un modelo de precisión completa totalmente entrenado y reduce la precisión de los pesos del modelo, luego ajusta el rendimiento de este modelo modificado. Esto produce un modelo significativamente más pequeño, así como modelos más pequeños, a costa de realizar algunos ajustes.

La destilación es el proceso de entrenar un nuevo modelo para que coincida con el rendimiento de uno existente. Esto significa crear un nuevo modelo que esté diseñado desde cero como cuantizado, y luego usar el modelo existente para generar tantos datos de entrenamiento como sea necesario para entrenarlo hasta que funcione lo más cerca posible del modelo existente.

Los beneficios de estos cuatro enfoques se resumen en la siguiente tabla:

Los cuatro producen modelos más compactos, pero aparte de PTQ, todos requieren algún entrenamiento adicional, mientras que solo Full QAT y Distillation producen modelos nuevos y más rápidos. Full QAT y Distillation son mucho más caros de implementar porque requieren mucho más entrenamiento que Output QAT.

En este artículo, solo vamos a ver PTQ y Output QAT, que no cambian el tamaño ni la velocidad del modelo de modelo.

tagConfiguración experimental

Para estos experimentos, nuestro modelo de referencia es jina-embeddings-v4 con el adaptador de recuperación, que produce vectores de coma flotante (FP32) de precisión de 32 bits en 2048 dimensiones. Por lo tanto, cada modelo tiene un tamaño de 8196 bytes, u 8kB.

Estudiamos varias condiciones experimentales utilizando tareas de evaluación comparativa de recuperación de documentos de consulta del conjunto de NanoBEIR benchmark. El proceso de recuperación utiliza la similitud coseno entre vectores para encontrar y clasificar los documentos que mejor coinciden con las consultas.

• Línea base: el rendimiento de los vectores de modelo jina-embeddings-v4 sin ninguna cuantización. Todos estos experimentos utilizaron una versión beta del modelo, y el rendimiento de la versión es algo mejor.
• PTQ: cuantificamos los vectores de salida en vectores binarios sin cambiar el modelo.
• Output QAT: cuantificamos los vectores de salida y aplicamos un ajuste fino al adaptador de recuperación para mejorar su rendimiento en condiciones cuantificadas.

tagNiveles de cuantización

Experimentamos con cuatro niveles diferentes de cuantización.

• Enteros de 8 bits: los valores FP32 se reducen a enteros en el rango de -128 a 127, lo que reduce los modelos 4 veces a 2048 bytes.
• Enteros de 4 bits: igual que para los enteros de 4 bits, pero asignamos al rango de -8 a 7, lo que reduce el tamaño de los vectores en un factor de 8, a 1024 bytes.
• Cuantización ternaria: todos los valores se asignan a uno de los tres valores: -1, 0, 1. Almacenado de forma óptima, esto reduce cada dimensión a 1,6 bits, lo que reduce el tamaño de los vectores de modelo aproximadamente 40 veces a aproximadamente 230 bytes.
• Cuantización binaria: convertimos los valores escalares FP32 a un bit, utilizando el tipo de datos torch.sign, que proporciona solo dos valores, que tardan un bit en almacenarse. Esto reduce los vectores de modelo de 2048 dimensiones de 8192 bytes a 128 bytes, una reducción de 64 veces.

tagEscalado

Para la cuantización binaria, la cuantización es muy simple: si un valor de vector es superior a 0 o positivo, se asigna a 1. De lo contrario, se asigna a -1.

Para los otros escenarios de cuantización, normalizamos los valores a un rango y luego redondeamos al valor más cercano permitido por el nivel de cuantización. Los vectores de modelo consisten en números de escala entre -∞ y +∞ (o, en la práctica, números positivos y negativos realmente grandes). Usamos dos números, $max$ y $min$, para escalar los valores para la cuantización.

Para la cuantización ternaria, tomamos cada componente de vector $v$ y lo traducimos de la siguiente manera:

• si $v$ ≥ $max$, $v$ se convierte en 1.
• si $v$ ≤ $min$, $v$ se convierte en -1.
• si $min$ < $v$ < $max$, $v$ se convierte en 0.

Para enteros de 4 bits:

• si $v$ ≥ $max$, $v$ se convierte en 7.
• si $v$ ≤ $min$, $v$ se convierte en -8.
• si $min$ < $v$ < $max$, $v$ se convierte en $16*(v - min)/(max - min) - 8$, luego se redondea al entero más cercano. Esto escala el valor al rango $[-8,7]$.

Para enteros de 8 bits:

• si $v$ ≥ $max$, $v$ se convierte en 127.
• si $v$ ≤ $min$, $v$ se convierte en -128.
• si $min$ < $v$ < $max$, $v$ se convierte en $256*(v - min)/(max - min) - 128$, redondeado al entero más cercano. Esto escala el valor al rango $[-128,127]$.

Para calcular $max$ y $min$, utilizamos dos enfoques:

• Mín./máx.: procesamos nuestros datos en lotes y, para cada lote, identificamos el componente vectorial más alto y más bajo, estableciendo $max$ en el más alto y $min$ en el más bajo.
• Promedio móvil en lotes: para cada lote, calculamos el promedio y la desviación estándar de los componentes del vector. Mantuvimos un promedio móvil tanto del promedio como de la desviación estándar a medida que procesábamos todos los lotes. Si $avg$ es el promedio móvil actual de los valores promedio del lote, y $std$ es el promedio móvil actual de las desviaciones estándar, entonces para cada lote:

$max = avg + std$
$min = avg - std$

tagAjuste fino de QAT

Para los experimentos de PTQ, utilizamos el modelo tal cual y cuantificamos los modelos que produjo utilizando los métodos descritos anteriormente.

Para el Output QAT, ajustamos el modelo utilizando la estimación directa. Esto significa que invertimos el proceso de cuantización, restaurando la precisión completa de los valores, antes de calcular la pérdida (es decir, el error), y luego usamos esa métrica de pérdida para ajustar el modelo.

En cada caso, realizamos un ajuste fino durante 10.000 pasos, guardando un punto de control cada 500 pasos. Luego, conservamos el punto de control con la puntuación más alta en el NanoBEIR benchmark.

tagCuantificación Asimétrica

PTQ y Output QAT reducen el tamaño de los vectores de "embeddings" (vectores modelo), pero no reducen el tamaño del modelo ni la velocidad de inferencia; todo el ahorro se encuentra en el tamaño de los "embeddings" (vectores modelo) de documentos almacenados y la velocidad de recuperación.

Como resultado, probamos tanto la cuantificación de los vectores de consulta como el dejarlos sin cuantificar en el momento de la recuperación, ya que de cualquier manera no cambia el tamaño de los vectores de "embeddings" (vectores modelo) almacenados.

tagResultados

Probamos nueve condiciones en total, que se resumen en las tablas siguientes:

Tabla 2: Condiciones Experimentales

Tabla 3: Puntuación promedio (en % correcto) para cada condición en los doce NanoBEIR benchmarks.

Como puede ver en la tabla anterior, el ajuste fino para la cuantificación mejora las puntuaciones. La única diferencia entre las condiciones PTQ Binario y QAT Binario es el ajuste fino, y la diferencia en la puntuación es significativa. Del mismo modo, vemos una mejora de casi el 2% en las puntuaciones entre las condiciones PTQ Binario Docs Only y QAT Binario Docs Only, que solo se distinguen por el mismo ajuste fino.

Como era de esperar, también vemos que las puntuaciones generalmente mejoran cuanto menos cuantificamos, con la cuantificación de 4 bits obteniendo mejores resultados que la ternaria, y la ternaria mejor que la binaria. Sin embargo, ir más allá a 8 bits no parece haber mejorado nada.

Solo probamos dejar las consultas sin cuantificar en los casos binarios, pero esto parece mejorar el rendimiento.

Finalmente, nuestras pruebas sugieren que el método de escala de promedio móvil supera al enfoque simplista de min/max.

tagConclusión

La cuantificación tiene algunas ventajas operativas importantes para los modelos de "embeddings" (vectores modelo), al reducir significativamente el tamaño de los vectores de "embeddings" (vectores modelo) y acelerar la recuperación de información. Si bien la cuantificación simple posterior al entrenamiento (PTQ) proporciona beneficios inmediatos en términos de memoria y almacenamiento, nuestros experimentos demuestran que el entrenamiento con reconocimiento de cuantificación (QAT) mitiga significativamente las inevitables pérdidas de precisión. El ajuste fino produjo consistentemente mejores puntuaciones.

El grado de cuantificación impacta directamente en el rendimiento, que es lo que cabría esperar de un método basado en reducir la precisión de los valores. La cuantificación menos agresiva (por ejemplo, 4 bits) generalmente supera a los métodos más agresivos (por ejemplo, binario), pero sorprendentemente, no hubo una diferencia significativa en el rendimiento entre la cuantificación de 8 bits y la de 4 bits. Parece que hasta que no se alcanza un cierto umbral de imprecisión, hay poca diferencia entre una cuantificación mayor y menor.

Las estrategias de escala también son significativas, con el método de promedio móvil mostrando resultados superiores en comparación con un enfoque fijo de min/max. El uso de valores de escala que son relativos a los datos parece funcionar significativamente mejor y merece una mayor exploración.

La cuantificación puede obtener más de sus modelos de "embeddings" (vectores modelo) por menos. Aunque este artículo no explora todas las opciones de cuantificación, explora dos que son fácilmente accesibles y que tienen beneficios reales que ofrecer. Estamos trabajando para refinar y mejorar las estrategias de cuantificación para poder reducir aún más los costes de los usuarios, y esperamos lanzar soporte binario para jina-embeddings-v4 en un futuro próximo.