Pyramid-JiT elimina el VAE y entrena modelos de imagen 4,3 veces más rápido
Training Text-to-Image Models Without a VAE
Linum presenta Pyramid-JiT, una arquitectura de espacio de píxeles solo con decodificador que predice la imagen objetivo en varias resoluciones a lo largo del tronco DiT. Al eliminar el VAE y usar múltiples cabezas de salida, alcanza el FD-DINOv2 de Linum v2 con 11,3 veces menos muestras y entrena en 4,3 veces menos horas de GPU a 4 veces más píxeles. El código y los pesos están disponibles bajo licencia Apache 2.0.
Si podemos reducir el número de tokens, podemos recortar el costo de la atención y lograr ahorros masivos en entrenamiento e inferencia.
- schopra909
¡Hola HN, aquí el autor!
Para dar contexto, somos un laboratorio de 2 personas entrenando modelos generativos de video. El objetivo es un nuevo conjunto de herramientas de animación controlables (pueden leer más sobre eso aquí https://www.linum.ai/about si tienen curiosidad).
El mayor cuello de botella de nuestro último modelo de texto-a-video en términos de costo de entrenamiento e inferencia es la atención. Los modelos de video son increíblemente densos en tokens (por ejemplo, 110K tokens para un clip de varios segundos). Si podemos condensar esa ventana de contexto de forma más agresiva, podemos entrenar modelos más grandes por mucho menos $$ y ofrecérselos a los prosumers a precios razonables (a diferencia de los grandes modelos actuales como Seedance, que cuestan un ojo de la cara de ejecutar).
Tradicionalmente, los modelos de imagen y video tienen dos componentes disjuntos: VAE (Variational Autoencoder) y Diffusion Transformer (DiT). Se entrenan por separado, y empíricamente los VAE parecen tener dificultades para superar la reducción de tokens de 16x16.
Aquí, estamos cambiando al espacio de píxeles, desechando el VAE, y logrando una reducción de tokens de 32x32 (ventanas de contexto 4x más pequeñas) mientras aprendemos un modelo general mejor con una fracción de las muestras de entrenamiento.
La tesis central es "lo más simple es mejor". Si podemos poner el problema de compresión en el Diffusion Transformer (DiT) más potente, deberíamos poder aprender un "espacio latente" optimizado para la generación y obtener una mejor compresión sin perjudicar la calidad de la generación.
Estaré revisando este post de vez en cuando en las próximas horas, así que siéntanse libres de dejar preguntas abajo. Y trataré de responderlas lo mejor […]
- bitpush
¿Hay alguna forma de usar esto en ComfyUI hoy?
Requisitos: