Los modelos de lenguaje grande de razonamiento (LLM) están diseñados para resolver problemas complejos dividiéndolos en una serie de pasos más pequeños. Estos potentes modelos son particularmente buenos para tareas desafiantes como programación avanzada y planificación de varios pasos.
Pero desarrollar modelos de razonamiento exige una enorme cantidad de cálculo y energía debido a ineficiencias en el proceso de formación. Mientras algunos de los procesadores de alta potencia trabajan continuamente en consultas complicadas, otros en el grupo permanecen inactivos.
Investigadores del MIT y de otros lugares encontraron una manera de utilizar este tiempo de inactividad computacional para acelerar de manera eficiente el entrenamiento del modelo de razonamiento.
Su nuevo método entrena automáticamente un modelo más pequeño y más rápido para predecir los resultados del LLM de razonamiento más grande, que el modelo más grande verifica. Esto reduce la cantidad de trabajo que debe realizar el modelo de razonamiento, acelerando el proceso de formación.
La clave de este sistema es su capacidad para entrenar e implementar el modelo más pequeño de forma adaptativa, de modo que se active sólo cuando algunos procesadores estén inactivos. Al aprovechar recursos computacionales que de otro modo se habrían desperdiciado, acelera la capacitación sin incurrir en gastos generales adicionales.
Cuando se probó en múltiples LLM de razonamiento, el método duplicó la velocidad de entrenamiento preservando la precisión. Esto podría reducir el coste y aumentar la eficiencia energética del desarrollo de LLM avanzados para aplicaciones como la previsión de tendencias financieras o la detección de riesgos en las redes eléctricas.
“La gente quiere modelos que puedan manejar tareas más complejas. Pero si ese es el objetivo del desarrollo de modelos, entonces debemos priorizar la eficiencia. Encontramos una solución sin pérdidas para este problema y luego desarrollamos un sistema completo que puede ofrecer aceleraciones bastante dramáticas en la práctica”, dice Qinghao Hu, postdoctorado del MIT y coautor principal de un artículo sobre esta técnica.
En el artículo lo acompañan el coautor principal Shang Yang, un estudiante de posgrado en ingeniería eléctrica e informática (EECS); Junxian Guo, estudiante de posgrado de EECS; el autor principal Song Han, profesor asociado en EECS, miembro del Laboratorio de Investigación de Electrónica y científico distinguido de NVIDIA; así como otros en NVIDIA, ETH Zurich, el MIT-IBM Watson AI Lab y la Universidad de Massachusetts en Amherst. La investigación se presentará en la Conferencia Internacional ACM sobre Soporte Arquitectónico para Lenguajes de Programación y Sistemas Operativos.
Cuello de botella en el entrenamiento
Los desarrolladores quieren que los LLM de razonamiento identifiquen y corrijan errores en su proceso de pensamiento crítico. Esta capacidad les permite resolver consultas complicadas que harían tropezar a un LLM estándar.
Para enseñarles esta habilidad, los desarrolladores entrenan a los LLM de razonamiento utilizando una técnica llamada aprendizaje por refuerzo (RL). El modelo genera múltiples respuestas potenciales a una consulta, recibe una recompensa para el mejor candidato y se actualiza en función de la respuesta principal. Estos pasos se repiten miles de veces a medida que el modelo aprende.
Pero los investigadores descubrieron que el proceso de generar múltiples respuestas, llamado implementación, puede consumir hasta el 85 por ciento del tiempo de ejecución necesario para el entrenamiento de RL.
“En comparación, actualizar el modelo, que es la parte real de ‘capacitación’, consume muy poco tiempo”, dice Hu.
Este cuello de botella ocurre en los algoritmos RL estándar porque todos los procesadores del grupo de entrenamiento deben finalizar sus respuestas antes de poder pasar al siguiente paso. Debido a que algunos procesadores pueden estar trabajando en respuestas muy largas, otros que generaron respuestas más cortas esperan a que terminen.
“Nuestro objetivo era convertir este tiempo de inactividad en una aceleración sin desperdiciar costes”, añade Hu.
Intentaron utilizar una técnica existente, llamada decodificación especulativa, para acelerar las cosas. La decodificación especulativa implica entrenar un modelo más pequeño llamado redactor para adivinar rápidamente los resultados futuros del modelo más grande.
El modelo más grande verifica las conjeturas del redactor y las respuestas que acepta se utilizan para el entrenamiento.
Debido a que el modelo más grande puede verificar todas las conjeturas del redactor a la vez, en lugar de generar cada resultado secuencialmente, acelera el proceso.
Una solución adaptativa
Pero en la decodificación especulativa, el modelo del redactor normalmente se entrena sólo una vez y permanece estático. Esto hace que la técnica sea inviable para el aprendizaje por refuerzo, ya que el modelo de razonamiento se actualiza miles de veces durante el entrenamiento.
Un redactor estático rápidamente se volvería obsoleto e inútil después de unos pocos pasos.
Para superar este problema, los investigadores crearon un sistema flexible conocido como “Domar la cola larga” o TLT.
La primera parte de TLT es un entrenador de redactor adaptativo, que utiliza tiempo libre en procesadores inactivos para entrenar el modelo de redactor sobre la marcha, manteniéndolo bien alineado con el modelo de destino sin utilizar recursos computacionales adicionales.
El segundo componente, un motor de implementación adaptativo, gestiona la decodificación especulativa para seleccionar automáticamente la estrategia óptima para cada nuevo lote de entradas. Este mecanismo cambia la configuración de decodificación especulativa según las características de la carga de trabajo de entrenamiento, como la cantidad de entradas procesadas por el modelo preliminar y la cantidad de entradas aceptadas por el modelo de destino durante la verificación.
Además, los investigadores diseñaron el borrador del modelo para que fuera liviano y pudiera entrenarse rápidamente. TLT reutiliza algunos componentes del proceso de entrenamiento del modelo de razonamiento para entrenar al redactor, lo que genera ganancias adicionales en aceleración.
“Tan pronto como algunos procesadores terminan sus consultas breves y quedan inactivos, inmediatamente los cambiamos para que realicen un entrenamiento de modelo preliminar utilizando los mismos datos que están usando para el proceso de implementación. El mecanismo clave es nuestra decodificación especulativa adaptativa; estas ganancias no serían posibles sin ella”, dice Hu.
Probaron TLT en múltiples LLM de razonamiento que fueron capacitados utilizando conjuntos de datos del mundo real. El sistema aceleró el entrenamiento entre un 70 y un 210 por ciento preservando al mismo tiempo la precisión de cada modelo.
Como ventaja adicional, el modelo de redactor pequeño podría utilizarse fácilmente para una implementación eficiente como subproducto gratuito.
En el futuro, los investigadores quieren integrar TLT en más tipos de marcos de entrenamiento e inferencia y encontrar nuevas aplicaciones de aprendizaje por refuerzo que podrían acelerarse utilizando este enfoque.
“A medida que el razonamiento continúa convirtiéndose en la principal carga de trabajo que impulsa la demanda de inferencia, el TLT de Qinghao es un gran trabajo para hacer frente al cuello de botella computacional del entrenamiento de estos modelos de razonamiento. Creo que este método será muy útil en el contexto de la informática de IA eficiente”, afirma Han.
Este trabajo está financiado por el MIT-IBM Watson AI Lab, el MIT AI Hardware Program, el MIT Amazon Science Hub, Hyundai Motor Company y la National Science Foundation.