Los ingenieros de DeepMind aceleran nuestra investigación mediante la creación de herramientas, la ampliación de algoritmos y la creación de desafiantes mundos virtuales y físicos para entrenar y probar sistemas de inteligencia artificial (IA). Como parte de este trabajo, evaluamos constantemente nuevas bibliotecas y marcos de aprendizaje automático.
Recientemente, hemos descubierto que un número cada vez mayor de proyectos están bien atendidos por JAXun marco de aprendizaje automático desarrollado por Investigación de Google equipos. JAX resuena bien con nuestra filosofía de ingeniería y ha sido ampliamente adoptado por nuestra comunidad de investigación durante el último año. Aquí compartimos nuestra experiencia de trabajo con JAX, describimos por qué lo encontramos útil para nuestra investigación de IA y brindamos una descripción general del ecosistema que estamos construyendo para apoyar a los investigadores de todo el mundo.
¿Por qué JAX?
JAX es una biblioteca de Python diseñada para computación numérica de alto rendimiento, especialmente investigación de aprendizaje automático. Su API para funciones numéricas se basa en NumPy, una colección de funciones utilizadas en informática científica. Tanto Python como NumPy son ampliamente utilizados y familiares, lo que hace que JAX sea simple, flexible y fácil de adoptar.
Además de su API NumPy, JAX incluye un sistema extensible de transformaciones de funciones componibles que ayudan a respaldar la investigación sobre aprendizaje automático, que incluyen:
- Diferenciación: La optimización basada en gradientes es fundamental para el aprendizaje automático. JAX admite de forma nativa el modo directo e inverso diferenciación automática de funciones numéricas arbitrarias, mediante transformaciones de funciones como grad, hessian, jacfwd y jacrev.
- Vectorización: En la investigación de ML a menudo aplicamos una única función a muchos datos, por ejemplo, calculando la pérdida en un lote o evaluación de gradientes por ejemplo para el aprendizaje diferencialmente privado. JAX proporciona vectorización automática a través de la transformación vmap que simplifica esta forma de programación. Por ejemplo, los investigadores no necesitan razonar sobre el procesamiento por lotes al implementar nuevos algoritmos. JAX también admite el paralelismo de datos a gran escala a través de la transformación pmap relacionada, distribuyendo elegantemente datos que son demasiado grandes para la memoria de un solo acelerador.
- Compilación JIT: XLA se utiliza para compilar y ejecutar programas JAX justo a tiempo (JIT) en GPU y TPU en la nube aceleradores. La compilación JIT, junto con la API compatible con NumPy de JAX, permite a los investigadores sin experiencia previa en informática de alto rendimiento escalar fácilmente a uno o varios aceleradores.
Hemos descubierto que JAX ha permitido una experimentación rápida con algoritmos y arquitecturas novedosos y ahora sustenta muchas de nuestras publicaciones recientes. Para obtener más información, considere unirse a nuestra mesa redonda JAX, el miércoles 9 de diciembre a las 7:00 p. m. GMT, en el NeurIPS conferencia virtual.
JAX en DeepMind
Apoyar la investigación de IA de última generación significa equilibrar la creación rápida de prototipos y la iteración rápida con la capacidad de implementar experimentos a una escala tradicionalmente asociada con los sistemas de producción. Lo que hace que este tipo de proyectos sea particularmente desafiante es que el panorama de la investigación evoluciona rápidamente y es difícil de pronosticar. En cualquier momento, un nuevo avance en la investigación puede cambiar, y regularmente lo hace, la trayectoria y los requisitos de equipos enteros. Dentro de este panorama en constante cambio, una responsabilidad principal de nuestro equipo de ingeniería es garantizar que las lecciones aprendidas y el código escrito para un proyecto de investigación se reutilicen de manera efectiva en el siguiente.
Un enfoque que ha demostrado ser exitoso es la modularización: extraemos los componentes más importantes y críticos desarrollados en cada proyecto de investigación en sistemas bien probados y eficientes. componentes. Esto permite a los investigadores centrarse en su investigación y, al mismo tiempo, beneficiarse de la reutilización de código, la corrección de errores y las mejoras de rendimiento en los ingredientes algorítmicos implementados por nuestras bibliotecas principales. También hemos descubierto que es importante asegurarnos de que cada biblioteca tenga un alcance claramente definido y que sean interoperables pero independientes. Compra incrementalla capacidad de elegir funciones sin estar limitado a otras, es fundamental para brindar la máxima flexibilidad a los investigadores y apoyarlos siempre en la elección de la herramienta adecuada para el trabajo.
Otras consideraciones que se han incluido en el desarrollo de nuestro ecosistema JAX incluyen asegurarse de que siga siendo coherente (cuando sea posible) con el diseño de nuestro ecosistema existente. TensorFlow bibliotecas (por ejemplo Soneto y TRFL). También hemos intentado crear componentes que (cuando sea relevante) coincidan lo más posible con sus matemáticas subyacentes, para que sean autodescriptivos y minimicen los saltos mentales “del papel al código”. Finalmente, hemos elegido fuente abierta nuestras bibliotecas para facilitar el intercambio de resultados de investigación y alentar a la comunidad en general a explorar el ecosistema JAX.
Nuestro ecosistema hoy
haikus
El modelo de programación JAX de transformaciones de funciones componibles puede complicar el manejo de objetos con estado, por ejemplo, redes neuronales con parámetros entrenables. Haiku es una biblioteca de redes neuronales que permite a los usuarios utilizar modelos familiares de programación orientada a objetos mientras aprovecha el poder y la simplicidad del paradigma funcional puro de JAX.
Haiku es utilizado activamente por cientos de investigadores de DeepMind y Google, y ya ha sido adoptado en varios proyectos externos (p. ej. Engatusar, química profunda, NumPyro). Se basa en la API para Sonetonuestro modelo de programación basado en módulos para redes neuronales en TensorFlow, y nuestro objetivo es hacer que la migración de Sonnet a Haiku sea lo más simple posible.
Obtenga más información en GitHub
optax
La optimización basada en gradientes es fundamental para el aprendizaje automático. Optax proporciona una biblioteca de transformaciones de gradiente, junto con operadores de composición (por ejemplo, cadena) que permiten implementar muchos optimizadores estándar (por ejemplo, RMSProp o Adam) en una sola línea de código.
La naturaleza composicional de Optax admite naturalmente la recombinación de los mismos ingredientes básicos en optimizadores personalizados. Además, ofrece una serie de utilidades para la estimación de gradientes estocásticos y la optimización de segundo orden.
Muchos usuarios de Optax han adoptado Haiku pero, de acuerdo con nuestra filosofía de aceptación incremental, se admite cualquier biblioteca que represente parámetros como estructuras de árbol JAX (p. ej. Elegía, Linaza y estax). Por favor mira aquí para obtener más información sobre este rico ecosistema de bibliotecas JAX.
Obtenga más información en GitHub
RLax
Muchos de nuestros proyectos más exitosos se encuentran en la intersección del aprendizaje profundo y el aprendizaje por refuerzo (RL), también conocido como aprendizaje por refuerzo profundo. RLax es una biblioteca que proporciona componentes básicos útiles para construir agentes RL.
Los componentes de RLax cubren un amplio espectro de algoritmos e ideas: TD-learning, gradientes de políticas, actores críticos, MAP, optimización de políticas próximas, transformación de valor no lineal, funciones de valor generales y una serie de métodos de exploración.
Aunque algunas introducciones agentes de ejemplo se proporcionan, RLax no pretende ser un marco para construir e implementar sistemas completos de agentes RL. Un ejemplo de un marco de agente con todas las funciones que se basa en componentes de RLax es Cumbre.
Obtenga más información en GitHub
chex
Las pruebas son fundamentales para la confiabilidad del software y el código de investigación no es una excepción. Sacar conclusiones científicas de experimentos de investigación requiere tener confianza en la exactitud de su código. Chex es una colección de utilidades de prueba utilizadas por los autores de bibliotecas para verificar que los componentes básicos comunes sean correctos y sólidos y por los usuarios finales para verificar su código experimental.
Chex proporciona una variedad de utilidades que incluyen pruebas unitarias compatibles con JAX, afirmaciones de propiedades de tipos de datos JAX, simulaciones y falsificaciones, y entornos de prueba multidispositivo. Chex se utiliza en todo el ecosistema JAX de DeepMind y en proyectos externos como Engatusar y MineRL.
Obtenga más información en GitHub
Jraph
Graficar redes neuronales (GNN) son un área de investigación apasionante con muchas aplicaciones prometedoras. Véase, por ejemplo, nuestro trabajo reciente sobre predicción de tráfico en Google Maps y nuestro trabajo en simulación de física. Jraph (pronunciado “jirafa”) es una biblioteca liviana que admite el trabajo con GNN en JAX.
Jraph proporciona una estructura de datos estandarizada para gráficos, un conjunto de utilidades para trabajar con gráficos y un “zoológico” de modelos de redes neuronales de gráficos fácilmente bifurcables y extensibles. Otras características clave incluyen: procesamiento por lotes de GraphTuples que aprovechan eficientemente los aceleradores de hardware, soporte de compilación JIT de gráficos de forma variable mediante relleno y enmascaramiento, y pérdidas definidas en particiones de entrada. Al igual que Optax y nuestras otras bibliotecas, Jraph no impone restricciones a la elección del usuario de una biblioteca de redes neuronales.
Obtenga más información sobre el uso de la biblioteca de nuestra rica colección de ejemplos.
Obtenga más información en GitHub
Nuestro ecosistema JAX está en constante evolución y alentamos a la comunidad de investigación de ML a explorar nuestras bibliotecas y el potencial de JAX para acelerar su propia investigación.