Los generadores de imágenes con IA ya entrenados aprenden a rellenar cualquier región con una fracción del coste

Los modelos autorregresivos generan una imagen igual que los modelos de lenguaje escriben texto: token a token. Primero la imagen se divide en códigos discretos y después un transformer los predice uno por uno, normalmente por filas, de izquierda a derecha y de arriba abajo. Este orden se conoce como raster scan.

Investigadores de Stony Brook University y ÉTS Montréal han propuesto un método para convertir un modelo raster ya entrenado en un generador de orden arbitrario. El modelo puede rellenar huecos dentro del encuadre (inpainting), completar la imagen más allá de sus bordes (outpainting) y editar regiones seleccionadas. Para ello no hace falta reentrenar desde cero.

Según los autores, la adaptación adicional cuesta más de un orden de magnitud menos que entrenar desde cero un modelo de «cualquier orden».

Por qué un modelo por filas no sabe arreglar el centro

Con el entrenamiento raster, el modelo solo ve lo que está dibujado más arriba y más a la izquierda. No aprendió a predecir un token cuyos vecinos de la derecha y de abajo ya se conocen. Por eso, una máscara en mitad del encuadre lo deja bloqueado.

Existen dos formas de sortear esta limitación, y ambas son caras:

• Entrenar desde cero con una mezcla de órdenes. Por ejemplo, RandAR requiere 300 épocas, mientras que LlamaGen-XL, con orden por filas, necesita 100. El presupuesto se triplica.
• Difusión discreta, categoría a la que pertenece también MaskGIT. El modelo restaura de forma iterativa los tokens ocultos por la máscara con atención bidireccional. La calidad es buena, pero la KV-cache (almacenamiento de datos intermedios ya calculados para no recalcularlos en cada paso), la decodificación especulativa y servidores como vLLM no funcionan con este esquema.

Las empresas que ya tienen un pipeline autorregresivo y sus servidores se ven obligadas a construir un segundo stack en paralelo. Los autores consideran superfluo ese gasto.

Dos piezas sin las que el reentrenamiento se rompe

La idea se apoya en una propiedad matemática: por la regla de la cadena de probabilidades, la distribución conjunta de los tokens se puede descomponer en cualquier orden. Por tanto, un modelo que ha asimilado bien las imágenes en un orden ya sabe casi todo lo necesario. Lo que le falta es saber usar ese conocimiento con otro orden. Para lograrlo hicieron falta dos componentes, y por separado no funcionan.

Embeddings posicionales «con visión de futuro» (future-aware). En el orden raster, el modelo siempre sabe qué va a dibujar a continuación. En un orden aleatorio, un mismo contexto puede conducir a objetivos distintos.

El token del cielo y el token de un rostro se predicen de forma diferente a partir de los mismos datos visibles, y los gradientes entran en conflicto. En el modelo de lenguaje XLNet, este problema se resolvía con un segundo flujo de atención, que duplica el cómputo de la capa y rompe la KV-cache. Aquí la solución es más sencilla.

A la entrada de cada paso se suma un vector aprendido de la posición que hay que predecir a continuación. La tabla parte de ceros y la suma se aplica solo en la capa de entrada. La arquitectura sigue siendo un transformer corriente de un solo flujo.

Aparte, se ha decidido qué posición usar para la atención. En LlamaGen, que emplea embeddings posicionales rotatorios RoPE, y en RAR, que usa tablas de posiciones absolutas, se toma el índice espacial del token y no el número de paso. El modelo conserva todo lo que sabía sobre la vecindad de los píxeles.

Complejidad gradual de los órdenes (curriculum). Si se aplican permutaciones aleatorias desde el principio, las pérdidas son enormes, los gradientes son ruidosos y el modelo se sale del buen punto de preentrenamiento y olvida lo acumulado. Los autores introducen un coeficiente que crece de forma lineal: al principio casi todos los ejemplos van en orden raster y después, cada vez más, en orden aleatorio.

Cada ejemplo del batch recibe su propia permutación. Reducir simplemente el learning rate no ayuda aquí, porque el gradiente seguiría apuntando en la misma dirección incómoda. El calendario cambia la propia composición de los datos de entrenamiento, y los primeros gradientes concuerdan con la solución preentrenada.

Si se elimina cualquiera de los dos componentes, el entrenamiento diverge o la calidad cae de forma notable.

Qué se logró y dónde está el precio

El método se aplicó a cuatro modelos base de dos arquitecturas y dos tamaños: LlamaGen-L, LlamaGen-XL, RAR-L y RAR-XL. Resultados:

• En inpainting y outpainting sobre ImageNet, los modelos adaptados superan con claridad a MaskGIT en FID. Esta métrica compara la distribución de las imágenes generadas con la de las reales, y cuanto más baja, mejor.
• Con MAR-L, que trabaja con tokens continuos y un cabezal de difusión, los resultados son comparables. El modelo no se entrenó específicamente para ninguna de estas tareas.
• Los mismos pesos sirven para inpainting, outpainting y edición por regiones. Lo único que cambia es el conjunto de tokens conocidos en la entrada.
• Hay un precio. En LlamaGen-XL, el FID de la generación estándar por clases aumentó. El Inception Score, la precision y el recall se mantuvieron cerca de los valores originales. Los autores lo dicen sin rodeos: es un intercambio de calidad por flexibilidad, no una conversión gratuita.

La arquitectura en sí sigue siendo causal, por lo que la KV-cache y la decodificación paralela funcionan. El modelo se puede ejecutar con vLLM mediante el registro de un modelo propio. Para la difusión discreta, según los autores, esto es inviable por principio.

Qué aporta a los ingenieros

Quien ya dispone de un modelo autorregresivo de imágenes preentrenado y de la infraestructura para ejecutarlo obtiene edición sin un segundo stack y sin reentrenar. El inpainting y el outpainting pasan a ser un modo de un modelo ya existente. Las limitaciones también están claras. La validación se limita a ImageNet y a cuatro modelos base, y la generación estándar pierde algo de calidad. Pero la dirección es nítida: se pueden extraer capacidades útiles de modelos que ya están entrenados. El código está publicado en GitHub, en el repositorio Ordered-AR/unlockedadaptation.