El problema fundamental que Needle 2 aborda es la ejecución eficiente de modelos de lenguaje en dispositivos de borde con recursos extremadamente limitados, como microcontroladores, Raspberry Pis y teléfonos de bajo costo. Tradicionalmente, la inferencia de modelos de lenguaje grandes (LLMs) requiere hardware potente con GPUs o NPUs y gigabytes de RAM. Sin embargo, la vasta mayoría de los dispositivos conectados en el mundo son de bajo costo y no poseen estas capacidades. Needle 2 postula que para tareas específicas como la llamada a funciones (tool calling) y la extracción estructurada, no se necesita el 'conocimiento del mundo' ni la capacidad de generar prosa abierta de un LLM generalista. En su lugar, un modelo mucho más pequeño, altamente optimizado para la eficiencia energética y de memoria, puede ser igualmente efectivo.
La relevancia de esta tesis radica en la democratización de la IA en el borde, permitiendo experiencias de usuario privadas, instantáneas y offline en miles de millones de dispositivos IoT y móviles. Esto contrasta con la tendencia actual de la 'Edge AI' que a menudo se limita a PCs y Macs de gama alta. Al redefinir el problema a uno de mapeo de lenguaje natural a funciones con parámetros tipados, Needle 2 reduce drásticamente los requisitos computacionales, haciendo viable la IA en dispositivos con presupuestos de potencia y memoria de unos pocos cientos de milivatios y decenas de megabytes.
Arquitectura del Sistema
Needle 2 se basa en una arquitectura de transformador altamente modificada y un motor de inferencia co-diseñado. El modelo de 45M de parámetros es pre-entrenado en un corpus de 115B tokens y post-entrenado en 38B tokens, con un enfoque en trazas de razonamiento compactas. Las decisiones arquitectónicas clave incluyen:
1. Hadamard MLP: Reemplaza las proyecciones densas tradicionales con una transformación de Walsh fija y diagonales aprendidas. Esto reduce significativamente el conteo de parámetros en la mezcla de canales, que es un cuello de botella en modelos pequeños.
2. Engram: Mueve el 'conocimiento del mundo' fuera del stack de atención a tablas de n-gramas hasheadas. Estas tablas se leen unas pocas filas por token, ofreciendo capacidad de conocimiento casi 'gratis' en tiempo de decodificación, crucial para dispositivos con memoria flash lenta y limitada.
3. Multi-lane Residual Streams: Proporcionan flexibilidad de enrutamiento a una red de 27 capas y 512 de ancho, con un costo de unos pocos productos punto por capa, evitando la necesidad de más volumen de atención o MLP.
4. Sistema de Memoria Bounded: La atención utiliza una ventana deslizante de 256 tokens para asegurar que el KV cache tenga un tamaño fijo y determinista (28MB). El prompt del sistema y las declaraciones de herramientas se 'fijan' como 'sinks' permanentes, impidiendo su desalojo y asegurando que el modelo nunca olvide sus capacidades.
5. Cuantificación de 2 bits (Cactus Quants): A diferencia de la cuantificación post-hoc, Needle 2 se entrena con Cactus Quants desde el pre-entrenamiento hasta el post-entrenamiento, aplicando la cuantificación a pesos, activaciones y el KV cache. Esto permite que el modelo de 2 bits desplegado sea el mismo que fue entrenado, minimizando la pérdida de calidad.
El motor de inferencia es un binario C++ auto-contenido y sin dependencias que incluye el modelo, el tokenizador y el compilador de gramática. Optimiza la velocidad al negarse a computar innecesariamente: los pesos 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 y la ruta aritmética en int8 de principio a fin. Una gramática compilada a nivel de byte restringe cada token, permitiendo al motor calcular puntuaciones de salida solo para tokens candidatos, saltándose hasta el 98% de la proyección del vocabulario en tokens estructurales. El binario universal detecta la CPU al inicio y selecciona sus kernels (SDOT, NEON, AVX2, RISC-V vectors, wasm SIMD o escalar), y el pool de hilos procesa las secciones seriales cortas de un token sin dormir, lo que casi duplica la velocidad de decodificación. Todas estas optimizaciones son exactas o validadas token por token contra una ruta de referencia para garantizar la fidelidad de la salida.
Flujo de Inferencia de Needle 2 en Dispositivo
- 1 Entrada de Usuario Texto de lenguaje natural (ej. 'enciende la luz de la sala')
- 2 Motor de Inferencia Binario C++ universal, auto-selecciona kernels (SDOT, NEON, AVX2, etc.)
- 3 Tokenizador Convierte el texto de entrada en tokens
- 4 Modelo Needle 2 (CQ2-bit) Procesa tokens, pesos de 2 bits expandidos en registros vectoriales
- 5 Grammar Compiler Restringe la generación de tokens de salida según el esquema de funciones
- 6 KV Cache (256-token window) Gestiona el estado de la sesión, prompt y herramientas fijadas como sinks
- 7 Generación de Salida Calcula logits solo para tokens válidos, int8 end-to-end
- 8 Llamada a Función Salida estructurada (ej. {function: 'turn_on_light', args: {room: 'sala'}})
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Needle 2 Engine (C++) | Motor de inferencia co-diseñado y optimizado para hardware de borde, incluyendo microcontroladores y CPUs sin aceleración de GPU/NPU. Maneja la carga del modelo, tokenización, aplicación de gramática y ejecución de kernels específicos de CPU. vs TensorFlow Lite, ONNX Runtime, PyTorch Mobile Binario universal que auto-selecciona kernels (SDOT, NEON, AVX2, RISC-V vectors, wasm SIMD) en tiempo de ejecución. Compilado single-threaded para bare metal o como librería estática para Cortex-M4/M7/M55. |
| storage | Cactus Quants (CQ2-bit) | Esquema de cuantificación de 2 bits integrado en el entrenamiento del modelo, aplicado a pesos, activaciones y KV cache. Permite un tamaño de modelo de 14MB y un consumo de RAM de 28MB, crucial para dispositivos con memoria flash y DRAM limitadas. vs Post-training quantization (PTQ) a int8/int4, Model pruning, Knowledge distillation Entrenamiento de extremo a extremo con cuantificación (QAT), asegurando que el modelo desplegado de 2 bits sea el mismo que fue entrenado. |
| compute | Hadamard MLP | Componente de la arquitectura del modelo que reemplaza las proyecciones densas tradicionales en el MLP con una transformación de Walsh fija y diagonales aprendidas, reduciendo drásticamente el conteo de parámetros para la mezcla de canales. vs MLP densos estándar, MLP con técnicas de sparse attention |
| compute | Engram (Hashed n-gram tables) | Mecanismo para almacenar 'conocimiento del mundo' fuera del stack de atención principal, permitiendo una recuperación de información casi sin costo computacional durante la decodificación. Esto libera parámetros del modelo para tareas de razonamiento y llamada a funciones. vs Aumentar el tamaño del modelo para incluir más conocimiento, Retrieval-Augmented Generation (RAG) con bases de datos externas Tablas de n-gramas hasheadas, leídas unas pocas filas por token. |
| cache | KV Cache con Ventana Deslizante de 256 tokens | Gestiona la memoria del estado de atención, asegurando que el consumo de RAM sea fijo y determinista (28MB) independientemente de la duración de la sesión. Los prompts del sistema y las declaraciones de herramientas se 'fijan' para evitar su desalojo. vs KV cache de tamaño ilimitado, Estrategias de eviction LRU/LFU sin fijación Sinks permanentes para herramientas y prompt del sistema. |
Trade-offs
Ganancias
- ▲▲ Eficiencia de memoria (RAM)
- ▲▲ Tamaño del modelo (binario)
- ▲ Latencia de inferencia
- ▲▲ Eficiencia energética (FLOPs/token)
- ▲ Privacidad y fiabilidad offline
Costes
- ▲ Capacidad de lenguaje general (chat, prosa abierta, conocimiento del mundo)
- △ Complejidad de entrenamiento (requiere QAT)
- ▲ Flexibilidad para tareas fuera de tool calling/extracción estructurada
Fundamentos Teóricos
El concepto de cuantificación de modelos para reducir su huella y acelerar la inferencia tiene raíces profundas en la investigación de redes neuronales, con trabajos pioneros como 'Quantized Neural Networks: Training Neural Networks with Low-Precision Weights and Activations' de Courbariaux et al. (2016) explorando la viabilidad de pesos y activaciones de baja precisión. La aproximación de Needle 2, donde la cuantificación es parte integral del proceso de entrenamiento (Quantization-Aware Training o QAT), se alinea con la evolución de estas técnicas para mitigar la pérdida de precisión inherente a la cuantificación post-hoc. Esto se relaciona con el principio de co-diseño de hardware/software, donde las limitaciones del hardware objetivo (memoria, ciclos de CPU) informan directamente el diseño del algoritmo y la arquitectura del modelo.
La idea de separar el 'conocimiento del mundo' del modelo principal, como se ve en el 'engram' de Needle 2, puede conectarse con conceptos de sistemas de memoria externa o bases de conocimiento en IA simbólica, donde la información se almacena y recupera de manera eficiente fuera de la red neuronal principal. Esto recuerda a los primeros sistemas expertos o a los enfoques de 'retrieval-augmented generation' (RAG) en LLMs, aunque implementado a una escala mucho más granular y optimizada para el borde. La gestión de la memoria con ventanas deslizantes y 'sinks' permanentes para el KV cache es una aplicación práctica de principios de gestión de memoria en sistemas operativos y bases de datos, adaptados a las restricciones de los transformadores y la inferencia en tiempo real.