La expresividad, la entrenabilidad y la generalización en el aprendizaje automático

Esta es una traducción de mi publicación anterior. Muchas gracias a Joaquin Antonio Ruales (@_jruales) por traducir esta publicación.

Cuando leo trabajos de investigación sobre el aprendizaje automático, intento clasificarlos en base al tipo de contribuciones que hacen: contribuciones en 1) Expresividad 2) Entrenabilidad y/o 3) Generalización. Aprendí este sistema de categorización de mi colega Jascha Sohl-Dickstein en Google Brain, y esta terminología también es propuesta por este trabajo de investigación. Esta categorización me ha resultado efectiva cuando trato de entender de qué manera cierto trabajo de investigación (especialmente si es teórico) conecta varios sub-campos de la investigación en inteligencia artificial (por ej.: robótica, modelos generadores, procesamiento de lenguas naturales) en un solo panorama general [1].

En esta entrada de blog, discuto la aplicación de estos conceptos a la investigación actual (noviembre de 2017) en el área del aprendizaje supervisado, aprendizaje no supervisado, y aprendizaje por refuerzo. Considero que la Generalización se compone de dos categorías — generalización "débil" y "fuerte" — y discutiré a cada una por separado. La siguiente tabla resume mi opinión de dónde están las cosas al momento:

	Expresividad ¿Qué computaciones puede hacer este modelo?	Entrenabilidad ¿Qué tan fácil es ajustar el modelo a los datos?	Generalización débil ¿Qué tan bien funciona el modelo cuando ocurren pequeñas perturbaciones en la distribución de los datos?	Generalización fuerte ¿Qué tan bien funciona el modelo cuando ocurren grandes cambios en la distribución de los datos?
Aprendizaje supervisado Aprende usando etiquetas creadas por humanos. Predice estas etiquetas	Fácil	OK	Ya mismo	Difícil
Aprendizaje no supervisado Aprende sin necesidad de etiquetas creadas por humanos. Predice estas etiquetas.	Fácil	OK	Ya mismo	Difícil
Aprendizaje por refuerzo (RL) Aprende a maximizar la recompensa en un entorno	Fácil	Difícil	Dificilísimo	Dificilísimo

Quisiera expresar mi más profundo agradecimiento a Jascha Sohl-Dickstein y a Ben Poole por sus comentarios y correcciones a esta entrada de blog, así como a Marcin Moczulski por las útiles discusiones sobre la entrenabilidad de las redes neuronales recurrentes.

Esta entrada de blog cubre un amplio rango de investigación desde mi punto de vista en particular. Cabe recalcar que cualquier error factual aquí es mío y no representa las opiniones de mis colegas o correctores. No duden en contribuir su opinión en los comentario o en enviarme un email si quieren discutir el artículo o sugerir correcciones — estoy aquí para aprender.

Expresividad

¿Qué computaciones puede hacer este modelo?

Existen medidas de expresividad que caracterizan lo complejas que pueden llegar a ser las funciones computadas por una red neuronal u otra función paramétrica. Las redes neuronales son exponencialmente expresivas con respecto a su profundidad, lo que significa que incluso las redes neuronales de tamaño moderado son lo suficientemente expresivas para la mayoría de problemas de aprendizaje supervisado, no supervisado y de refuerzo que se investigan hoy en día [2]. Una demostración de esto es que las redes neuronales profundas son capaces de memorizar conjuntos de datos muy grandes.

Las redes neuronales pueden expresar todo tipo de cosas: variables continuas, complejas, discretas e incluso estocásticas. Gracias a la investigación en modelado generativo y en aprendizaje profundo bayesiano en los últimos años, las redes neuronales profundas han producido increíbles resultados de modelado generativo.

Los recientes avances en el modelado generativo ilustran la profunda expresividad de las redes neuronales: las redes neuronales pueden producir geometrías multidimensionales de datos extremadamente complejas (audio, imágenes) que son casi indistinguibles de las reales. Aquí están unos resultados de generación de imágenes de caras usando una arquitectura de red generativa antagónica (GAN) reciente, propuesta por investigadores en NVIDIA:

Los resultados aún no son perfectos (nótese el fondo borroso), pero estamos muy cerca. Así mismo, en el campo de la síntesis de audio, las muestras de audio generadas por los más recientes modelos WaveNet suenan a gente real.

El aprendizaje no supervisado no se limita a los modelos generadores. Algunos investigadores, como Yann LeCun, quieren renombrar al aprendizaje no supervisado "aprendizaje predictivo", donde el modelo deduce el pasado, imputa el presente o predice el futuro. Sin embargo, dado que bastante del aprendizaje no supervisado se ha enfocado en predecir distribuciones conjuntas extremadamente complejas (imágenes del pasado/futuro, audio), creo que los modelos generadores son un punto de referencia bastante bueno para medir la expresividad en el dominio no supervisado.

Las redes neuronales también parecen ser lo suficientemente expresivas como para el aprendizaje por refuerzo. Incluso redes bastante pequeñas (2 capas convolucionales y 2 capas totalmente conectadas) son lo suficientemente potentes como para resolver las tareas de control de Atari y MuJoCo (aunque su "entrenabilidad" deja mucho que desear — consúltese la siguiente sección).

Por sí sola, la expresividad no es un asunto muy interesante: es simple aumentarla agregando más capas, más conectividad, etc. El desafío es hacer que las redes neuronales sean lo suficientemente expresivas como para tener un buen desempeño en el conjunto de entrenamiento y el de pruebas, y al mismo tiempo evitar que sea muy difícil el entrenamiento. Por ejemplo, el uso de convoluciones 2D parece ser necesario para que los modelos de clasificación de imágenes logren generalizar, a pesar de que una red profunda totalmente conectada tiene la capacidad suficiente para memorizar el conjunto de entrenamiento.

La expresividad es el problema más fácil (solo añade más capas!), pero es al mismo tiempo el más misterioso: no tenemos una buena manera de medir cuánta expresividad (y de qué tipo) se requiere en una tarea en particular. ¿Qué tipos de problemas necesitarán redes neuronales que sean varios órdenes de magnitud más grandes que las que se usan hoy en día? Por qué necesitarán tanta potencia de cómputo? ¿Será que las redes neuronales actuales ya son lo suficientemente expresivas para implementar inteligencia parecida a la humana? ¿Acaso resolver la generalización para nuevos y más difíciles problemas requiere modelos considerablemente más expresivos?

El cerebro tiene muchos órdenes de magnitud más "unidades neuronales" (1e11) que el número de unidades en nuestras redes neuronales grandes (Inception-ResNet-V2 tiene aproximadamente 25e6 unidades ReLU). Esta comparación ya es drástica, incluso cuando le damos a las neuronas artificiales ReLU el beneficio de la duda de que sean remotamente comparables a las neuronas biológicas. Incluso una sola neurona biológica (con sus diversos neurotransmisores, bosques dendríticos que integran la señal de otras 10,000 neuronas en una integración variable en el tiempo) también es increíblemente expresiva. Una langosta (el insecto) implementa su sistema de detección de colisiones con tan solo una neurona y ya vuela mejor que cualquier sistema de drones que hayamos construido. ¿De dónde viene toda esta expresividad y a dónde va? ¿Cuánta más expresividad necesitamos?

Entrenabilidad

Una vez encontrada una familia de modelos matemáticos lo suficientemente expresiva, ¿qué tan fácil es encontrar un buen modelo?

En el campo del aprendizaje automático, un modelo es cualquier programa de computadora cuya funcionalidad es aprendida en base a datos. Durante el "aprendizaje", buscamos un modelo razonablemente bueno que utilice los datos para tomar decisiones, de entre un espacio de modelos (potencialmente enorme). Este proceso de búsqueda usualmente se formula como la solución a un problema de optimización sobre el espacio de modelos.

Varios tipos de optimización

Una estrategia común, especialmente en el aprendizaje profundo, es establecer una métrica escalar que evalúe la "bondad" de un modelo (es decir, qué tan bueno es su desempeño). Después usamos una técnica de optimización numérica para maximizar tal "bondad" (o equivalentemente, "minimizar la maldad").

Un ejemplo concreto: la minimización del error de entropía cruzada es una manera estándar de entrenar redes neuronales para que clasifiquen imágenes. Esto se hace con la esperanza de que cuando el modelo haya minimizado la pérdida de entropía cruzada con respecto al conjunto de datos, también logrará lo que queremos, por ejemplo clasificar imágenes correctamente con cierto grado de precisión y exhaustividad en las imágenes de prueba.

La búsqueda de buenos modelos (el entrenamiento) no es, a la final, nada más que optimización — ¡no hay lugar a dudas! Sin embargo... los objetivos de optimización pueden ser difíciles de especificar. Un escenario clásico en el aprendizaje supervisado es el submuestreo (sub-resolución) de imágenes, en donde es difícil definir una sola métrica escalar que esté alineada exactamente con cómo los humanos perciben la "pérdida perceptual" de un algoritmo de submuestreo en particular. Asimismo, la super-resolución y la síntesis de imágenes también resultan difíciles, porque es difícil definir una fórmula para la "bondad" del modelo como un objetivo de maximización. ¡Imagínate descubrir una función matemática que determine de manera objetiva lo "fotorealista" que sea una imágen! El debate de cómo medir la calidad de los modelos generadores (para imágenes, audio) sigue y sigue incluso hasta hoy.

La técnica más popular de los últimos años que se ocupa de este problema es la estrategia de la co-adaptación. En esta estrategia, se formula el problema de optimización como la búsqueda de un punto de equilibrio entre dos distribuciones no estacionarias que evolucionan simultáneamente [3]. Una analogía intuitiva que explica por qué esta estrategia es "natural" es la evolución ecológica de una especie de predadores y una especie de presas. Primero, la especie de predadores se adaptan para poder cazar la presa efectivamente. Después, la especie de presas se adapta para evitar a los predadores. Y así sucesivamente. Las dos especies co-evolucionan como antagonistas el uno del otro, y el resultado es que ambas especies se vuelven más inteligentes.

Las redes generativas antagónicas (GAN) operan acorde a un principio similar para así no tener que explícitamente definir objetivos de pérdida perceptual. Asimismo, el auto-juego competitivo en el aprendizaje por refuerzo emplea este principio para aprender comportamientos complejos. Aunque el objetivo de optimización ahora tiene una especificación implícita, sigue siendo un problema de optimización, y los practicantes del aprendizaje automático pueden reutilizar herramientas familiares como las redes neuronales y el método estocástico del descenso más rápido (SGD, Stochastic Gradient Descent).

Las estrategias evolutivas típicamente aplican un método de optimización-como-simulación. El usuario especifica un sistema dinámico sobre una población de modelos, y en cada paso temporal de la simulación, la población se actualiza de acuerdo a las reglas del sistema dinámico. Los modelos pueden o no interactuar el uno con el otro. La simulación corre adelante en el tiempo, y se espera que la dinámica del sistema cause que la población eventualmente converja a un buen modelo.

Para un excelente tutorial sobre las estrategias evolutivas en el contexto del aprendizaje por refuerzo, véase el artículo A Visual Guide to Evolution Strategies (Una guía visual de las estrategias evolutivas) por David Ha.

La situación actual

La entrenabilidad ya básicamente ha sido resuelta para las redes neuronales prealimentadas (feedforward) y los objetivos de optimización "directos" que son comunes en el aprendizaje supervisado (afirmo esto empíricamente, no teóricamente). Algunos de los grandes avances tecnológicos publicados en el 2015 (normalización por lotes (Batch Norm), redes residuales (ResNets), buena inicialización) se usan ampliamente hoy en día para que las redes prealimentadas sean super entrenables. De hecho, hay redes neuronales con cientas de capas que logran minimizar el error de entrenamiento a cero en conjuntos de datos grandes. Véase éste trabajo de investigación para ver una síntesis completa de la infraestructura algorítmica y de hardware usadas por las redes neuronales profundas en la actualidad.

Las redes neuronales recurrentes todavía resultan un poco difíciles de entrenar, pero la comunidad de investigación sigue avanzando en el área, al punto que ya no se considera loco usar una red neuronal de gran memoria de corto plazo (LSTM) en un sistema robótico complejo y esperar que "simplemente funcione". Esto es increíble, considerando que en el 2014 no muchos confiaban en la entrenabilidad de las redes neuronales recurrentes, y en las décadas previas hubo un montón de trabajos que demostraban que las redes recurrentes eran horribles para entrenar. Existe evidencia que una gran cantidad de las arquitecturas de redes neuronales recurrentes son equivalentemente expresiva la una a la otra, y que cualquier diferencia en desempeño es a causa de que unas arquitecturas son más fáciles de entrenar que otras [4].

En el aprendizaje no supervisado, los datos de salida de los modelos suelen ser (¡no siempre!) más grandes — por ejemplo, 1024 x 1024 píxeles para imágenes, o secuencias gigantescas para la generación de habla y texto. Desafortunadamente, esto los hace más difíciles de entrenar.

Uno de los avances importantes del 2017 fue que la redes generativas antagónicas (GANs) se volvieron más fáciles de entrenar. Las mejoras más populares han sido simples modificaciones al objetivo de divergencia de Jensen-Shannon: mínimos cuadrados, desviación absoluta con margen y la desviación de Wasserstein (1, 2). El trabajo reciente de NVIDIA extiende el trabajo de las GANs basadas en Wasserstein, haciendo que el modelo sea menos sensible a una gran variedad de hiperparámetros, tales como los parámetros de normalización por lotes (BatchNorm), de elección de arquitectura, etc. La estabilidad es increíblemente importante para aplicaciones prácticas e industriales — es lo que nos da confianza de que un modelo será compatible con nuestras ideas de investigación o aplicaciones a futuro. Estos resultados son emocionantes porque indican que nuestras redes neuronales generadoras son lo suficientemente expresivas para generar las imágenes correctas, y el obstáculo que las detenía en el pasado era uno de entrenabilidad. Es posible que la entrenabilidad siga siendo un obstáculo — desafortunadamente con las redes neuronales es difícil saber si es que un modelo no es lo suficientemente expresivo y/o si es que no lo hemos entrenado lo suficiente.

La inferencia con variables latentes discretas dentro de las redes neuronales también solía ser difícil de entrenar a causa de la gran varianza producida por los estimadores de gradiente con técnicas Monte Carlo. Sin embargo, este tipo de inferencia ha resurgido en estos últimos años en todo tipo de arquitecturas, desde redes generativas antagónicas (GANs) a modelos de lenguajes, a redes neuronales con memoria aumentada, a aprendizaje por refuerzo. Las representaciones discretas son muy útiles desde un punto de vista de expresividad y es excelente que ahora podamos entrenarlas de manera bastante fiable.

Trágicamente, el aprendizaje profundo por refuerzo sigue bastante atrasado cuando se trata de la pura entrenabilidad, sin siquiera considerar el aspecto de la generalización. El aprendizaje por refuerzo es difícil porque para ambientes con más de 1 paso temporal, estamos buscando un modelo que realice la optimización (de la recompensa) en el momento de la inferencia. Hay un proceso interno de optimización en el cual el modelo induce el control óptimo, y un bucle externo de optimización que aprende este modelo óptimo, usando solo una base de datos de lo que el agente vio.

Hace poco, añadí una dimensión adicional a una tarea de control robótico contínuo, y el desempeño de mi algoritmo de aprendizaje por refuerzo de derrumbó de >80% a 10%. ¡El aprendizaje por refuerzo no es solo difícil, sino que también es poco fiable! Como el paisaje matemático de optimización es tan estocástico, no podemos siquiera obtener el mismo resultado cuando usamos diferentes semillas para el generador de números aleatorios. Por lo tanto, la estrategia es reportar la distribución de curvas de recompensa a lo largo de varios intentos, cada uno con diferente semilla aleatoria. Diferentes implementaciones del mismo algoritmo varían en desempeño en diferentes entornos, así que las métricas reportadas en la literatura deben ser tomadas con cautela.

¡Qué ridiculez! La entrenabilidad en el aprendizaje por refuerzo definitivamente sigue sin resolverse, porque aún no podemos aplicar nuestro modelo a un problema de escala un poquito más grande y esperar que el mismo aprendizaje haga la misma cosa 10 sobre 10 veces.

Si tratamos al aprendizaje por refuerzo como a un problema puramente de optimización (y lidiamos con la generalización y las tareas complejas más luego), el problema sigue siendo intratable. Digamos que hay un ambiente en donde una recompensa dispersa se da solo al final de un episodio (por ej. el cuidado de niños en el cual se le paga a la niñera cuando vuelven los padres a casa). El número de posibles acciones (y los resultados correspondientes) crecen exponencialmente en relación a la duración del episodio, pero solo pocas de esas secuencias de acciones corresponden al éxito.

Por lo tanto, estimar la gradiente de política en un punto en particular en el panorama matemático de optimización de modelos requiere una cantidad exponencial de muestras con respecto al espacio de acciones antes de que se obtenga una señal útil de aprendizaje. Esto es tan malo como tratar de calcular una expectativa Monte Carlo de una distribución de probabilidad (sobre todas las secuencias de acciones) en donde la masa de la distribución se concentra en una "distribución delta de Dirac" (ver diagrama a continuación). Cuando hay una superposición desvanecedora entre la distribución de propuestas (exploración) y la distribución de recompensas, los estimados de la gradiente de política en base a métodos Monte Carlo de muestras finitas simplemente no funcionarán, sin importar cuántas muestras se recolecten.

Además, si la distribución de datos es no estacionaria (tal como en el caso de los algoritmos fuera-de-política con un buffer de reproducción), recolectar "malos datos" puede introducir bucles de retroalimentación inestables en el proceso externo de optimización.

Aquí está la misma idea pero desde la perspectiva de la optimización en vez de de la estimación Monte Carlo: sin ninguna distribución a priori sobre el espacio de estados (tales como un entendimiento del mundo o una instrucción explícita proveída al agente), el panorama matemático de la optimización se ve como "queso suizo" — pequeños hoyos de soluciones óptimas convexas rodeadas de vastos altiplanos de espacio de parámetros en los cuales la información del gradiente de política es inútil. Esto significa que el espacio completo del modelo tiene básicamente cero información (porque la señal de aprendizaje es efectivamente uniforme a lo largo del espacio de modelos, a causa de que el área que no es cero es extremadamente pequeña).





Si no tenemos buenas representaciones sobre las cuales entrenar el modelo, nos resulta mejor ir ciclando por semillas aleatorios y tomando muestras aleatorias de políticas hasta que descubramos un buen modelo que aterrice en uno de estos hoyos de queso suizo. El hecho que esta técnica aleatoria sea sorprendentemente buena en el aprendizaje por refuerzo sugiere que es muy posible que nuestros paisajes de optimización se vean así.

Según yo, los entornos comparativos como Atari y MuJoCo del aprendizaje por refuerzo, aunque sean interesantes desde el punto de vista de la optimización, en verdad no van hasta el límite del aprendizaje automático, sino más bien simplemente resuelven una política monolítica que optimiza el desempeño en cierto entorno estéril. Hay poca presión selectiva que promueva que el modelo "generalice", así que el problema se vuelve uno de pura optimización, y no un problema del aprendizaje automático en verdad difícil.

No es que quiera complicar los problemas del aprendizaje por refuerzo con asuntos de generalización (¡ciertamente no hará las cosas fáciles de depurar!), pero creo que el aprendizaje de ambos—el entorno y la tarea—es la única manera que hará que el aprendizaje por refuerzo sea manejable para los problemas de robótica del mundo real.

Contrasta esto con el aprendizaje supervisado y el no supervisado, en los cuales podemos obtener señales de aprendizaje de manera barata, sin importar dónde estemos en el espacio de búsqueda de modelos. La distribución de propuesta para gradientes de mini-lotes tiene algo en común con la distribución de gradientes. Si estamos usando el método estocástico del descenso más rápido (Stochastic Gradient Descent) con mini-lotes de tamaño=1, entonces la probabilidad de muestrear la transición con una señal útil de aprendizaje es en el peor de los casos 1/N, en donde N es el tamaño del conjunto de datos (así que está garantizado que se aprenderá algo después de cada época de entrenamiento). Podemos llegar a una buena solución con la fuerza bruta simplemente echando al problema un montón de datos y potencia de cálculo. Además, mejorar la generalización perceptual de las capas más bajas puede de hecho tener un efecto de reducción de la varianza usando "boostrapping" encima de las características (features) de las capas de nivel bajo.

Para resolver los problemas de aprendizaje por refuerzo con alta dimensionalidad y complejidad, debemos primero considerar la generalización y el entendimiento perceptual general antes de atacar el problema de la optimización numérica. Intentar algo diferente es un gasto insensato de esfuerzo. Debemos llegar al punto en el que cada punto de datos provea un número positivo de bits al algoritmo de aprendizaje por refuerzo, y necesitamos tener una manera fácil de hacer un muestreo por importancia de gradientes cuando la tarea sea muy complicada (sin tener que recolectar exponencialmente más datos). Solamente entonces será razonable asumir que podemos tener éxito usando la fuerza bruta en este problema.

El aprendizaje por demostración, por imitación y por refuerzo inverso, y la interfaz con instrucciones de lenguaje natural quizá provean maneras de llevar la política inicial rápidamente a un punto que obtenga algo de señal de aprendizaje (por ejemplo, el entorno otorga una recompensa de 0 pero las observaciones ceden un poco de información útil para los sesgos inductivos del módulo de planificación del modelo).

En resumen: la entrenabilidad es: fácil para el aprendizaje supervisado, difícil-pero-ya-mismo-llega para el aprendizaje no supervisado, y sigue siendo terriblemente difícil para el aprendizaje por refuerzo.

Generalización

La generalización es el más profundo de los 3 problemas, y es la esencia misma del aprendizaje automático. En términos generales, es qué tan bueno es el desempeño del modelo en el conjunto de datos de prueba cuando ha sido entrenado en un conjunto de datos de entrenamiento.

Cuando se habla de la generalización, se distinguen dos escenarios: 1) los conjuntos de datos de entrenamiento y de prueba son muestreados de la misma distribución (y debemos aprender esta distribución solo en base a los datos de entrenamiento) o 2) los datos de entrenamiento y de prueba son muestreados de distribuciones distintas (y necesitamos generalizar de los datos de entrenamiento a los de prueba). Llamemos al (1) "generalización débil" y al (2) "generalización fuerte". Esta manera de clasificar la generalización puede también usar la nomenclatura "interpolación y extrapolación", o "robustez y entendimiento".

Generalización débil: Dos paisajes

¿Qué tan bien funciona el modelo cuando ocurren pequeñas perturbaciones en la distribución de los datos?

En la "generalización débil", asumimos que las muestras del conjunto de entrenamiento y del conjunto de prueba provienen de la misma distribución. Sin embargo, en el mundo real, casi siempre hay alguna diferencia entre la distribución de entrenamiento y la de prueba, incluso en el límite de un gran número de muestras.

Estas diferencias pueden venir del ruido del sensor, del deterioro gradual de los objetos, de variaciones en iluminación (puede que haya estado nublado el día en el que el fotógrafo tomó las imágenes del conjunto de prueba). Otra situación es que las diferencias pueden ser generadas por un antagonista. Las perturbaciones antagónicas son casi imperceptibles al ojo humano, así que podríamos considerar a los ejemplos antagónicos como que fueran muestras de la "misma distribución".

Consecuentemente, es útil en la práctica pensar sobre la "generalización débil" como la evaluación en una versión "perturbada" de la distribución de entrenamiento:

Una distribución perturbada de datos de prueba puede también inducir un paisaje de optimización perturbado (el punto más bajo es el mejor).

El hecho de que no conozcamos la distribución de prueba por adelantado presenta ciertas dificultades para la optimización. Si somos muy agresivos en optimizar el paisaje de entrenamiento (el mínimo *global* puntiagudo de la parte izquierda de la curva azul), terminamos con un modelo que es sub-óptimo con respecto a los datos de prueba (el mínimo *local* puntiagudo de la parte izquierda de la curva roja). Aquí, hemos sobreajustado ("overfitting" en inglés) a la distribución de entrenamiento o a muestras de datos de entrenamiento, y hemos fallado en cuanto a generalizar a la distribución perturbada de prueba.

La "regularización" es una técnica que empleamos para prevenir el sobreajuste. No tenemos ninguna información previa de qué será la perturbación en el conjunto de prueba, así que a menudo lo mejor que podemos hacer es tratar de entrenar en varias versiones aleatoriamente perturbadas de la distribución de entrenamiento, y tener la esperanza que estas incluyan las perturbaciones características de la distribución de prueba. Algunos regularizadores comúnmente usados en el aprendizaje profundo son: el método estocástico del descenso más rápido (Stochastic Gradient Descent), dropout, ruido en los pesos de la red, ruido en las activaciones, y aumento artificial de los datos. En el aprendizaje por refuerzo, es posible aleatorizar los parámetros de simulación para que se vuelva más robusto el entrenamiento. En su charla en ICLR 2017, Chiyuan Zhang comentó que él considera como regularización "cualquier cosa que haga que el entrenamiento sea más difícil" (en vez de la vista convencional de "limitar la capacidad del modelo"). Básicamente, *dificultar* la optimización mejora, de alguna manera, la generalización.

Esto es inquietante — nuestros métodos de "generalizar" son muy rudimentarios, equivalentes a una "lobotomía de optimizadores". Básicamente le lanzamos una llave de tuercas al optimizador con la esperanza que interfiera con el entrenamiento justo lo suficiente como para prevenir el sobreajuste (overfitting). Además, ¡mejorar la entrenabilidad puede hacernos pagar un precio en generalización! Esta manera de ver la generalización (débil) ciertamente complica la manera en la que procedemos con la investigación en el área de la entrenabilidad.

Pero si los mejores optimizadores sobreajustan, entonces ¿cómo explicamos por qué ciertos optimizadores parecen bajar ambos el error de entrenamiento y el de prueba? La realidad es que cualquier combinación de paisaje de optimización y de optimizador causa cierto balance entre 1) encontrar una mejor región de modelos y 2) sobreajustar a una solución específica, y no tenemos buenos métodos para controlar este balance.

La prueba más desafiante de la generalización débil es probablemente la de los ataques antagónicos, en donde las perturbaciones vienen de un "adversario" o "antagonista" que provee la peor posible perturbación para tus datos, de tal manera que con alta probabilidad, tu modelo se desempeñe mal. Aún no tenemos ninguna técnica de aprendizaje automático que sea muy robusta a ataques antagónicos, pero tengo la corazonada de que eventualmente será posible resolver este problema [5].

Hay algunos trabajos teóricos sobre la aplicación de la teoría de la información para mostrar que las redes neuronales aparentemente van a través de una transición de fases durante el proceso de entrenamiento en la cual el modelo cambia de "memorizar" los datos a "comprimirlos". Esta teoría está comenzando a tomar velocidad, pero sigue habiendo un debate académico sobre si es una teoría válida o no. Hay que mantener los ojos abiertos con respecto a esta teoría — la intuición de la "memorización" y la "compresión" es convincente.

Generalización fuerte: superficie geométrica natural

En las pruebas de la "generalización fuerte", el modelo es evaluado en una distribución de datos completamente diferente a la de prueba, pero con datos proviniendo de la misma superficie geométrica subyacente (o proceso generativo)

Cómo puede uno tener la esperanza de aprender un buen modelo para la distribución de prueba si es que es "completamente diferente" a la distribución de entrenamiento? La respuesta a esto es que estos vistazos de datos en verdad provienen de la misma superficie geométrica de "datos naturales". Aún puede haber un montón de información en común entre la distribución de entrenamiento y la de prueba, con tal de que provengan del mismo proceso generativo subyacente.

El espacio de datos observables en el mundo puede ser descrito como una "superficie matemática natural" contínuamente variante. La tupla (videoclip de un par de platillos chocando el uno con el otro, el sonido de los platillos chocando) is un punto en esta superficie matemática. La tupla (videoclip de un par de platillos chocando el uno con el otro, el sonido de ranas croando) no lo es — esos puntos de datos juntos simplemente son inconsistentes con nuestra realidad. Como puedes ver, aunque no hay duda que la superficie matemática es gigantesca, también es altamente estructurada. Por ejemplo las leyes físicas como la gravedad se obedecen en todos los datos que observamos, los objetos no desaparecen y reaparecen repentinamente, y así sucesivamente.

La generalización fuerte puede ser considerada la cantidad de esta "super-superficie" geométrica que es capturada por un modelo en particular que es solo entrenada en una diminuta muestra de puntos de datos de la superficie. Nótese que un clasificador de imágenes no necesita descubrir las ecuaciones de Maxwell — simplemente necesita tener un entendimiento de la realidad que sea consistente con los datos en la superficie geométrica.

Se puede argumentar que los modelos modernos de clasificación entrenados en ImageNet son buenos en la generalización fuerte — los modelos entrenados en ImageNet en verdad entienden conceptos tales como bordes, contornos y objetos, y por eso es que es tan común transferir sus pesos pre-entrenados a otros conjuntos de datos aparte de ImageNet para realizar aprendizaje "few-shot" ("de pocas tomas", en el que se entrena un modelo en muy pocos ejemplos en vez de cientos o más ejemplos) y aprendizaje de métricas. Sin embargo, podrían ser mucho mejores: los clasificadores entrenados en ImageNet no funcionan universalmente, el aprendizaje "few-shot" aún no ha sido resuelto completamente, y estos clasificadores todavía son susceptibles a ejemplos antagónicos. Obviamente nuestros modelos no entienden lo que están viendo de la misma manera en la que los humanos lo hacen. Sin embargo, es un comienzo.

Así como con la generalización débil, la distribución de prueba puede ser muestreada antagónicamente de tal manera que cause la máxima diferencia entre las distribuciones de entrenamiento y de prueba. El programa AlphaGo Zero (un programa que aprende a jugar el juego de mesa Go sin necesidad de observar juegos humanos, y que sin embargo logra vencer a los campeones mundiales de este juego) es mi ejemplo favorito: a la hora de la prueba, este programa observa datos de jugadores humanos completamente diferentes a los datos de su distribución de entrenamiento (nunca ha "visto" a un humano antes). Además, el humano está usando todo su inteligencia para llevar a AlphaGo a regímenes que no hayan sido observados en los datos de prueba. A pesar de que AlphaGo no entiende explícitamente nada sobre matemáticas abstractas, la psicología de oponentes, o qué significa el color verde, claramente entiende lo suficiente acerca del mundo como para ser más listo que un jugador humano en una tarea bastante específica. Si un sistema de inteligencia artificial es robusto en contra a un adversario humano experto, yo considero que tiene la suficiente capabilidad de generalización.

Tristemente, la investigación en el área del aprendizaje por refuerzo ha ignorado en gran medida el problema de la generalización fuerte. La mayoría de las tareas estandarizadas del aprendizaje por refuerzo utilizan ambientes estáticos, con muy poca riqueza perceptual (por ej., el humanoide no entiende su mundo o cómo se ve un cuerpo, sino que simplemente lo interpreta como un conjunto de articulaciones flotantes cuyas posiciones pueden llevar a obtener cierta recompensa).

Yo en verdad creo que resolver la generalización es esencial para resolver la entrenabilidad en el aprendizaje por refuerzo. Mientras más "entiende" el mundo nuestro sistema inteligente, mejor le irá obteniendo señales de aprendizaje, y posiblemente con menos muestras. Por esto es que el aprendizaje "few-shot", el aprendizaje por imitación, y el aprender-a-aprender son importantes: nos alejan de las soluciones de fuerza bruta en donde la varianza es alta y la información es baja.

Yo creo que son necesarias dos cosas para lograr formas más fuertes de generalización:

Primero, necesitamos modelos que activamente deduzcan las leyes fundamentales del mundo en base a la observación y la experimentación. El razonamiento simbólico y la inferencia causal parecen ser temas de estudio con bastante potencial, pero es probable que cualquier tipo de aprendizaje no supervisado pueda ayudar. Esto me recuerda de la larga búsqueda de la humanidad por entender el movimiento de los cuerpos celestes derivando las leyes físicas del universo por medio de un sistema de razonamiento lógico (las matemáticas). Es interesante tomar en cuenta que antes de la Revolución de Copérnico los humanos probablemente dependían al comienzo de heurísticas bayesianas ("supersticiones"), y estos modelos "bayesianos" fueron descartados una vez que descubrimos la mecánica clásica.

Nuestros métodos de aprendizaje automático en base a modelos (modelos que intentan "predecir" aspectos de su ambiente) son al momento "pré-copérnicos" en el sentido que solo interpolan en base a supersticiones estadísticas muy superficiales, en vez de desarrollar principios generales y profundos para explicar y extrapolar datos que pueden estar a millones de años luz de distancia o a muchos pasos temporales en el futuro. Nótese que los humanos en verdad no necesitaron entender la teoría de la probabilidad para derivar la mecánica determinística de los cuerpos celestes, lo cual plantea la pregunta de si habrá cómo usar el aprendizaje automático y la inferencia causal sin necesidad de usar explícitamente la estadística.

Una manera de drásticamente reducir el peso de la complejidad es hacer que nuestros sistemas de aprendizaje automático sean más adaptativos. Necesitamos ir más allá de los modelos de optimización que predigan o actúen en maneras estáticas. En vez de eso, tenemos que optimizar modelos que puedan pensar, recordar, y aprender, y todo esto en tiempo real.

Segundo, necesitamos echarle suficiente diversidad de datos al problema, para empujar al modelo a que desarrolle estas representaciones abstractas. Un ambiente necesita ser lo suficientemente rico como para forzar el desarrollo de la representación correcta (aunque AlphaGo Zero nos hace dudar de qué cantidad de la superficie geométrica natural de datos en verdad es necesaria para el agente). Sin estas restricciones, el problema está fundamentalmente subespecificado y la probabilidad de que accidentalmente descubramos la solución real es nula. Quizá los humanos nunca se hubieran vuelto seres inteligentes si no se hubieran podido parar en sus patas traseras y preguntado por qué las estrellas se mueven en tan raros patrones elípticos.

Me pregunto si la civilización trisolariana (del libro "El problema de los tres cuerpos") evolucionó a tal nivel de avance tecnológico porque su supervivencia dependía de su entendimiento de la mecánica celeste compleja. Quizá necesitemos incluir un poco de movimiento de cuerpos celestes en nuestros ambientes de simulación MuJoCo y Bullet :)

Notas

[1] Hay algunos campos de investigación que no caben tan bien en esta estructura de expresividad, entrenabilidad y generalización. Por ejemplo, la investigación en el campo de la interpretabilidad trata de entender por qué cierto modelo produce una respuesta en particular. Ésta respuesta no solo es exigida, en ciertos campos de alto riesgo (medicina, aplicación de la ley), por las autoridades responsables de las políticas relacionadas al aprendizaje automático y por los clientes del aprendizaje automático, sino que también pueden ayudarnos a entender mejor la generalización: si descubrimos que el modelo determina las diagnosis de manera muy distinta de la cual lo hacen los profesionales médicos humanos, esto podría significar que hay ciertos casos en los que el proceso deductivo del modelo fallará al momento de generalizar.

[2] Un bosquejo de la prueba: Una capa completamente conectada con N neuronas seguida de una no linealidad ReLU puede tajar a un espacio vectorial en N regiones lineales a trozos. La introducción de una segunda capa ReLU subdivide el espacio en N más pedazos, resultando en N^2 regiones lineales a trozos en el espacio de entrada. 3 capas son N^3, y así sucesivamente. Para un análisis más detallado, véase Raghu et al. 2017.

[3] Esto a veces se llama un problema de optimización multi-nivel. Sin embargo, esto implica un bucle de optimización "externo" y uno "interno", mientras que la co-adaptación puede ocurrir simultáneamente. Por ejemplo, procesos concurrentes en una sola máquina comunicándose asincrónicamente, o especies evolucionando contínuamente con respecto el una a la otra en un ecosistema. En estos casos, no hay bucles claros de optimización "externos" e "internos".

[4] Los modelos de secuencia-a-secuencia con atención lograron resultados de punta cuando fueron presentados por primera vez, pero sospecho que la ventaja que otorgan es una de entrenabilidad, no de expresividad o generalización. Es posible que los modelos de secuencia-a-secuencia sin atención obtengan resultados igual de buenos si son inicializados adecuadamente.

[5] Una idea para proteger en contra de los métodos antagónicos, sin resolver la generalización fuerte: hacer que sea extremadamente costoso computar las perturbaciones antagónicas. Los modelos y los datos son parcialmente como cajas negras. En cada llamada al modelo durante el momento de la inferencia, elige un modelo al azar de un conjunto de modelos entrenados, y entrega este modelo al adversario sin decirle qué modelo recibió. Los modelos son entrenados independientemente el uno del otro y pueden adoptar arquitecturas diferentes. Esto hace que sea difícil computar gradientes de diferencias finitas, por lo que f(x+dx) - f(x) puede tener una varianza arbitrariamente grande. Además, los gradientes entre cómputos sucesivos de gradiente aún tendrán una alta varianza porque diferentes pares de modelos pueden ser muestreados. Otra manera de mejorar las cosas sería usar datos multimodales (vídeo, multi-vista, imágenes + sonido), para que sea difícil perturbar entradas en una manera que preserve la consistencia de las entradas.

Normalizing Flows Tutorial, Part 2: Modern Normalizing Flows

This tutorial will show you how to use normalizing flows like MAF, IAF, and Real-NVP to deform an isotropic 2D Gaussian into a complex cloud of points spelling the words "SIGGRAPH" in space. Like stretching taffy.

I'm looking for help translate these posts into different languages! Please email me at <myfirstname><mylastname>2004<at>gmail.com if you are interested. 

Xiaoyi Yin (尹肖贻) has kindly translated this post into Chinese (中文). 

Jaeseong You has kindly translated this post into Korean (한국어)

Kieran Didi has kindly translated this post into German (Deutsch)

In my previous blog post, I described how simple distributions like Gaussians can be “deformed” to fit complex data distributions using normalizing flows. We implemented a simple flow by chaining 2D Affine Bijectors with PreLU nonlinearities to build a small invertible neural net.

However, this MLP flow is pretty weak: there are only 2 units per hidden layer. Furthermore, the non-linearity is monotonic and piecewise linear, so all it does is slightly warp the data manifold around the origin. This flow completely fails to implement more complex transformations like separating an isotropic Gaussian into two modes when trying to learn the “Two Moons” dataset below:

Fortunately, there are several more powerful normalizing flows that have been introduced in recent Machine Learning literature. We will explore several of these techniques in this tutorial.

Autoregressive Models are Normalizing Flows

Autoregressive density estimation techniques like WaveNet and PixelRNN model learn complex joint densities $p(x_{1:D})$ by decomposing the joint density into a product of one-dimensional conditional densities, where each $x_i$ depends on only the previous $i-1$ values:

$$p(x) = \prod_i{p(x_i \,|\, x_{1:i-1})}$$

The conditional densities usually have learnable parameters. For example, a common choice is an autoregressive density $p(x_{1:D})$ whose conditional density is a univariate Gaussian, whose mean and standard deviations are computed by neural networks that depend on the previous $x_{1:i-1}$.

$$p(x_i \,|\, x_{1:i-1}) = \mathcal{N}(x_i  \,|\,\mu_i, (\exp\alpha_i)^2)$$

$$\mu_i = f_{\mu_i}(x_{1:i-1})$$

$$\alpha_i = f_{\alpha_i}(x_{1:i-1})$$

Learning data with autoregressive density estimation makes the rather bold inductive bias that the ordering of variables are such that your earlier variables don’t depend on later variables. Intuitively, this shouldn’t be true at all for natural data (the top row of pixels in an image does have a causal, conditional dependency on the bottom of the image). However it’s still possible to generate plausible images in this manner (to the surprise of many researchers!). 

To sample from this distribution, we compute $D$ “noise variates” $u_{1:D}$ from the standard Normal, $N(0,1)$, then apply the following recursion to get $x_{1:D}$.

$$x_i = u_i\exp{\alpha_i} + \mu_i$$
$$u_i \sim \mathcal{N}(0, 1)$$

The procedure of autoregressive sampling is a deterministic transformation of the underlying noise variates (sampled from $\mathcal{N}(0, \mathbb{I})$) into a new distribution, so autoregressive samples can actually be interpreted as a TransformedDistribution of the standard Normal!

Armed with this insight, we can stack multiple autoregressive transformations into a normalizing flow. The advantage of doing this is that we can change the ordering of variables $x_1,...x_D$ for each bijector in the flow, so that if one autoregressive factorization cannot model a distribution well (due to a poor choice of variable ordering), a subsequent layer might be able to do it.

The Masked Autoregressive Flow (MAF) bijector implements such a conditional-Gaussian autoregressive model. Here is a schematic of the forward pass for a single entry in a sample of the transformed distribution, $x_i$:

The gray unit $x_i$ is the unit we are trying to compute, and the blue units are the values it depends on. $\alpha_i$ and $\mu_i$ are scalars that are computed by passing $x_{1:i-1}$ through neural networks (magenta, orange circles). Even though the transformation is a mere scale-and-shift, the scale and shift can have complex dependencies on previous variables. For the first unit $x_1$, $\mu$ and $\alpha$ are usually set to learnable scalar variables that don’t depend on any $x$ or $u$.

More importantly, the transformation is designed this way so that computing the inverse $u = f^{-1}(x)$ does not require us to invert $f_\alpha$ or $f_\mu$. Because the transformation is parameterized as a scale-and-shift, we can recover the original noise variates by reversing the shift and scale: $u = (x-f_\mu(x))/\exp(f_\alpha(x))$. The forward and inverse pass of the bijector only depend on the forward evaluation of $f_\alpha(x)$ and $f_\mu(x)$, allowing us to use non-invertible functions like ReLU and non-square matrix multiplication in the neural networks $f_\mu$ and $f_\alpha$.

The inverse pass of the MAF model is used to evaluate density:

distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x))

Runtime Complexity and MADE

Autoregressive models and MAF can be trained “quickly” because all conditional likelihoods $p(x_1), p(x_2\,|\, x_1), ... p(x_D\,|\, x_{1:D-1}))$ can be evaluated simultaneously in a single pass of D threads, leveraging the batch parallelism of modern GPUs. We are operating under the assumption that parallelism, such as SIMD vectorization on CPUs/GPUs, has zero runtime overhead.

On the other hand, sampling autoregressive models is slow because you must wait for all previous $x_{1:i-1}$ to be computed before computing new $x_i$. The runtime complexity of generating a single sample is D sequential passes of a single thread, which fails to exploit processor parallelism.

Another issue: in the parallelizable inverse pass, should we use separate neural nets (with differently-sized inputs) for computing each $\alpha_i$ and $\mu_i$? That's inefficient, especially if we consider that learned representations between these D networks should be shared (as long as the autoregressive dependency is not violated). In the Masked Autoencoder for Distribution Estimation (MADE) paper, the authors propose a very nice solution: use a single neural net to output all values of $\alpha$ and $\mu$ simultaneously, but mask the weights so that the autoregressive property is preserved.

This trick makes it possible to recover all values of $u$ from all values of $x$ with a single pass through a single neural network (D inputs, D outputs). This is far more efficient than processing D neural networks simultaneously (D(D+1)/2 inputs, D outputs).

To summarize, MAF uses the MADE architecture as an efficiency trick for computing nonlinear parameters of shift-and-scale autoregressive transformations, and casts these efficient autoregressive models into the normalizing flows framework.

Inverse Autoregressive Flow (IAF)

In Inverse Autoregressive Flow, the nonlinear shift/scale statistics are computed using the previous noise variates $u_{1:i-1}$, instead of the data samples:

$$x_i = u_i\exp{\alpha_i} + \mu_i$$
$$\mu_i = f_{\mu_i}(u_{1:i-1})$$
$$\alpha_i = f_{\alpha_i}(u_{1:i-1})$$

The forward (sampling) pass of IAF is fast: all the $x_i$ can be computed in a single pass of $D$ threads working in parallel. IAF also uses MADE networks to implement this parallelism efficiently.

However, if we are given a new data point and asked to evaluate the density, we need to recover $u$ and this process is slow: first we recover $u_1 = (x-\mu_1) * \exp(-\alpha_1)$, then $u_i = (x-\mu_i(u_{1:i-1})) * \exp(-\alpha_i(u_{1:i-1}))$ sequentially. On the other hand, it’s trivial to track the (log) probability of samples generated by IAF, since we already know all of the $u$ values to begin with without having to invert from $x$.

The astute reader will notice that if you re-label the bottom row as x_1, .. x_D, and the top row as u_1, … u_D, this is exactly equivalent to the Inverse Pass of the MAF bijector! Likewise, the inverse of IAF is nothing more than the forward pass of MAF (with $x$ and $u$ swapped). Therefore in TensorFlow Distributions, MAF and IAF are actually implemented using the exact same Bijector class, and there is a convenient “Invert” feature for inverting Bijectors to swap their inverse and forward passes.

iaf_bijector = tfb.Invert(maf_bijector)

IAF and MAF make opposite computational tradeoffs - MAF trains quickly but samples slowly, while IAF trains slowly but samples quickly. For training neural networks, we usually demand way more throughput with density evaluation than sampling, so MAF is usually a more appropriate choice when learning distributions.

Parallel Wavenet

An obvious follow-up question is whether these two approaches can be combined to get the best of both worlds, i.e. fast training and sampling.

The answer is yes! The much-publicized Parallel Wavenet by DeepMind does exactly this: an autoregressive model (MAF) is used to train a generative model efficiently, then an IAF model is trained to maximize the likelihood of its own samples under this teacher. Recall that with IAF, it is costly to compute density of external data points (such as those from the training set), but it can cheaply compute density of its own samples by caching the noise variates $u_{1:D}$, thereby circumventing the need to call the inverse pass. Thus, we can train the “student” IAF model by minimizing the divergence between the student and teacher distributions.

This is an incredibly impactful application of normalizing flows research - the end result is a real-time audio synthesis model that is 20 times faster to sample, and is already deployed in real-world products like the Google Assistant.

NICE and Real-NVP

Finally, we consider is Real-NVP, which can be thought of as a special case of the IAF bijector.

In a NVP “coupling layer”, we fix an integer $0 < d < D$. Like IAF, $x_{d+1}$ is a shift-and-scale that depends on previous $u_{d}$ values. The difference is that we also force $x_{d+2}, x_{d+3}, … x_{D}$ to only depend on these $u_{d}$ values, so a single network pass can be used to produce $\alpha_{d+1:D}$ and $\mu_{d+1:D}$.

As for $x_1:d$ they are “pass-through” units that are set equivalently to $u_{1:d}$. Therefore, Real-NVP is also a special case of the MAF bijector (since $\alpha(u_{1:d}) = \alpha(x_{1:d})$).

Because the shift-and-scale statistics for the whole layer can be computed from either $x_{1:d}$ or $u_{1:d}$ in a single pass, NVP can perform forward and inverse computations in a single parallel pass (sampling and estimation are both fast). MADE is also not needed.

However, empirical studies suggest that Real-NVP tends to underperform MAF and IAF and my experience has been that NVP tends to fit my toy 2D datasets (e.g. SIGGRAPH dataset) more poorly when using the same number of layers. Real-NVP and IAF are nearly equivalent in the 2D case, except the first unit of IAF is still transformed via a scale-and-shift that does not depend on $u_1$, while Real-NVP leaves the first unit unmodified.

Real-NVP was a follow-up work to the NICE bijector, which is a shift-only variant that assumes $\alpha=0$. Because NICE does not scale the distribution, the ILDJ is actually constant!

Batch Normalization Bijector

The Real-NVP paper proposes several novel contributions, one of which is a Batch Normalization bijector used to stabilize training. Conventionally, Batch Norm is applied to training neural networks where the forward statistics are mean-centered and scaled to diagonal unit covariance, and the batchnorm statistics (running mean, running variance) are accumulated via an exponential moving average. At test time, the accumulated statistics are used to normalize data.

In normalizing flows, batch norm is used in bijector.inverse during training, and the accumulated statistics are used to de-normalize data at “test time” (bijector.forward). Concretely, BatchNorm Bijectors are typically implemented as follows:


Inverse pass:

1. Compute the current mean and standard deviation of the data distribution $x$.
2. Update running mean and standard deviation
3. Batch normalize the data using current mean/std

Forward pass:

1. Use running mean and standard deviation to un-normalize the data distribution. 

Thanks to TF Bijectors, this can be implemented with only a few lines of code:




The ILDJ can be derived easily by simply taking the log derivative of inverse function (consider the univariate case).

Code Example

Thanks to the efforts of Josh Dillon and the Google Bayesflow team, there is already a flexible implementation of MaskedAutoregressiveFlow Bijector that uses MADE networks to implement efficient recovery of $u$ for training.

I’ve created a complex 2D distribution, which is a point cloud in the shape of the letters “SIGGRAPH” using this blender script. We construct our dataset, bijector, and transformed distribution in a very similar fashion to the first tutorial, so I won’t repeat the code snippets here - you can find the Jupyter notebook here. This notebook can train a normalizing flow using MAF, IAF, Real-NVP with/without BatchNorm, for both the "Two Moons" and "SIGGRAPH" datasets.

One detail that’s easy to miss / introduce bugs on is that this doesn’t work at all unless you permute the ordering of variable at each flow. Otherwise, none of the layers’ autoregressive factorization will be learn structure of $p(x1 | x2)$. Fortunately, TensorFlow has a Permute bijector specially made for doing this.





Here’s the learned flow, along with the final result. It reminds me a lot of a taffy pulling machine.

Discussion

TensorFlow distributions makes normalizing flows easy to implement, and automatically accumulate all the Jacobians determinants in a chain for us in a way that is clean and highly readable. When deciding which Normalizing Flow to use, consider the design tradeoff between a fast forward pass and a fast inverse pass, as well as between an expressive flow and a speedy ILJD.

In Part 1 of the tutorial, I motivated Normalizing Flows by saying that we need availability of more powerful distributions that can be used in reinforcement learning and generative modeling. In the big picture of things, it’s not clear whether having volume-tracking normalizing flows is actually the best tool for AI applications like robotics, structured prediction, when techniques like variational inference and implicit density models already work extremely well in practice. Even still, normalizing flows are a neat family of methods to have in your back pocket and they have demonstrable real-world applications, such as in real-time generative audio models deployed on Google Assistant.

Although explicit-density models like normalizing flows are amenable to training via maximum likelihood, this is not the only way they can be used and are complementary to VAEs and GANs. It’s possible to use normalizing flow as a drop-in replacement for anywhere you would use a Gaussian, such as VAE priors and latent codes in GANs. For example, this paper use normalizing flows as flexible variational priors, and the TensorFlow distributions paper presents a VAE that uses a normalizing flow as a prior along with a PixelCNN decoder. Parallel Wavenet trains an IAF "student" model via KL divergence.

One of the most intriguing properties of normalizing flows is that they implement reversible computation (i.e. have a defined inverse of an expressive function). This means that if we want to perform a backprop pass, we can re-compute the forward activation values without having to store them in memory during the forward pass (potentially expensive for large graphs). In a setting where credit assignment may take place over very long time scales, we can use reversible computation to “recover” past decision states while keeping memory usage bounded. In fact, this idea was utilized in the RevNets paper, and was actually inspired by the invertibility of the NICE bijector. I’m reminded of the main character from the film Memento who is unable to store memories, so he uses invertible compute to remember things.

Thank you for reading.

Code on Github

Acknowledgements

I’m grateful to Dustin Tran, Luke Metz, Jonathan Shen, Katherine Lee, and Samy Bengio for proofreading this post.

References and Further Reading

• The content and outline of this blog post was heavily influenced by the Masked Autoregressive Flow for Density Estimation paper, which is very well-written and is more or less my primary source for understanding this topic. Give it a read!
  
• Some earlier work on NFs: https://math.nyu.edu/faculty/tabak/publications/Tabak-Turner.pdf and https://arxiv.org/pdf/1302.5125.pdf and https://arxiv.org/abs/1505.05770
• Talk by Laurent Dinh & discussion by Twitter Cortex researchers. Some neat ideas and discussion here.
• Tutorial on Normalizing Flows using PyMC.
• There’s a body of work that I don’t fully understand yet, bridging Normalizing Flows to Langevin Flow and Hamiltonian Flow. As the number of Bijectors in a normalizing flow goes to infinity, one arrives at a Continuous-Time Flow, which apparently can express even richer transformations. 

Normalizing Flows Tutorial, Part 1: Distributions and Determinants

I'm looking for help translate these posts into different languages! Please email me at <myfirstname><mylastname>2004<at>gmail.com if you are interested.

Xiaoyi Yin (尹肖贻) has kindly translated this post into Chinese (中文).
Jaeseong You has kindly translated this post into Korean (한국어)

Kieran Didi has kindly translated this post into German (Deutsch)

If you are a machine learning practitioner working on generative modeling, Bayesian deep learning, or deep reinforcement learning, normalizing flows are a handy technique to have in your algorithmic toolkit. Normalizing flows transform simple densities (like Gaussians) into rich complex distributions that can be used for generative models, RL, and variational inference. TensorFlow has a nice set of functions that make it easy to build flows and train them to suit real-world data.

This tutorial comes in two parts:

• Part 1: Distributions and Determinants. In this post, I explain how invertible transformations of densities can be used to implement more complex densities, and how these transformations can be chained together to form a “normalizing flow”.
• Part 2: Modern Normalizing Flows: In a follow-up post, I survey recent techniques developed by researchers to learn normalizing flows, and explain how a slew of modern generative modeling techniques -- autoregressive models, MAF, IAF, NICE, Real-NVP, Parallel-Wavenet -- are all related to each other.

This series is written for an audience with a rudimentary understanding of linear algebra, probability, neural networks, and TensorFlow. Knowledge of recent advances in Deep Learning, generative models will be helpful in understanding the motivations and context underlying these techniques, but they are not necessary.

Background

Statistical Machine Learning algorithms try to learn the structure of data by fitting a parametric distribution $p(x; \theta)$ to it. Given a dataset, if we can represent it with a distribution, we can:

1. Generate new data “for free” by sampling from the learned distribution in silico; no need to run the true generative process for the data. This is a useful tool if the data is expensive to generate, i.e. a real-world experiment that takes a long time to run [1]. Sampling is also used to construct estimators of high-dimensional integrals over spaces.
2. Evaluate the likelihood of data observed at test time (this can be used for rejection sampling or to score how good our model is).
3. Find the conditional relationship between variables. For example, learning the distribution $p(x_2 | x_1)$ allows us to build discriminative classification or regression models.
4. Score our algorithm by using complexity measures like entropy, mutual information, and moments of the distribution.

We’ve gotten pretty good at sampling (1), as evidenced by recent work on generative models for images and audio. These kinds of generative models are already being deployed in real commercial applications and Google products. 

However, the research community currently directs less attention towards unconditional & conditional likelihood estimation (2, 3) and model scoring (4). For instance, we don’t know how to compute the support of a GAN decoder (how much of the output space has been assigned nonzero probability by the model), we don’t know how to compute the density of an image with respect to a DRAW distribution or even a VAE, and we don’t know how to analytically compute various metrics (KL, earth-mover distance) on arbitrary distributions, even if we know their analytic densities.

Generating likely samples isn’t enough: we also care about answering “how likely is the data?” [2], having flexible conditional densities (e.g. for sampling/evaluating divergences of multi-modal policies in RL), and being able to choose rich families of priors and posteriors in variational inference. 

Consider for a moment, your friendly neighborhood Normal Distribution. It’s the Chicken Soup of distributions: we can draw samples from it easily, we know its analytic density and KL divergence to other Normal distributions, the central limit theorem gives us confidence that we can apply it to pretty much any data, and we can even backprop through its samples via the reparameterization trick. The Normal Distribution’s ease-of-use makes it a very popular choice for many generative modeling and reinforcement learning algorithms.

Unfortunately, the Normal distribution just doesn’t cut it in many real-world problems we care about. In Reinforcement Learning -- especially continuous control tasks such as robotics -- policies are often modeled as multivariate Gaussians with diagonal covariance matrices. 

By construction, uni-modal Gaussians cannot do well on tasks that require sampling from a multi-modal distribution. A classic example of where uni-modal policies fail is an agent trying to get to its house across a lake. It can get home by circumventing the lake clockwise (left) or counterclockwise (right), but a Gaussian policy is not able to represent two modes. Instead, it chooses actions from a Gaussian whose mean is a linear combination of the two modes, resulting in the agent going straight into the icy water. Sad!

The above example illustrates how the Normal distribution can be overly simplistic. In addition to bad symmetry assumptions, Gaussians have most of their density concentrated at the edges in high dimensions and are not robust to rare events. Can we find a better distribution with the following properties?

1. Complex enough to model rich, multi-modal data distributions like images and value functions in RL environments?
2. … while retaining the easy comforts of a Normal distribution: sampling, density evaluation, and with re-parameterizable samples?

The answer is yes! Here are a few ways to do it:

• Use a mixture model to represent a multi-modal policy, where a categorical represents the “option” and the mixture represents the sub-policy. This provides samples that are easy to sample and evaluate, but samples are not trivially re-parameterizable, which makes them hard to use for VAEs and posterior inference. However, using a Gumbel-Softmax / Concrete relaxation of the categorical “option” would provide a multi-modal, re-parameterizable distribution.
• Autoregressive factorizations of policy / value distributions. In particular,  the Categorical distribution can model any discrete distribution.
• In RL, one can avoid this altogether by symmetry-breaking the value distribution via recurrent policies, noise, or distributional RL. This helps by collapsing the complex value distributions into simpler conditional distributions at each timestep. 
• Learning with energy-based models, a.k.a undirected graphical models with potential functions that eschew an normalized probabilistic interpretation. Here’s a recent example of this applied to RL.
• Normalizing Flows: learn invertible, volume-tracking transformations of distributions that we can manipulate easily.

Let's explore the last approach - Normalizing Flows.

Change of Variables, Change of Volume

Let's build up some intuition by examining linear transformations of 1D random variables. Let $X$ be the distribution $\text{Uniform}(0,1)$. Let random variable $Y = f(X) = 2X + 1$. $Y$ is a simple affine (scale & shift) transformation of the underlying “source distribution” $X$. What this means is that a sample $x^i$ from $X$ can be converted into a sample from $Y$ by simply applying the function $f$ to it. 

The green square represents the shaded probability mass on $\mathbb{R}$ for both $p(x)$ and $p(y)$ - the height represents the density function at that value. Observe that because probability mass must integrate to 1 for any distribution, the act of scaling the domain by 2 everywhere means we must divide the probability density by 2 everywhere, so that the total area of the green square and blue rectangle are the same (=1).

If we zoom in on a particular x and an infinitesimally nearby point $x+dx$, then applying f to them takes us to the pair $(y, y+dy)$.

On the left, we have a locally increasing function ($dy/dx > 0$) and on the right, a locally decreasing function ($dy/dx < 0$). In order to preserve total probability, the change of $p(x)$ along interval $dx$ must be equivalent to the change of $p(y)$ along interval $dy$:

$$p(x) dx = p(y) dy$$

In order to conserve probability, we only care about the amount of change in y and not its direction (it doesn’t matter if $f(x)$ is increasing or decreasing at x, we assume the amount of change in y is the same regardless). Therefore, $p(y) = p(x) | dx/dy |$. Note that in log-space, this is equivalent to $\log p(y) = \log p(x) + \log | dx/dy |$. Computing log-densities is more well-scaled for numerical stability reasons.

Now let’s consider the multivariate case, with 2 variables. Again, zooming into an infinitesimally small region of our domain, our initial “segment” of the base distribution is now a square with width dx.

Note that a transformation that merely shifts a rectangular patch $(x1,x2, x3,x4)$ does not change the area. We are only interested in the rate of change per unit area of x, so the displacement $dx$ can be thought of as a unit of measure, which is arbitrary. To make the following analysis simple and unit-less, let’s investigate a unit square on the origin, i.e. 4 points $(0,0), (1,0), (0,1), (1,1)$.

Multiplying this by the matrix $[[a, b];[c, d]]$ will take points on this square into a parallelogram, as shown on the figure to the right (below). $(0,0)$ is sent to $(0,0)$, $(1,0)$ is sent to $(a,b)$, $(1,0)$ sent to $(c,d)$, $(1,1)$ sent to $(a+c,b+d)$.

Thus, a unit square in the domain of $X$ corresponds to a deformed parallelogram in the domain of $Y$, so the per-unit rate of change in area is the area of the parallelogram, i.e. $ad - bc$.

The area of a parallelogram, $ad - bc$, is nothing more than the absolute value of the determinant of the linear transformation!

In 3 dimensions, the “change in area of parallelogram” becomes a “change in volume of parallelpiped”, and even higher dimensions, this becomes “change in volume of a n-parallelotope”. But the concept remains the same - determinants are nothing more than the amount (and direction) of volume distortion of a linear transformation, generalized to any number of dimensions.

What if the transformation f is nonlinear? Instead of a single parallelogram that tracks the distortion of any point in space, you can picture many infinitesimally small parallelograms corresponding to the amount of volume distortion for each point in the domain. Mathematically, this locally-linear change in volume is $|\text{det}(J(f^{-1}(x)))|$, where J(f^-1(x)) is the Jacobian of the function inverse - a higher-dimensional generalization of the quantity dx/dy from before.


$$y = f(x)$$

$$p(y) = p(f^{-1}(y)) \cdot |\text{det} J(f^{-1}(y))|$$
$$\log p(y) = \log p(f^{-1}(y)) + \log |\text{det}(J(f^{-1}(y)))|$$

When I learned about determinants in middle & high school I was very confused at the seemingly arbitrary definition of determinants. We were only taught how to compute a determinant, instead of what a determinant meant: the local, linearized rate of volume change of a transformation.

Transformed Distributions in TensorFlow

TensorFlow has an elegant API for transforming distributions. A TransformedDistribution is specified by a base distribution object that we will transform, and a Bijector object that implements 1) a forward transformation $y = f(x)$, where $f : \mathbb{R}^d → \mathbb{R}^d$  2) its inverse transformation $x = f^-1(y)$, and 3) the inverse log determinant of the Jacobian $\log |\text{det}J (f^-1(y))|$. For the rest of this post, I will abbreviate this quantity as ILDJ.

Under this abstraction, forward sampling is trivial:

bijector.forward(base_dist.sample())

To evaluate log-density of the transformed distribution:

distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x)

Furthermore, if bijector.forward is a differentiable function, then Y = bijector.forward(x) is a re-parameterizable distribution with respect to samples x = base_distribution.sample(). This means that normalizing flows can be used as a drop-in replacement for variational posteriors in a VAE (as an alternative to a Gaussian).

Some commonly used TensorFlow distributions are actually implemented using these TransformedDistributions.

Source Distribution	Bijector.forward	Transformed Distribution
Normal	exp(x)	LogNormal
Exp(rate=1)	-log(x)	Gumbel(0,1)
Gumbel(0,1)	Softmax(x)	Gumbel-Softmax / Concrete

Under standard convention, TransformedDistributions are named as $\text{Bijector}^{-1}\text{BaseDistribution}$ so an ExpBijector applied to a Normal distribution becomes LogNormal. There are some exceptions to this naming scheme - the Gumbel-Softmax distribution is implemented as the RelaxedOneHotCategorical distribution, which applies a SoftmaxCentered bijector to a Gumbel distribution.

Normalizing Flows and Learning Flexible Bijectors

Why stop at 1 bijector? We can chain any number of bijectors together, much like we chain layers together in a neural network [3]. This is construct is known as a “normalizing flow”. Additionally, if a bijector has tunable parameters with respect to bijector.log_prob, then the bijector can actually be learned to transform our base distribution to suit arbitrary densities. Each bijector functions as a learnable “layer”, and you can use an optimizer to learn the parameters of the transformation to suit our data distribution we are trying to model. One algorithm to do this is maximum likelihood estimation, which modifies our model parameters so that our training data points have maximum log-probability under our transformed distribution. We compute and optimize over log probabilities rather than probabilities for numerical stability reasons.

This slide from Shakir Mohamed and Danilo Rezende’s UAW talk (slides) that illustrates this concept:

However, computing the determinant of an arbitrary $N \times N$ Jacobian matrix has runtime complexity $O(N^3)$, which is very expensive to put in a neural network. There is also the trouble of inverting an arbitrary function approximator. Much of the current research on Normalizing Flows focuses on how to design expressive Bijectors that exploit GPU parallelism during forward and inverse computations, all while maintaining computationally efficient ILDJs.

Code Example

Let’s build a basic normalizing flow in TensorFlow in about 100 lines of code. This code example will make use of:

• TF Distributions - general API for manipulating distributions in TF. For this tutorial you’ll need TensorFlow r1.5 or later.
• TF Bijector - general API for creating operators on distributions
• Numpy, Matplotlib.

We are trying to model the distribution $p(x_1, x_2) = \mathcal{N}(x1|\mu=1/4x_2^2, \sigma=1) \cdot N(x_2|\mu=0, \sigma=4)$. We can generate samples from the target distribution using the following code snippet (we generate them in TensorFlow to avoid having to copy samples from the CPU to the GPU on each minibatch):

For our base distribution, we’ll use an Isotropic Gaussian.



Next, we construct the bijector and create a TransformedDistribution from it. Let’s build a flow that resembles a standard fully-connected network, i.e. alternating matrix multiplication with nonlinearities.

The Jacobian of an affine function is trivial to compute, but worst case determinants are $O(n^3)$, which is unacceptably slow to compute. Instead, TensorFlow provides a structured affine transformation whose determinant can be computed more efficiently. This Affine transform is parameterized as a lower triangular matrix $M$ plus a low rank update:

$$M + V \cdot D \cdot V^T$$

To compute $\text{det}(M + V \cdot D \cdot V^T)$ cheaply, we use the matrix determinant lemma.

Next, we need an invertible nonlinearity in order to express non-linear functions (otherwise the chain of affine bijectors remains affine). Sigmoid / tanh may seem like good choices, but they are incredibly unstable to invert - small changes in the output near -1 or 1 correspond to massive changes in input. In my experiments I could not chain 2 saturating nonlinearities together without gradients exploding. Meanwhile, ReLU is stable, but not invertible for $x < 0$.

I chose to implement PReLU (parameterized ReLU), which is the same as Leaky ReLU but with a learnable slope in the negative regime. The simplicity of PReLU and its straightforward Jacobian makes for a nice exercise in implementing your own custom Bijectors: notice that the ILDJ is 0 when $x > 0$ (no volume change) and $1/\alpha$ otherwise (compensating for the contraction in volume from multiplying x by $\alpha$).

PReLU is an element-wise transformation, so the Jacobian is diagonal. The determinant of a diagonal matrix is just the product of the diagonal entries, so we compute the ILDJ by simply summing the diagonal entries of the log-Jacobian [4]. We build the “MLP Bijector” by using tfb.Chain(), then apply it to our base distribution to create the transformed distribution:


Finally, we’ll train the model using Maximum Likelihood estimation: maximize the expected log probability of samples from the real data distribution, under our choice of model.

We can visualize the (slow) deformation of space by coloring samples from base distribution according to their starting quadrant,

And that’s it! TensorFlow distributions makes normalizing flows to implement, and automatically accumulate all the Jacobians determinants in a way that is clean and highly readable. Full source code for this post can be found on Github.

You might notice that the deformation is rather slow, and it takes a lot of layers to learn a rather simple transformation [5]. In the next post, I will cover more modern techniques for learning normalizing flows.

Acknowledgements

I am very grateful to Dustin Tran for clarifying my understanding of normalizing flows, Luke Metz, Katherine Lee, and Samy Bengio for proofreading this post, and to Ben Poole, Rif A. Saurous, Ian Langmore for helping me to debug my code. You rock!

Footnotes

[1] The notion that we can augment our dataset with *new* information from a finite set of data is a rather disturbing one, and it remains to be shown whether probabilistic machine learning can truly replace true generative processes (e.g. simulation of fluid dynamics), or whether at the end of the day it is only good for amortizing computation and any generalization we get on the training / test distribution is a lucky accident.
[2] See A note on the evaluation of generative models for a thought-provoking discussion about how high log-likelihood is neither sufficient nor necessary to generate “plausible” images. Still, it’s better than nothing and in practice a useful diagnostic tool.
[3]There’s a connection between Normalizing Flows and GANs via encoder-decoder GAN architectures that learn the inverse of the generator (ALI / BiGAN). Since there is a separate encoder trying to recover $u = G^{-1}(X)$ such that $X = G(u)$, the generator can be thought of as a flow for the simple uniform distribution. However, we don’t know how to compute the amount of volume expansion/contraction w.r.t. X, so we cannot recover density from GANs. However, it’s probably not entirely unreasonable to model the log-det-jacobian numerically or enforce some kind of linear-time Jacobian by construction.
[4] The lemma “Determinant of diagonal matrices is the product of the diagonal entries” is quite intuitive from a geometric point of view: each dimension’s length distortion is independent of the other dimensions, so the total volume change is just the product of changes in each direction, as if we were computing the volume of a high-dimensional rectangular prism.
[5] This MLP is rather limited in capacity because each affine transformation is only a 2x2 matrix, and the PReLU “warps” the underlying distribution very slowly (so several PreLUs are needed to bend the data into the right shape). For low dimensional distributions, this MLP is a very poor choice of a normalizing flow, and is meant for educational purposes.