NVIDIA publica receta para afinar embeddings en un día | Keryc
Puedes convertir un modelo de embeddings general en uno que entienda de verdad tu dominio con una sola GPU y menos de un día de entrenamiento. ¿Suena a magia? No lo es: es una receta práctica que combina generación sintética de datos, minería de negativos difíciles, entrenamiento contrastivo y despliegue optimizado.
Qué ofrece esta receta de NVIDIA
NVIDIA publica una pipeline completa (NeMo Data Designer, NeMo Automodel y Nemotron) que va desde documentos crudos hasta un servicio de embeddings listo para producción. Los puntos clave:
Generación automática de pares (pregunta, documento) usando un LLM para crear datos sintéticos de alta calidad.
Minería de negativos difíciles para enseñar al modelo a distinguir pasajes confusos.
Soporte para preguntas multi-hop (1 a 3 saltos) y unrolling para entrenamiento contrastivo.
Exportación a ONNX/TensorRT y despliegue en NVIDIA NIM con API compatible OpenAI.
¿El resultado? En sus pruebas vieron mejoras de doble dígito en métricas como Recall@10 y nDCG@10. Atlassian aplicó la receta a su conjunto Jira y llevó Recall@60 de 0.751 a 0.951 usando una sola GPU A100 80GB.
Requisitos y herramientas
Necesitas:
NeMo Data Designer para generación sintética.
NeMo Automodel y Nemotron para entrenar embeddings.
BEIR para evaluación de recuperación de información.
NeMo Export-Deploy para convertir a ONNX/TensorRT.
NVIDIA NIM para servir en producción.
Un directorio con tus documentos de dominio (.txt, .md, etc.).
Una clave de API de NVIDIA (gratuita en build.nvidia.com).
Una GPU NVIDIA Ampere o más nueva con al menos 80 GB VRAM (A100/H100 80GB probadas).
Si no tienes todo, la pipeline es modular: puedes empezar en la etapa que te convenga.
Preparar datos: split, minería de negativos difíciles, unrolling
nemotron embed prep -c default
Afinar el modelo embedding
nemotron embed finetune -c default
Evaluar base vs afinado (BEIR)
nemotron embed eval -c default
Exportar a ONNX/TensorRT
nemotron embed export -c default
Desplegar en NIM
nemotron embed deploy -c default
Cada etapa puede tomar minutos u horas; el total está pensado para caber en un día con la configuración adecuada.
Cómo funciona la generación sintética (SDG)
En vez de etiquetar manualmente miles de pares query-documento, el pipeline usa un LLM (por ejemplo nvidia/nemotron-3-nano-30b-a3b) para leer tus documentos y generar pares QA. El proceso produce distintos tipos de preguntas: consultas contextuales, preguntas factuales y multi-hop.
Ejemplo de chunk de documento:
El TDP del GPU H100 en formato SXM es 700W. La solución de enfriamiento debe mantener la temperatura de unión por debajo de 83°C. Para despliegues densos superiores a 4 GPUs por nodo se recomienda refrigeración líquida.
Ejemplos de pares generados incluyen preguntas de lookup simples y preguntas multi-hop que requieren conectar información de varias secciones. Cada par recibe puntuaciones de calidad (relevancia, precisión, soporte contextual y claridad) y solo los pares que superan umbrales entran al entrenamiento.
Minería de negativos difíciles: por qué importa
Si entrenas solo con pares positivos, el modelo aprende a separar lo muy distinto pero no aprende a rechazar pasajes que parecen relevantes pero no lo son. La minería de negativos difíciles identifica esos "casi correctos" para que el modelo aprenda distinciones finas.
El proceso automatizado:
Embedea todas las queries y pasajes con el modelo base.
Calcula similitudes y enmascara los positivos reales.
Aplica un filtro de margen: elimina candidatos que estén por encima del 95% del mínimo score positivo (evita falsos negativos).
Selecciona los top-k como negativos difíciles (por defecto 5 por query).
¿Por qué usar un techo del 95%? Para evitar enseñar al modelo que un pasaje que podría ser relevante es negativo por accidente.
Multi-hop y unrolling
Las preguntas multi-hop pueden tener múltiples documentos positivos. La técnica de unrolling convierte una pregunta multi-hop en múltiples ejemplos (query, cada documento positivo) para que la pérdida contrastiva vea cada positivo por separado. Así el modelo aprende que varios pasajes pueden ser relevantes para una sola consulta compuesta.
Fine-tuning: arquitectura y parámetros recomendados
La receta afina un biencoder (dos encoders: uno para queries y otro para documentos) usando pérdida contrastiva. Parámetros por defecto útiles:
Parámetro
Valor por defecto
Notas
Epochs
3
Para datasets grandes baja a 1 o 2
Learning rate
1e-5
Prueba 5e-6 o 2e-5 si hace falta
Warmup steps
5
5-10% del total de pasos funciona bien
Global batch size
128
Se escala automáticamente si tu dataset es pequeño
Pasajes por query
5
1 positivo + 4 negativos difíciles
Temperatura
0.02
Baja temperatura = distribución muy afilada
La temperatura agresiva (0.02) funciona porque los negativos difíciles producen gradientes fuertes y precisos.
Si tienes menos de 2000 ejemplos, la pipeline ajusta batch size, frecuencia de checkpoint y validación para que el entrenamiento sea estable.
Evaluación con BEIR y resultados esperados
La evaluación estándar usa BEIR y calcula nDCG@k, Recall@k, Precision@k y MAP@k en k = 1, 5, 10, 100. En sus pruebas con el dataset sintético Retrieval Synthetic NVDocs, los resultados fueron:
NDCG:
NDCG@1: 0.55178 -> 0.60796 (+10.2%)
NDCG@5: 0.51894 -> 0.57689 (+11.2%)
NDCG@10: 0.55506 -> 0.61559 (+10.9%)
NDCG@100: 0.60617 -> 0.66567 (+9.8%)
Recall:
Recall@1: 0.28478 -> 0.31547 (+10.8%)
Recall@5: 0.54486 -> 0.60288 (+10.6%)
Recall@10: 0.62979 -> 0.69296 (+10.0%)
Recall@100: 0.81421 -> 0.87020 (+6.9%)
Un buen fine-tune típicamente entrega ~15% de mejora en nDCG@10 y Recall@10 en menos de un día, aunque los números reales dependen de calidad del corpus y tamaño del dataset.
Atlassian reportó un caso real: Recall@60 pasó de 0.751 a 0.951 (un 26.7% de ganancia) en su dataset de Jira con una sola A100 80GB.
Export, quantización y despliegue
Para producción conviene exportar a ONNX y (opcionalmente) compilar un engine TensorRT para máxima latencia/throughput. La pipeline soporta:
Export a ONNX (opset 17).
Compilación TensorRT con perfiles de optimización por batch y longitud de secuencia.
Cuantización FP8 para acelerar más.
Empaquetado en un contenedor NIM que expone un endpoint /v1/embeddings compatible con APIs tipo OpenAI.
Además hay una verificación de precisión que compara las métricas BEIR en el endpoint NIM y marca desviaciones mayores que tolerancia como alertas.
Consejos prácticos y problemas comunes
Datos desordenados = resultados mediocres. Limpia y formatea tus documentos antes de SDG.
Si ves overfitting, baja epochs o sube el umbral de calidad en SDG.
Si no hay suficiente datos, agrega más documentos o prueba una LLM más potente para SDG.
Ajusta learning rate (ejemplos: 5e-6 para datasets grandes, 2e-5 para muy pequeños).
Empieza con 50-100 documentos para un POC rápido; la pipeline escala bien.
Reflexión final
No necesitas un centro de datos ni un equipo enorme para obtener embeddings que realmente entiendan tu dominio. Con la receta de NVIDIA puedes generar datos de entrenamiento sin etiquetado manual, enseñarle a tu modelo a distinguir casos confusos y desplegarlo en producción en menos de un día. ¿Tienes documentos de dominio y una GPU con VRAM suficiente? Entonces tienes todo para experimentar y mejorar la relevancia de tus búsquedas o sistemas RAG.
