Olmix: framework para mezclar datos en el desarrollo de LMs | Keryc
Olmix llega como respuesta práctica a un problema que seguro te suena: mezclar textos, código, PDFs y matemáticas en la receta correcta para entrenar un modelo de lenguaje es más arte que ciencia. ¿Cuánto de cada fuente usas? ¿Cómo sabes si estás desperdiciando datos o compute? Olmix propone un flujo repetible y eficiente para decidir y actualizar esa mezcla a lo largo del desarrollo de un LM.
Qué problema resuelve Olmix
Cuando construyes un lenguaje modelo, los conjuntos de datos cambian todo el tiempo: añades corpus nuevos, filtras basura, particionas dominios. Recalcular la mezcla desde cero cada vez es caro. Además, la literatura no da una guía consistente sobre cómo configurar la mezcla: tamaño de los modelos proxy, número de experimentos, tipo de regresión, y otras decisiones quedan a la intuición.
Olmix aborda dos retos: ofrece defaults empíricamente validados para no adivinar configuraciones, y propone técnicas de reutilización de mezclas para actualizar la mezcla eficientemente a medida que tu corpus evoluciona.
OlmixBase: un punto de partida con respaldo experimental
OlmixBase encapsula los hallazgos de un estudio amplio pensado para responder preguntas prácticas: ¿qué tan pequeños pueden ser tus modelos proxy? ¿Cuántas corridas necesitas? ¿Qué regresor usar?
En el paper entrenaron modelos target de 1B parámetros sobre el corpus DCLM, particionado en 24 dominios temáticos, y evaluaron en 52 tareas downstream (math, code, commonsense QA). La métrica principal fue bits-per-byte (BPB), que mide log-likelihood negativo normalizado por bytes en UTF-8.
Estos son los puntos clave que incorpora OlmixBase:
Proxy suficientemente grandes para ser confiables. Proxies por encima de ~15M parámetros mantienen una alta correlación de ranking con modelos target de 1B (ρ > 0.89). Proxies muy pequeños (1M) caen a ρ = 0.73, demasiado ruidosos para decisiones de mezcla confiables. Si estás usando proxies extremadamente pequeños para ahorrar compute, podrías estar sacrificando señales útiles.
Costes que escalan linealmente con el número de dominios. El tamaño del "swarm" (las corridas proxy) debe crecer en O(m) para m dominios. Esto te da una regla práctica para presupuestar compute cuando tu conjunto de dominios crece.
Regresión log-lineal como default robusto. Probaron varios modelos regresores y encontraron que, en general, la regresión log-lineal ofrece el mejor ajuste global y se mantiene competitiva en validación downstream. La elección del regresor puede depender del tamaño del swarm, pero log-lineal es un punto de partida sólido.
Restringir la repetición de datos escasos. Un fallo común es que el optimizador asigne demasiadas iteraciones a un dominio pequeño (por ejemplo, código), forzando repetición dañina. OlmixBase incluye restricciones de factibilidad para evitar asignar más entrenamiento que la disponibilidad real de datos.
Todo esto viene "preconfigurado" en OlmixBase para darte un inicio razonable y reproducible en la etapa temprana de desarrollo.
Reutilización de mezcla: cómo actualizar sin recomputar todo
La segunda gran pieza de Olmix es la reutilización de mezcla. El insight práctico es simple: cuando cambias tu corpus normalmente alteras unos pocos dominios, no todos. Entonces no necesitas recomputar una mezcla global si puedes reutilizar las relaciones entre dominios que no cambiaron.
La técnica central es agrupar los dominios sin cambios en un "dominio virtual" y resolver un problema de mezcla mucho más pequeño que incluya ese virtual y los dominios cambiados. Luego expandes la solución de vuelta al espacio completo. Esto reduce drásticamente el número de corridas proxy porque el coste escala linealmente con la cantidad de dominios optimizados.
Olmix presenta dos estrategias:
Full Mixture Reuse: mantener fijas las proporciones de todos los dominios sin cambios y recomputar solo los cambiados. Es extremadamente eficiente y en muchos escenarios obtiene la mayor parte del beneficio de recomputar todo.
Partial Mixture Reuse: cuando hay acoplamientos entre dominios (por ejemplo, la sección web de "software development" compite con un nuevo corpus de código para resolver tareas de programación), recomputas también un subconjunto seleccionado de dominios sin cambios que están acoplados con los nuevos. Esto reduce el efecto de acoplamiento y cierra la brecha respecto a recomputación completa con solo unas pocas corridas extra.
Resultados prácticos (resumen de experimentos)
Para una secuencia realista de 5 actualizaciones que termina en 64 dominios, entrenando modelos de 1B parámetros sobre 100B tokens:
Full Mixture Reuse alcanza 95% de la mejora que obtendrías con recomputación completa, usando 74% menos corridas proxy (216 vs 832).
Partial Mixture Reuse llega al 98% con 67% menos corridas.
La mejor mezcla obtenida mediante las técnicas de reutilización supera en 12.2% a la distribución natural (baseline sin mezcla) y es 3.05x más eficiente en datos.
La mezcla final tiende a sobreponderar dominios de alto valor como arXiv, FineMath y código, algo que tiene sentido cuando te fijas en tareas académicas y matemáticas complejas.
¿Por qué importa esto para tu proyecto? (consejos prácticos)
Si entrenas LMs con datos heterogéneos, Olmix te ayuda a dejar de adivinar y a gastar menos compute cuando tu corpus cambia.
Recomendación rápida:
Empieza con OlmixBase para obtener una mezcla inicial: usa proxies ≈ 15M parámetros, aplica regresión log-lineal y define restricciones de repetición según disponibilidad de datos.
Cuando el corpus cambie, usa Full Mixture Reuse si los cambios tocan pocos dominios sin grandes solapamientos semánticos.
Si detectas acoplamientos (nuevos datos que compiten con partes del corpus existente por las mismas tareas), aplica Partial Mixture Reuse e incluye esos dominios en la recomputación.
Mide en BPB y en métricas downstream relevantes para tus tareas específicas; no asumas que la mezcla óptima según BPB será siempre idéntica para cada tarea.
Reflexión final
Olmix no es solo otro algoritmo académico: es una caja de herramientas pensada para el flujo real de desarrollo de LMs. Si reconoces el dolor de recalcular mezclas a cada iteración, o de confiar en proxies demasiado chicas que te dan señales ruidosas, Olmix ofrece reglas prácticas y técnicas de actualización que escalan con tus necesidades.
La mezcla de datos es una palanca de primer orden en la calidad del modelo. Olmix nos recuerda que con buenas decisiones experimentales y un poco de ingeniería se puede convertir ese arte desordenado en un flujo repetible y eficiente.
