Wasserstein Loss/es
| Article | |
|---|---|
| Topic area | Generative Models |
| Prerequisites | Generative Adversarial Network, Optimal Transport, Lipschitz Continuity, Kullback-Leibler Divergence |
Visión general
La pérdida de Wasserstein es una función objetivo para el entrenamiento de modelos generativos que mide la distancia entre dos distribuciones de probabilidad utilizando la distancia del transportista de tierra (EMD), también conocida como la 1-distancia de Wasserstein. Introducida para redes generativas adversarias (GAN) por Arjovsky et al. en 2017, sustituyó el criterio de Jensen-Shannon implícito en la GAN original por una métrica que se mantiene bien definida y continua incluso cuando las distribuciones del modelo y de los datos tienen soportes disjuntos. El marco resultante de Wasserstein GAN (WGAN) mejoró sustancialmente la estabilidad del entrenamiento y proporcionó una pérdida cuya magnitud se correlaciona con la calidad de las muestras.
La pérdida de Wasserstein es ahora una herramienta estándar siempre que se necesita una comparación distribucional y los soportes pueden no superponerse. Más allá de las GAN, aparece en adaptación de dominio, optimización distribucionalmente robusta y estimación de densidad. Su atractivo central es geométrico: en lugar de preguntar cuánto se superpone una distribución con otra, pregunta cuánta masa de probabilidad debe moverse, y a qué distancia, para transformar una en la otra.
Intuición
Considere dos pilas de tierra dispuestas sobre una recta, que representan dos distribuciones de masa total unitaria. La distancia del transportista de tierra es el costo mínimo de transformar la primera pila en la segunda transportando tierra, donde mover una unidad de masa a lo largo de una distancia $ d $ incurre en un costo $ d $. Si las pilas son idénticas, el costo es cero. Si una pila está centrada en $ 0 $ y la otra en $ \theta $, el costo es $ |\theta| $, independientemente de cuán concentrada esté cada pila.
Esta última propiedad expone la debilidad central de la divergencia KL y la divergencia de Jensen-Shannon en el mismo escenario. Si ambas distribuciones son masas puntuales (deltas de Dirac) en $ 0 $ y $ \theta $, la KL es infinita para $ \theta \neq 0 $, y Jensen-Shannon es la constante $ \log 2 $. Ninguna tiene un gradiente útil con respecto a $ \theta $. La distancia de Wasserstein, en cambio, es igual a $ |\theta| $ con un gradiente informativo $ \mathrm{sign}(\theta) $. Esta es exactamente la situación que enfrenta un generador al inicio del entrenamiento, cuando su distribución de salida está soportada en una variedad de baja dimensión que casi con certeza no se superpone con la variedad de los datos. Una pérdida que proporciona gradiente en este régimen puede impulsar el aprendizaje donde las pérdidas basadas en divergencias no pueden.
Formulación
Para dos medidas de probabilidad $ P_r $ (real) y $ P_g $ (generada) en un espacio métrico $ (\mathcal{X}, d) $, la 1-distancia de Wasserstein es
$ {\displaystyle W_1(P_r, P_g) = \inf_{\gamma \in \Pi(P_r, P_g)} \mathbb{E}_{(x, y) \sim \gamma} [\, d(x, y) \,]} $
donde $ \Pi(P_r, P_g) $ es el conjunto de distribuciones conjuntas con marginales $ P_r $ y $ P_g $. Cada $ \gamma $ es un plan de transporte: $ \gamma(x, y) $ especifica cuánta masa en $ x $ en la fuente se envía a $ y $ en el objetivo.
Esta forma primal es intratable para distribuciones continuas de alta dimensión. La dualidad de Kantorovich-Rubinstein proporciona un equivalente manejable:
$ {\displaystyle W_1(P_r, P_g) = \sup_{\|f\|_L \leq 1} \, \mathbb{E}_{x \sim P_r}[f(x)] - \mathbb{E}_{x \sim P_g}[f(x)]} $
donde el supremo se toma sobre todas las funciones 1-Lipschitz $ f : \mathcal{X} \to \mathbb{R} $. La función $ f $ se denomina crítico (en lugar de discriminador) porque produce una puntuación con valor real no acotada en lugar de una probabilidad.
En una GAN de Wasserstein, el crítico $ f_w $ se parametriza mediante una red neuronal y se entrena para maximizar el objetivo dual, mientras que el generador $ g_\theta $ se entrena para minimizar $ -\mathbb{E}[f_w(g_\theta(z))] $. El problema min-max es
$ {\displaystyle \min_\theta \max_{w : \|f_w\|_L \leq 1} \mathbb{E}_{x \sim P_r}[f_w(x)] - \mathbb{E}_{z \sim P_z}[f_w(g_\theta(z))].} $
La cuestión práctica crítica es cómo imponer la restricción 1-Lipschitz sobre $ f_w $.
Imposición de la restricción de Lipschitz
Recorte de pesos
La WGAN original recorta cada peso del crítico a un intervalo pequeño $ [-c, c] $ después de cada actualización de gradiente, con $ c $ típicamente $ 0.01 $. Esto garantiza una constante de Lipschitz acotada, pero restringe la capacidad expresiva del crítico y hace que la restricción efectiva sea sensible a $ c $. Las redes recortadas frecuentemente se saturan o no logran aprender características de alta frecuencia.
Penalización de gradiente
WGAN-GP (Gulrajani et al., 2017) reemplaza el recorte por una penalización suave que impone que la norma del gradiente del crítico sea aproximadamente $ 1 $:
$ {\displaystyle \mathcal{L}_{\mathrm{GP}} = \mathbb{E}_{\hat{x} \sim P_{\hat{x}}} \left[ (\|\nabla_{\hat{x}} f_w(\hat{x})\|_2 - 1)^2 \right]} $
donde $ \hat{x} $ se muestrea uniformemente a lo largo de líneas rectas entre pares de $ P_r $ y $ P_g $. La pérdida total del crítico se convierte en el objetivo dual menos $ \lambda \mathcal{L}_{\mathrm{GP}} $, con $ \lambda $ típicamente $ 10 $. WGAN-GP es incompatible con la normalización por lotes en el crítico, ya que la penalización es por muestra, pero funciona bien con normalización de capa o instancia.
Normalización espectral
La normalización espectral (Miyato et al., 2018) divide cada matriz de pesos por su mayor valor singular, aproximado mediante iteración de potencias. Dado que la constante de Lipschitz de una composición está acotada por el producto de las constantes de Lipschitz de las capas, normalizar cada capa para que tenga norma espectral $ 1 $ limita la red a 1-Lipschitz. Esto es computacionalmente económico y se desacopla del optimizador.
Otros enfoques
También se han propuesto penalizaciones de consistencia, pérdidas de tipo hinge y proyección directa sobre clases de funciones de Lipschitz. La elección interactúa con el optimizador (típicamente Adam o RMSProp), la proporción de actualizaciones crítico-a-generador (a menudo 5:1) y la arquitectura.
Entrenamiento e inferencia
Un paso de entrenamiento típico de una GAN de Wasserstein alterna:
- Muestrear un minilote de los datos reales y un lote de vectores latentes $ z \sim P_z $.
- Actualizar el crítico $ n_{\mathrm{critic}} $ veces ascendiendo el objetivo dual (y restando la penalización de gradiente si corresponde).
- Actualizar el generador una vez descendiendo $ -\mathbb{E}[f_w(g_\theta(z))] $.
El crítico se entrena hasta casi la convergencia en cada paso del generador en el artículo original de WGAN, en contraste con la GAN original donde se contiene al discriminador. Con la pérdida de Wasserstein, un crítico más fuerte proporciona una mejor señal de gradiente, porque el objetivo dual ofrece una estimación de $ W_1 $ que el generador puede seguir de manera útil.
En el momento de la inferencia, la pérdida no juega ningún papel. La calidad de las muestras se evalúa mediante las métricas generativas habituales (puntuación de Inception, distancia de Frechet Inception), con el beneficio adicional de que el valor convergente del objetivo del crítico rastrea $ W_1 $ y puede servir como un indicador, en tiempo de entrenamiento, de la calidad de las muestras.
Variantes
- Distancia de Wasserstein por cortes calcula $ W_1 $ sobre proyecciones unidimensionales y promedia sobre las direcciones, aprovechando la solución 1D en forma cerrada para evitar la estimación de un crítico neuronal.
- Divergencia de Sinkhorn aproxima Wasserstein con un regularizador entrópico, permitiendo una implementación diferenciable y paralelizable en GPU mediante el algoritmo de Sinkhorn.
- Distancia de energía y MMD son alternativas basadas en núcleos que comparten la propiedad de permanecer definidas en presencia de soportes disjuntos, con distintos compromisos sesgo-varianza.
- Críticos relativistas modifican el objetivo dual para puntuar diferencias relativas entre muestras reales y falsas, mejorando a menudo la estabilidad.
- $ p $-Wasserstein para $ p > 1 $ eleva el costo de transporte a la potencia $ p $; solo $ p = 1 $ admite el dual de Kantorovich-Rubinstein utilizado en WGAN, pero $ p = 2 $ aparece en la estimación de densidad basada en transporte óptimo y en modelos de difusión.
Comparación con otras pérdidas
| Pérdida | Soporte disjunto | Continua en $ \theta $ | Acotada |
|---|---|---|---|
| Divergencia KL | Infinita | No | No |
| Jensen-Shannon | Constante $ \log 2 $ | No | Sí |
| Wasserstein-1 | Finita, geométrica | Sí | No |
| MMD (con núcleo característico) | Finita | Sí | Sí |
La GAN original minimiza una cantidad proporcional a Jensen-Shannon cuando el discriminador es óptimo. Esto explica la patología empírica del colapso de modos y los gradientes que se desvanecen: cuando el generador está lejos de los datos, el discriminador se vuelve casi perfecto y proporciona casi ninguna señal útil. La pérdida de Wasserstein evita estructuralmente este régimen.
En comparación con las pérdidas de mínimos cuadrados o hinge, Wasserstein típicamente produce curvas de entrenamiento más interpretables, pero exige la restricción de Lipschitz y una mayor proporción de actualizaciones crítico-a-generador. En la práctica, muchas GAN de producción combinan normalización espectral con pérdidas hinge o no saturantes, difuminando la distinción categórica.
Limitaciones
La restricción de Lipschitz, sea cual sea la forma en que se imponga, es aproximada. El recorte de pesos limita la capacidad; la penalización de gradiente se impone solo a lo largo de líneas muestreadas; la normalización espectral acota las normas por capa, pero la verdadera constante de Lipschitz de la red puede ser menor. El objetivo dual estima por tanto $ W_1 $ hasta un factor multiplicativo que varía durante el entrenamiento, por lo que el valor de la pérdida es comparable entre ejecuciones solo de manera aproximada.
La pérdida de Wasserstein no elimina por completo el colapso de modos; aunque es menos propensa a él que la GAN original, todavía puede producirse pérdida parcial de modos, especialmente con críticos débiles. El costo computacional es mayor por actualización del generador debido a la mayor proporción crítico-a-generador. La penalización de gradiente también requiere un gradiente de segundo orden (gradiente de una norma del gradiente), lo que aumenta la memoria y el cómputo en un factor constante.
La 1-distancia de Wasserstein ignora la estructura geométrica de orden superior que captura $ p = 2 $, y en imágenes naturales de muy alta dimensión la métrica asumida en el espacio de píxeles puede no coincidir con la similitud perceptual, lo cual es una de las razones por las que a veces se utilizan adicionalmente pérdidas perceptuales o distancias en el espacio de características al estilo de FID.
Referencias
- ↑ Template:Cite arxiv
- ↑ Template:Cite arxiv
- ↑ Template:Cite arxiv
- ↑ Villani, Cedric. Optimal Transport: Old and New. Springer, 2009.
- ↑ Template:Cite arxiv