Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

1. El Problema: La Brecha entre Representaciones Continuas y Lógica Discreta

El objetivo de la extracción de algoritmos es sintetizar programas ejecutables directamente a partir de modelos entrenados en tareas algorítmicas, permitiendo el descubrimiento de algoritmos de novo sin depender de código escrito por humanos.

Sin embargo, aplicar este paradigma a la arquitectura Transformer dominante presenta un desafío fundamental:

• Entrelazamiento de Representaciones: Los Transformers estándar sufren de "superposición" (superposition), donde las características se codifican en direcciones superpuestas y no ortogonales en lugar de neuronas individuales.
• Incompatibilidad: Esto crea representaciones internas altamente entrelazadas y polisémicas, lo que hace imposible la recuperación directa de expresiones simbólicas explícitas (código legible) mediante regresión simbólica, a diferencia de lo que ocurre con las Redes Neuronales Recurrentes (RNN) donde se han logrado éxitos previos (ej. método MIPS).

2. Metodología: El Discrete Transformer

Los autores proponen el Discrete Transformer, una arquitectura diseñada específicamente para cerrar la brecha entre las representaciones continuas y la lógica simbólica discreta.

A. Arquitectura Especializada

El modelo se basa en el modelo de computación RASP (Restricted Access Sequence Processing) y se compone de tres módulos clave con disociación funcional estricta:

1. Residual Stream Numérico: A diferencia de los vectores latentes densos, el flujo residual se define como una colección de variables escalares explícitas. Esto preserva el historial computacional como coordenadas disociadas, evitando la interferencia de información.
2. Numerical Attention (Enrutador): Actúa estrictamente como un operador de enrutamiento de información. Utiliza una codificación lineal por partes (Piecewise Linear Encoding) para proyectar escalares y un mecanismo de atención "duro" (hard attention) mediante muestreo controlado por temperatura. Su función es recuperar valores específicos del historial (ej. "copiar el valor del token en la posición $t-1$").
3. Numerical MLP (Operador): Dedicado exclusivamente a transformaciones aritméticas y lógicas elementales. Se descompone en sub-módulos paralelos e independientes que realizan operaciones simples (ej. suma, multiplicación) sobre un número fijo y pequeño de operandos.
4. Cabeza de Salida Lineal: Agrega los pasos computacionales discretos para generar la predicción final.

B. Optimización y Extracción

• Muestreo con Recocido de Temperatura (Temperature Annealing): Se introduce un mecanismo de muestreo diferenciable en cada módulo. Durante el entrenamiento, la temperatura ($\tau$) se reduce gradualmente (de 10.0 a 0.1). Esto permite al modelo explorar el espacio de soluciones en un estado continuo y luego converger suavemente a un grafo computacional totalmente discreto.
• Pipeline de Extracción Modular: Una vez que el modelo converge a un estado discreto, se extrae el algoritmo mediante:
  1. Pruebas de Hipótesis (para Attention): Se analizan los patrones de atención para identificar si siguen patrones interpretables como "Desplazamiento Fijo" (Fixed Offset) o "Extremos en Ventana" (Windowed Extrema).
  2. Regresión Simbólica (para MLP): Se aplica regresión simbólica (usando la herramienta PySR) a los sub-módulos MLP para inferir las expresiones aritméticas exactas (ej. $y = 2x_1 + x_2$).
  3. Ensamblaje Global: Se reconstruye el programa completo mediante un recorrido hacia atrás en el grafo computacional, combinando los punteros de atención y las expresiones aritméticas en código Python ejecutable.

3. Contribuciones Clave

1. Arquitectura Discrete Transformer: Un diseño alineado con RASP que fuerza la disociación funcional entre el enrutamiento y el cálculo aritmético, permitiendo la optimización continua hacia la lógica discreta.
2. Pipeline de Extracción Híbrido: Desarrollo de un método que combina pruebas de hipótesis estadísticas y regresión simbólica para extraer programas de variables continuas, superando las limitaciones de los métodos anteriores que requerían aproximaciones de variables latentes discretas.
3. Controlabilidad mediante Sesgos Inductivos: Demostración de que las restricciones arquitectónicas actúan como sesgos inductivos fuertes, permitiendo a los investigadores guiar el descubrimiento de algoritmos hacia formas específicas o paradigmas alternativos.

4. Resultados Experimentales

El modelo fue evaluado en una suite diversa de tareas de razonamiento algorítmico, incluyendo aritmética lineal, composición no lineal y dinámicas continuas.

• Rendimiento Predictivo: El Discrete Transformer logra un rendimiento comparable al método MIPS (basado en RNN), con errores cuadráticos medios (MSE) cercanos a cero en la mayoría de las tareas.
• Fidelidad de Extracción: Los programas extraídos son 100% precisos (Accuracy = 1.00) en tareas discretas y logran un RMSE extremadamente bajo en tareas continuas.
  • Ejemplo: En la tarea parity_last2, el modelo extrae la expresión algebraica exacta $y_t = x_t + x_{t-1} - 2x_t x_{t-1}$, que es equivalente a la lógica XOR.
  • Ejemplo: En maximum_prev2, descubre la identidad $max(x_t, x_{t-1}) = x_{t-1} + ReLU(x_t - x_{t-1})$.
• Ventaja sobre MIPS: A diferencia de MIPS, que falla en datos de punto flotante debido a su dependencia de estados latentes discretos, el Discrete Transformer maneja nativamente variables continuas, permitiendo la extracción de algoritmos para dinámicas físicas (ej. caída libre, media móvil exponencial).
• Robustez: El modelo converge consistentemente a soluciones ejecutables incluso con variaciones en los hiperparámetros, siempre que la arquitectura cumpla los requisitos funcionales mínimos.

5. Significado e Impacto

• Interpretabilidad del Transformer: Este trabajo demuestra que es posible sintetizar programas ejecutables y legibles por humanos directamente desde los pesos de un Transformer, resolviendo el problema del entrelazamiento de características mediante un diseño arquitectónico específico.
• Descubrimiento Científico: Proporciona un marco para el descubrimiento de algoritmos de novo que no están sesgados por el conocimiento humano previo, capaz de revelar múltiples caminos lógicos equivalentes para resolver un mismo problema.
• Controlabilidad: Establece un nuevo paradigma donde la arquitectura del modelo no solo define su capacidad de aprendizaje, sino que actúa como un mecanismo de control para dirigir la forma y la estructura de los algoritmos descubiertos.

En resumen, el Discrete Transformer transforma la "caja negra" del Transformer en un grafo computacional transparente, permitiendo la extracción de algoritmos matemáticos y lógicos precisos, abriendo nuevas vías para la interpretabilidad mecánica y el descubrimiento automatizado de algoritmos.