Aprendizaje en dispositivo durante toda la vida, más cercano con una nueva técnica de entrenamiento

Un equipo de investigadores en MIT y el MIT-IBM Watson AI Lab desarrollaron una nueva técnica que permite el entrenamiento en dispositivo utilizando menos de un cuarto de megabyte de memoria. El nuevo desarrollo es un logro impresionante, ya que otras soluciones de entrenamiento suelen necesitar más de 500 megabytes de memoria, lo que supera la capacidad de 256 kilobytes de la mayoría de los microcontroladores.

Al entrenar un modelo de aprendizaje automático en un dispositivo de borde inteligente, puede adaptarse a nuevos datos y hacer mejores predicciones. Con eso dicho, el proceso de entrenamiento suele requerir mucha memoria, por lo que a menudo se lleva a cabo con computadoras en un centro de datos antes de que el modelo se despliegue en un dispositivo. Este proceso es mucho más costoso y plantea preocupaciones de privacidad en comparación con la nueva técnica desarrollada por el equipo.

Los investigadores desarrollaron los algoritmos y el marco de trabajo de una manera que reduce la cantidad de cálculo necesaria para entrenar un modelo, lo que hace que el proceso sea más rápido y eficiente en términos de memoria. La técnica puede ayudar a entrenar un modelo de aprendizaje automático en un microcontrolador en solo unos minutos.

La nueva técnica también ayuda con la privacidad, ya que mantiene los datos en el dispositivo, lo que es importante cuando se trata de datos sensibles. Al mismo tiempo, el marco de trabajo mejora la precisión del modelo en comparación con otros enfoques.

Song Han es profesor asociado en el Departamento de Ingeniería Eléctrica y Ciencias de la Computación (EECS), miembro del MIT-IBM Watson AI Lab y autor principal del artículo de investigación.

“Nuestro estudio permite que los dispositivos IoT no solo realicen inferencia, sino que también actualicen continuamente los modelos de IA con nuevos datos recopilados, allanando el camino para el aprendizaje en dispositivo durante toda la vida”, dijo Han. “La baja utilización de recursos hace que el aprendizaje profundo sea más accesible y puede tener un alcance más amplio, especialmente para dispositivos de borde de baja potencia”.

El artículo incluyó a los autores principales y estudiantes de doctorado de EECS Ji Lin y Ligeng Zhu, y a los investigadores postdoctorales de MIT Wei-Ming Chen y Wei-Chen Wang. También incluyó a Chuang Gan, miembro del personal de investigación principal en el MIT-IBM Watson AI Lab.

Hacer que el proceso de entrenamiento sea más eficiente

Para hacer que el proceso de entrenamiento sea más eficiente y menos intensivo en términos de memoria, el equipo se basó en dos soluciones algorítmicas. La primera se conoce como actualización dispersa, que utiliza un algoritmo que identifica los pesos más importantes para actualizar durante cada ronda de entrenamiento. El algoritmo congela los pesos uno a uno hasta que la precisión cae a un umbral determinado, en cuyo punto se detiene. Los pesos restantes se actualizan y las activaciones correspondientes a los pesos congelados no necesitan almacenarse en memoria.

“Actualizar todo el modelo es muy costoso porque hay mucha activación, por lo que la gente tiende a actualizar solo la última capa, pero como puede imaginar, esto perjudica la precisión”, dijo Han. “Para nuestro método, actualizamos selectivamente esos pesos importantes y nos aseguramos de que la precisión se preserve por completo”.

La segunda solución desarrollada por el equipo implica el entrenamiento cuantizado y la simplificación de los pesos. Un algoritmo primero redondea los pesos a solo ocho bits a través de un proceso de cuantización, lo que también reduce la cantidad de memoria para el entrenamiento y la inferencia, siendo la inferencia el proceso de aplicar un modelo a un conjunto de datos y generar una predicción. El algoritmo luego se basa en una técnica llamada escalado consciente de la cuantización (QAS), que actúa como un multiplicador para ajustar la relación entre el peso y el gradiente. Esto ayuda a evitar cualquier caída en la precisión que podría resultar del entrenamiento cuantizado.

Los investigadores desarrollaron un sistema llamado motor de entrenamiento pequeño, que ejecuta las innovaciones algorítmicas en un microcontrolador simple que carece de un sistema operativo. Para completar más trabajo en la etapa de compilación, antes de la implementación del modelo en el dispositivo de borde, el sistema cambia el orden de los pasos en el proceso de entrenamiento.

“Empujamos mucha computación, como la auto-diferenciación y la optimización del grafo, al tiempo de compilación. También recortamos agresivamente los operadores redundantes para admitir actualizaciones dispersas. Una vez en tiempo de ejecución, tenemos mucha menos carga de trabajo que hacer en el dispositivo”, dice Han.

Técnica altamente eficiente

Mientras que las técnicas tradicionales diseñadas para el entrenamiento ligero suelen necesitar alrededor de 300 a 600 megabytes de memoria, la optimización del equipo solo necesitó 157 kilobytes para entrenar un modelo de aprendizaje automático en un microcontrolador.

El marco de trabajo se probó entrenando un modelo de visión por computadora para detectar personas en imágenes, y aprendió a completar esta tarea en solo 10 minutos. El método también pudo entrenar un modelo más de 20 veces más rápido que otros métodos.

Los investigadores ahora buscarán aplicar las técnicas a modelos de lenguaje y diferentes tipos de datos. También quieren utilizar este conocimiento adquirido para reducir el tamaño de los modelos más grandes sin pérdida de precisión, lo que también podría ayudar a reducir la huella de carbono del entrenamiento de modelos de aprendizaje automático a gran escala.