Fast Exact Nearest-Neighbor Learning for High-Frequency Financial Time Series

Resumen Técnico: Aprendizaje de Vecinos Más Cercanos Exacto y Rápido para Series Temporales Financieras de Alta Frecuencia

Planteamiento del Problema
La IA financiera está entrando en una era de escalabilidad donde los modelos deben ingerir corpus históricos masivos (acciones, ETFs, FX, opciones) para soportar el trading en tiempo real, la gestión de riesgos y la valoración de derivados. Existe un cuello de botella crítico en el stack actual: mientras que las herramientas basadas en Python (NumPy, scikit-learn) ofrecen flexibilidad, sufren de sobrecarga del intérprete, saturación del ancho de banda de memoria y una pobre localidad de caché al consultar cientos de miles de estados históricos. Por el contrario, las soluciones en C++ y FPGA ofrecen velocidad pero crean una brecha entre la "investigación y la producción" debido a la necesidad de codebases separadas y optimización manual. El desafío específico abordado es la inferencia de vecinos más cercanos exactos (KNN) para series temporales de alta frecuencia, donde el costo lineal de la búsqueda por fuerza bruta ($O(n \cdot d)$) se vuelve prohibitivo para los presupuestos de latencia en tiempo real, aunque los métodos aproximados (como HNSW) sacrifican la exactitud requerida para ciertas aplicaciones financieras.

Metodología
Los autores proponen un árbol k-d con SIMD en Mojo diseñado específicamente para corpus de series temporales financieras. Mojo, un lenguaje de sistemas compilado con sintaxis compatible con Python, permite el co-diseño de algoritmo, disposición de memoria y hardware sin abandonar el ecosistema de Python. La metodología integra tres optimizaciones principales:

1. División Basada en la Varianza: A diferencia de las divisiones estándar basadas en la mediana, el árbol selecciona el eje de división basado en la dimensión con la máxima varianza ($j^* = \arg\max \text{Var}(\{x_{i,j}\})$). Esto concentra la poda temprana en las características financieras más discriminativas (retornos, volatilidad, momentum), creando subregiones más ajustadas.
2. Almacenamiento de Buffer Plano Contiguo: Los vectores de entrenamiento se almacenan en un buffer plano de ordenamiento de filas ($X \in \mathbb{R}^{n \times d}$) indexado por un arreglo permutado. Los nodos hoja se definen mediante desplazamientos enteros, lo que permite que los candidatos se lean como bloques de memoria contiguos. Esto elimina la sobrecarga de persecución de punteros (pointer-chasing) y asegura la alineación de la línea de caché.
3. Kernels Vectorizados en Tiempo de Compilación: Los cálculos de distancia utilizan vectorización SIMD (Instrucción Única, Múltiples Datos). El kernel se especializa en tiempo de compilación (comptime) para dimensiones específicas (por ejemplo, $d=16$ o $d=24$) y anchos de SIMD, generando código totalmente desenrollado y libre de escalares. La distancia euclidiana al cuadrado se descompone en bloques paralelos del ancho del registro.

El sistema fue evaluado en ocho conjuntos de datos financieros (acciones de EE. UU., ETFs, FX) en dos arquitecturas de hardware: x86 (Intel/AMD) y ARM64 (Apple M3). Las comparaciones se realizaron contra las implementaciones de fuerza bruta y k-d tree de scikit-learn, así como el enfoque de fuerza bruta SIMD de propio Mojo.

Contribuciones Clave

1. Implementación de Árbol k-d SIMD en Mojo: Una implementación novedosa que combina la poda basada en la varianza, la disposición de memoria de buffer plano y los kernels de distancia vectorizados en tiempo de compilación para KNN exacto en grandes conjuntos de datos financieros.
2. Análisis de Tiempo de Ejecución Teórico: Una demostración de que, bajo supuestos de poda estándar, el árbol k-d de Mojo domina asintóticamente tanto a la fuerza bruta SIMD de Mojo como al k-d tree de scikit-learn en el régimen de acciones fijas, gran-$n$ y dimensiones moderadas.
3. Co-diseño Finanzas-Hardware: Un enfoque estandarizado de ingeniería de características donde las características derivadas de OHLCV se alinean con los anchos SIMD (16 y 24 valores float32 mapeando a 2 y 3 anchos de registro, respectivamente), preservando la interpretabilidad financiera mientras se maximiza el rendimiento (throughput).
4. Validación Empírica: Benchmarks exhaustivos que muestran aceleraciones significativas a través de diversos activos y plataformas de hardware, con significancia estadística confirmada mediante pruebas de rango con signo de Wilcoxon.
5. Demostración de Escalabilidad: Un experimento secundario utilizando Extra Trees para la valoración de volatilidad implícita, mostrando que las ganancias de rendimiento permiten entrenar con 10 veces más datos, mejorando directamente la precisión de la predicción.

Resultados

• Aceleración en x86 ($d=16$): El árbol k-d de Mojo logró una aceleración de 17.5–21.6$\times$ sobre el k-d tree de scikit-learn y una aceleración de 1.2–1.3$\times$ sobre la fuerza bruta de scikit-learn en conjuntos de datos que oscilan entre 143K y 277K muestras.
• Aceleración en ARM64 ($d=24$): El método logró una aceleración de 28.1–43.5$\times$ sobre la fuerza bruta de scikit-learn en Apple M3. Notablemente, en ARM64, el kernel de fuerza bruta SIMD de Mojo fue de 780–1549$\times$ más lento que el k-d tree debido a un desajuste entre el ancho SIMD de tiempo de compilación (optimizado para registros de 256 bits de x86) y los registros nativos de 128 bits de ARM. Sin embargo, la poda algorítmica del k-d tree ($O(n^{1-1/d})$) lo hizo inmune a esta penalización específica del hardware, manteniendo ganancias masivas.
• Exponentes de Escalamiento: El ajuste empírico de la complejidad de tiempo $T(n) = Cn^\alpha$ arrojó un exponente de 1.755 para el árbol k-d de Mojo (solo consulta), comparado con 1.885 para fuerza bruta y 2.320 para el k-d tree de scikit-learn. Esto confirma un escalamiento sub-cuadrático que se amplía a medida que aumenta el tamaño del conjunto de datos.
• Calidad de Predicción: En el experimento de Extra Trees, entrenar con 200,000 muestras (10$\times$ el límite base de Python) redujo el RMSE de la Volatilidad Implícita (IV) de los Puts en un 8.0% (de 0.361 a 0.333). Separar calls y puts redujo aún más el RMSE en un 27–31%.
• Significancia Estadística: Una prueba de rango con signo de Wilcoxon confirmó que el árbol k-d de Mojo es estrictamente más rápido que el k-d tree de scikit-learn en los ocho conjuntos de datos ($p < 0.01$), con una media geométrica de aceleración de 10.8$\times$ (IC 95%: [5.94$\times$, 16.62$\times$]).

Significancia y Reivindicaciones
El artículo afirma ser el primer estudio de nivel de sistema, compatible con Python y de código abierto sobre la eficiencia de la inferencia de IA financiera a escala. Su principal significancia radica en demostrar que el co-diseño entre algoritmo, disposición de memoria, hardware y lenguaje puede generar una aceleración operativamente significativa sin sacrificar la exactitud ni la compatibilidad con Python.

Los autores argumentan que Mojo cierra la brecha entre la flexibilidad de la investigación en Python y el rendimiento de los sistemas de producción en C++. Al permitir un único código base que logra un rendimiento cercano al de C++ (aproximadamente el 60% de la velocidad de C++ optimizado en una comparación específica) mientras ofrece una aceleración de 10$\times$ sobre las librerías estándar de Python, este trabajo posiciona a Mojo como un stack escalable y listo para producción para la IA financiera. Los resultados sugieren que, a medida que los volúmenes de datos financieros crecen hacia escalas de nivel de tick, el escalamiento sub-cuadrático de este enfoque será cada vez más crítico para mantener las restricciones de latencia en tiempo real. Los principios también se reclaman aplicables a otros campos intensivos en datos que requieren la recuperación de vecinos más cercanos exactos, como la genómica y la ciberseguridad.