El problema fundamental que Needle 2 aborda es la ejecución eficiente de modelos de lenguaje en el 'verdadero edge': dispositivos con restricciones severas de hardware, como microcontroladores, Raspberry Pis y teléfonos de bajo costo, que carecen de GPUs o NPUs dedicadas y tienen memoria RAM limitada. Tradicionalmente, la IA en el edge se ha asociado con hardware más potente (laptops, PCs), dejando un vacío para los miles de millones de dispositivos IoT y móviles de bajo presupuesto.
Needle 2 postula que para tareas específicas como la llamada a funciones (ej. 'encender la luz'), la extracción estructurada o el control de dispositivos, no se requiere el conocimiento del mundo o la capacidad de generación de texto abierto de un modelo de lenguaje grande. Al enfocar el problema a una formulación más acotada, es posible reducir drásticamente el tamaño del modelo y los requisitos computacionales, haciendo la inferencia local, privada e instantánea una realidad en hardware de <$200. Esto representa un cambio de paradigma desde la dependencia de la nube o hardware de edge de gama alta hacia una IA verdaderamente ubicua y descentralizada.
Arquitectura del Sistema
Needle 2 es un modelo de 45M de parámetros preentrenado y post-entrenado con un corpus específico para acciones de dispositivos y extracción estructurada. Su arquitectura se desvía del transformer convencional para optimizar el consumo de recursos. Incorpora una Hadamard MLP que reemplaza las proyecciones densas con una transformada de Walsh fija y diagonales aprendidas, reduciendo significativamente los parámetros asociados a la mezcla de canales. Utiliza un 'engram' para mover el conocimiento del mundo fuera del stack de atención a tablas de n-gramas hasheadas, que se leen de forma eficiente en tiempo de decodificación, minimizando las lecturas de flash y el consumo de batería.
El sistema de memoria está diseñado para dispositivos con RAM fija. La atención utiliza una ventana deslizante de 256 tokens para limitar el tamaño del KV cache, asegurando que el consumo de RAM sea determinista (28MB). El prompt del sistema y las declaraciones de herramientas se 'fijan' como 'sinks' permanentes, impidiendo su desalojo. La cuantización de 2 bits (CQ2-bit) se aplica desde el preentrenamiento hasta el post-entrenamiento, incluyendo pesos, activaciones y KV cache, garantizando que el modelo desplegado sea el modelo entrenado, sin pérdidas de calidad por cuantización post-hoc. El motor de inferencia, un binario C++ sin dependencias, está co-diseñado con el modelo. No descomprime los pesos en RAM; los códigos de 2 bits se expanden directamente en registros vectoriales y se fusionan en productos punto enteros, manteniendo la memoria residente al tamaño del blob (14MB) y la ruta aritmética en int8 de principio a fin. Una gramática compilada a nivel de byte, derivada de los esquemas declarados, restringe cada token de salida, permitiendo al motor calcular scores solo para tokens candidatos y saltarse hasta el 98% de la proyección del vocabulario en tokens estructurales, optimizando el uso de FLOPs y bytes por token.
Flujo de Inferencia de Needle 2 en Dispositivo
- 1 Entrada de Usuario Sentencia de lenguaje natural (ej. 'enciende la luz de la sala')
- 2 Tokenizador Convierte la entrada en tokens para el modelo
- 3 Motor de Inferencia (C++) Carga el modelo CQ2-bit y selecciona kernels optimizados para CPU
- 4 Modelo Needle 2 (45M params) Procesa tokens, utilizando Hadamard MLP, engram y atención de ventana deslizante
- 5 KV Cache (256 tokens) Almacena claves y valores de atención, con prompts y herramientas fijadas
- 6 Gramática Compilada Restringe la generación de tokens de salida a un esquema válido
- 7 Generación de Salida Calcula logits solo para tokens válidos, produciendo un 'call envelope'
- 8 Acción del Dispositivo Ejecuta la función o extracción estructurada (ej. 'call light.turn_on(room='s...
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Needle 2 Engine (C++) | Motor de inferencia co-diseñado para ejecutar el modelo de forma ultraligera y eficiente en CPUs de bajo consumo. Optimizado para int8 y expansión de 2-bit directamente en registros vectoriales. vs TensorFlow Lite, ONNX Runtime, PyTorch Mobile Binario universal que auto-selecciona kernels (SDOT, NEON, AVX2, RISC-V vectors, wasm SIMD, scalar) en tiempo de ejecución. |
| storage | Cactus Quants (CQ2-bit) | Esquema de cuantización de 2 bits integrado en el entrenamiento del modelo, permitiendo un tamaño de archivo de 14MB y manteniendo la calidad. Aplica a pesos, activaciones y KV cache. vs Post-training quantization (PTQ) a 8-bit o 4-bit, Quantization-aware training (QAT) a 8-bit Cuantización de 2 bits entrenada de extremo a extremo, no post-hoc. |
| compute | Hadamard MLP | Reemplaza las proyecciones densas tradicionales en el MLP, reduciendo el conteo de parámetros y el costo computacional para la mezcla de canales. vs MLP densos estándar en transformers |
| storage | Engram (hashed n-gram tables) | Estructura de datos para almacenar 'conocimiento del mundo' fuera del stack de atención, permitiendo una recuperación eficiente y de bajo costo en tiempo de decodificación. vs Almacenar todo el conocimiento en los pesos del modelo Lectura de pocas filas por token, optimizado para dispositivos con latencia de flash/DRAM. |
| cache | KV Cache con ventana deslizante de 256 tokens | Gestiona la memoria del estado de atención de forma determinista, limitando el consumo de RAM a 28MB independientemente de la longitud de la sesión. Fija prompts y declaraciones de herramientas. vs KV cache de tamaño ilimitado (crecimiento lineal con la longitud de la secuencia) Prompts y declaraciones de herramientas fijadas como 'sinks' permanentes. |
Trade-offs
Ganancias
- ▲▲ Footprint de memoria (RAM)
- ▲▲ Tamaño del modelo (almacenamiento)
- ▲▲ Eficiencia energética (MFLOPs/token)
- ▲ Latencia de inferencia en edge
- ▲ Privacidad y fiabilidad offline
Costes
- ▲ Capacidad de lenguaje general (chat, prosa, conocimiento del mundo)
- ▲ Flexibilidad para tareas fuera de tool calling/extracción estructurada
Fundamentos Teóricos
La optimización de modelos para inferencia en dispositivos con recursos limitados se conecta con décadas de investigación en compresión de modelos, cuantización y diseño de arquitecturas eficientes. La cuantización, en particular, tiene raíces en el procesamiento de señales y la teoría de la información, buscando representar datos con menos bits manteniendo la fidelidad. Trabajos como los de "Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference" (Jacob et al., 2018) de Google, o "Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding" (Han et al., 2015) de Stanford, sentaron las bases para la cuantización post-entrenamiento. Sin embargo, Needle 2 va más allá al integrar la cuantización de 2 bits (Cactus Quants) directamente en el proceso de entrenamiento, abordando el desafío de la pérdida de calidad en cuantizaciones agresivas, un problema bien conocido en la literatura.
El uso de una ventana deslizante para la atención y la fijación de prompts en el KV cache se relaciona con técnicas de gestión de memoria y atención eficiente en transformers, como las propuestas en "Longformer: The Long-Document Transformer" (Beltagy et al., 2020) o "Reformer: The Efficient Transformer" (Kitaev et al., 2020), que buscan reducir la complejidad cuadrática de la atención. La idea de mover el 'conocimiento del mundo' a estructuras de datos externas (engram) para reducir el tamaño del modelo y el costo computacional en inferencia, aunque no directamente un algoritmo clásico, resuena con principios de sistemas expertos y bases de conocimiento, donde la información se almacena y recupera de manera estructurada en lugar de ser 'aprendida' implícitamente en los pesos del modelo, un enfoque que se ha explorado en la IA simbólica y híbrida.