The Discrete-Log Clock: How a Transformer Learns Modular Multiplication

Resumen Técnico: El Reloj de Logaritmo Discreto

Planteamiento del Problema
Si bien los mecanismos internos de los transformadores que aprenden la suma modular ($a + b \mod p$) han sido totalmente revertidos mediante representaciones de Fourier dispersas (el algoritmo "Clock"), el mecanismo para la multiplicación modular ($a \cdot b \mod p$) ha permanecido elusivo. Estudios empíricos previos (Doshi et al., 2024; Furuta et al., 2024) reportaron que los transformadores entrenados en multiplicación aprenden espectros de Fourier "densos" que requieren todas las frecuencias, lo que contrasta drásticamente con la estructura dispersa encontrada en la suma. Esto sugería que la multiplicación podría depender de un algoritmo fundamentalmente diferente y menos interpretable.

Metodología
Los autores entrenan un transformador de decodificador de una sola capa para predecir $a \cdot b \mod 113$ para entradas $a, b \in \{1, \dots, 112\}$. El modelo se entrena con el 30% de los datos y debe generalizar al 70% restante, exhibiendo el fenómeno de "grokking" donde la generalización ocurre repentinamente tras un largo periodo de memorización.

La innovación metodológica central es un cambio en la base de análisis. En lugar de aplicar la transformada de Fourier discreta (DFT) aditiva estándar a los embeddings de los tokens, los autores aplican la transformada de carácter multiplicativo. Esto implica:

1. Reconocer que los enteros no nulos módulo un primo $p$ forman un grupo cíclico $(\mathbb{Z}/p\mathbb{Z})^*$.
2. Utilizar una raíz primitiva $g$ (aquí $g=3$) para mapear las entradas a sus logaritmos discretos: $\alpha = \log_g(a)$ y $\beta = \log_g(b)$.
3. Reordenar las filas de la matriz de embeddings de acuerdo con estos logaritmos discretos.
4. Aplicar la DFT estándar a estos datos reordenados.

Este enfoque aprovecha el isomorfismo de grupo donde la multiplicación en el espacio original corresponde a la suma en el espacio del logaritmo discreto ($a \cdot b \equiv g^{\alpha + \beta} \mod p$).

Resultos Clave

• Dispersión Espectral: Cuando se analiza en la base aditiva, el espectro de los embeddings parece plano y denso (coeficiente de Gini 0.07). Sin embargo, en la base de carácter multiplicativo, el espectro se vuelve altamente disperso (coeficiente de Gini 0.58), con la energía concentrada en solo 4 frecuencias clave ($\{2, 8, 47, 56\}$).
• Sintonización de Neuronas: En la base multiplicativa, el 96.9% (496/512) de las neuronas del MLP están "limpiamente sintonizadas" a una sola frecuencia, con más del 85% de su energía de activación concentrada allí. En la base aditiva, cero neuronas cumplen este umbral.
• Visualización Estructural: Cuando los mapas de calor de activación de las neuronas se reordenan por logaritmo discreto, revelan claras franjas diagones sinusoidales 2D. Esto indica que las neuronas computan funciones de la forma $\cos(k(\alpha + \beta))$, realizando efectivamente la suma en el espacio logarítmico.
• Análisis SVD: Los componentes principales de la matriz de embeddings, cuando se reordenan por logaritmo discreto, exhiben estructuras sinusoidales claras, confirmando la codificación de caracteres multiplicativos.

El Algoritmo del "Reloj de Logaritmo Discreto"
El artículo concluye que el transformador implementa un algoritmo análogo al algoritmo "Clock" para la suma, denominado el Reloj de Logaritmo Discreto. El mecanismo procede de la siguiente manera:

1. Embedding: Los enteros se embeben como funciones sinusoidales de sus logaritmos discretos.
2. Enrutamiento: Los mecanismos de atención dirigen estos embeddings a la posición de salida.
3. Composición: El MLP aplica identidades trigonométricas para computar sumas de los logaritmos discretos (efectivamente sumando $\alpha + \beta$).
4. Decodificación: El modelo califica los candidatos de salida alineándolos con la suma computada en el espacio logarítmico, donde la respuesta correcta logra una interferencia constructiva a través de las frecuencias clave.

Significado y Reivindicaciones
La contribución principal es demostrar que el "espectro denso" observado en trabajos previos es un artefacto de analizar los datos en la base incorrecta, no un reflejo de la complejidad algorítmica. Al hacer coincidir la base de análisis con la estructura algebraica de la tarea (el grupo multiplicativo), los autores revelan que la multiplicación se reduce a la suma en el espacio del logaritmo discreto.

El artículo afirma que esta metodología es generalizable:

• El espectro multiplicativo disperso fue confirmado en 10 primos diferentes ($p=59$ a $113$).
• El mecanismo es universal a través de semillas aleatorias (aprox. 80% de las semillas).
• Resultados preliminares sugieren que el enfoque se extiende a la exponenciación modular ($a^b \mod p$) utilizando un transformador de 2 capas.

Los autores enfatizan que para cualquier tarea algebraica, aplicar la transformada de Fourier sobre el grupo relevante (en lugar de la DFT aditiva estándar) es esencial para revelar una estructura interpretable donde las herramientas estándar anteriormente solo veían ruido.