La generación de imágenes mediante modelos de difusión ha estado tradicionalmente confinada a la nube debido a los elevados requisitos computacionales y de memoria de los modelos de gran escala. Esto impone limitaciones inherentes en términos de latencia, costo operativo por inferencia y privacidad de los datos, especialmente en flujos de trabajo iterativos donde los usuarios refinan continuamente sus prompts.
El problema fundamental que aborda Bonsai Image 4B es cómo trasladar la capacidad de generación de imágenes de alta calidad a dispositivos edge, superando las restricciones de memoria y ancho de banda. Esto se logra mediante la cuantificación agresiva de los pesos del transformador de difusión, un componente crítico y el más grande del modelo, sin una pérdida significativa en la fidelidad de la imagen o la adherencia al prompt. La viabilidad de esta estrategia se basa en la observación de que, para muchas tareas de inferencia, la precisión de punto flotante completa (FP32 o FP16) puede ser excesiva, y representaciones de menor precisión pueden ser suficientes si se manejan cuidadosamente.
Históricamente, la compresión de modelos ha sido un área activa de investigación, con técnicas como la poda (pruning), la destilación (distillation) y la cuantificación (quantization) buscando reducir el tamaño y la complejidad computacional de las redes neuronales. Bonsai Image 4B se enfoca en la cuantificación de pesos a niveles extremos (1-bit y ternario) para habilitar un nuevo régimen de despliegue para modelos de difusión, democratizando el acceso a esta tecnología en hardware de consumo.
Arquitectura del Sistema
Bonsai Image 4B se construye a partir del modelo FLUX.2 Klein 4B, manteniendo su arquitectura subyacente pero modificando la representación de los pesos del transformador de difusión. El componente clave es la cuantificación de los pesos del transformador, que constituye la mayor parte del modelo y se invoca repetidamente durante el proceso de denoising. Se ofrecen dos variantes:
1. 1-bit Bonsai Image 4B: Utiliza pesos de transformador binarios ({-1, +1}) con un factor de escala group-wise en FP16. Esto resulta en una reducción de aproximadamente 14x en el tamaño de los pesos del transformador en comparación con la precisión completa. Un pequeño subconjunto de tensores de soporte sensibles a la precisión (~5%), denominados capas de proyección, se mantiene en FP16 para preservar la calidad. El transformador final de 1-bit ocupa 0.93 GB, una reducción de 8.3x respecto a los 7.75 GB del FLUX.2 Klein 4B original.
2. Ternary Bonsai Image 4B: Emplea pesos de transformador ternarios ({-1, 0, +1}) con un factor de escala group-wise en FP16. La inclusión del estado cero proporciona mayor flexibilidad representacional, mejorando la calidad visual y la fidelidad del prompt. Esta variante logra una reducción de aproximadamente 10x en los pesos del transformador, resultando en un transformador final de 1.21 GB, una reducción de 6.4x. Aunque ligeramente más grande que la variante de 1-bit, ofrece una mejor calidad.
Ambas variantes incluyen un codificador de texto comprimido y un VAE (Variational Autoencoder) en FP16. El payload total de despliegue para Apple Silicon es de 3.42 GB para la versión de 1-bit y 3.88 GB para la ternaria, frente a los 15.97 GB del FLUX.2 Klein 4B de precisión completa. En tiempo de ejecución, el codificador de texto se descarga después de la codificación del prompt, reduciendo aún más el uso de memoria activa. La implementación de bajo-bit se apoya en MLX para hardware Apple y en kernels GEMM de bajo-bit Gemlite para GPUs CUDA, optimizando las operaciones matriciales con pesos cuantificados.
Flujo de Inferencia de Imagen Local (Bonsai Image 4B)
- 1 Prompt Encoding El prompt de texto se procesa y codifica usando el text encoder comprimido.
- 2 Text Encoder Offload El text encoder se descarga de la memoria activa para liberar recursos.
- 3 Denoising Steps (Iterative) El transformador de difusión (cuantificado a 1-bit o ternario) se invoca repe...
- 4 Low-Bit GEMM Kernels Operaciones matriciales optimizadas usando MLX (Apple) o Gemlite (CUDA).
- 5 Image Generation Cada paso refina la imagen latente, reduciendo el ruido.
- 6 VAE Decoding El VAE en FP16 convierte la representación latente final en una imagen visible.
- 7 Display Output La imagen generada se muestra al usuario en el dispositivo.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | MLX | Framework de aprendizaje automático optimizado para hardware Apple Silicon, proporcionando rutas de bajo-bit para la inferencia cuantificada. |
| compute | Gemlite | Kernels GEMM (General Matrix Multiply) de bajo-bit optimizados para GPUs CUDA, permitiendo operaciones eficientes con pesos cuantificados. |
| storage | 1-bit Transformer Weights | Representación binaria ({-1, +1}) de los pesos del transformador para máxima compresión y reducción de footprint de memoria. vs FP32 weights, FP16 weights FP16 group-wise scaling factor |
| storage | Ternary Transformer Weights | Representación ternaria ({-1, 0, +1}) de los pesos del transformador para un equilibrio entre compresión y fidelidad. vs FP32 weights, FP16 weights, 1-bit weights FP16 group-wise scaling factor |
| data-processing | Diffusion Transformer | Componente principal del modelo de difusión encargado de los pasos de denoising, cuantificado para inferencia eficiente. |
| data-processing | VAE (Variational Autoencoder) | Utilizado para decodificar la representación latente final en una imagen de píxeles, manteniendo precisión FP16. FP16 precision |
Trade-offs
Ganancias
- ▲▲ Memory Footprint Reduction
- ▲ Inference Latency
- ▲▲ On-device Deployment Capability
- ▲ Privacy
- ▲ Operational Cost (Cloud)
Costes
- △ Model Accuracy/Fidelity
Fundamentos Teóricos
La cuantificación de redes neuronales es un campo de investigación activo que se remonta a trabajos tempranos sobre la reducción de la precisión de los pesos para acelerar la inferencia y reducir el footprint de memoria. Un hito importante fue el artículo "BinaryConnect: Training Deep Neural Networks with Binary Weights during Propagations" de Courbariaux et al. (2015), que exploró la binarización de pesos y activaciones. Este trabajo demostró que era posible entrenar redes con pesos binarios, sentando las bases para la cuantificación extrema.
El concepto de cuantificación ternaria, que añade un estado cero, también ha sido explorado para mejorar el trade-off entre compresión y precisión, como se ve en "Ternary Weight Networks" de Li et al. (2016). Estos enfoques buscan minimizar la pérdida de información al representar números de punto flotante con un número muy limitado de bits, a menudo utilizando técnicas como la cuantificación post-entrenamiento o la cuantificación consciente del entrenamiento.
El uso de factores de escala group-wise en FP16 es una técnica para compensar la pérdida de rango dinámico inherente a la cuantificación de bajo-bit, permitiendo que cada grupo de pesos tenga su propio factor de escala para mapear los valores cuantificados al rango adecuado. Esto se relaciona con principios de compresión de datos y codificación de información, donde la representación eficiente de datos es crucial para el rendimiento en sistemas con recursos limitados.