Resumo Técnico: Embeddings de Fourier Primos (PFE)

Declaração do Problema

Embeddings neurais padrão frequentemente falham em expor a estrutura algébrica inerente dos inteiros, particularmente no que diz respeito à aritmética modular. Embora o Teorema Chinês dos Restos (CRT) dite que a adição módulo um número composto $N$ se decompõe em adições independentes módulo seus fatores primos, os embeddings aprendidos de forma padrão não refletem naturalmente essa decomposição. Consequentemente, modelos treinados em tarefas de aritmética modular frequentemente sofrem de "grokking" — um fenômeno onde os modelos passam por longas fases de memorização antes de subitamente generalizar. Esse atraso é atribuído a uma falha representacional: quando o input carece de uma estrutura local prima explícita, a rede deve reconstruir essas relações algébricas do zero via gradiente descendente, um processo que é ávido por dados e frágil.

Abordagens existentes possuem limitações: codificações contínuas (ex: xVal) preservam a magnitude, mas descartam a periodicidade, enquanto abordagens periódicas como os Fourier Number Embeddings (FoNE) frequentemente utilizam bases compostas (ex: base-10) que emaranham estruturas primas distintas (ex: 2-ádica e 5-ádica) dentro de canais compartilhados, falhando em isolar a informação local prima.

Metodologia: Prime Fourier Embeddings (PFE)

Os autores introduzem o Prime Fourier Embeddings (PFE), um esquema de embedding principiado derivado da análise harmônica dos números racionais $\mathbb{Q}$.

• Construção: O PFE codifica um inteiro $a$ como uma coleção de pares $(\cos, \sin)$ indexados por primos. Para um primo $p$ e uma profundidade $d$ (onde $0 \le d < D$), o embedding é definido como:
  $$\text{PFE}_{p,d}(a) = \left[ \cos\left(\frac{2\pi a}{p^{d+1}}\right), \sin\left(\frac{2\pi a}{p^{d+1}}\right) \right] \in \mathbb{R}^2$$
  Isso corresponde à decomposição real do caractere $\chi_1$ do grupo $\mathbb{Z}/p^{d+1}\mathbb{Z}$.
• Base Teórica: O design é motivado pela dualidade de Pontryagin e pela fatoração de caracteres adélicos. O embedding implementa diretamente caracteres $p$-ádicos locais, garantindo que a estrutura de resíduo de qualquer inteiro seja legível a partir do correspondente canal primo sem acoplamento entre primos distintos.
• Implementação: Em experimentos, os autores utilizam um conjunto fixo de 16 primos ($P_0 = \{3, \dots, 59\}$) e um limite de profundidade de $D=3$, resultando em uma dimensão de embedding de 96. As características do embedding são determinísticas e não treináveis.

Contribuições Teóricas Principais

O artigo estabelece uma garantia estrutural rigorosa para qualquer mapa linear operando sobre o PFE:

1. Teorema de Decomposição em Bloco-Diagonal (Teorema 3.1): Qualquer mapa linear $W$ que seja equivariante em relação à ação do grupo produto no PFE deve ser block-diagonal.
   • Existe um bloco independente por primo $p$ e profundidade $d$.
   • Há zero acoplamento entre blocos de primos distintos.
   • Há zero acoplamento entre blocos de profundidades distintas dentro do mesmo primo.
2. Mecanismo de Prova: Este resultado segue do Lema de Schur. Cada bloco PFE representa uma representação irredutível não isomórfica (um caractere) do grupo aditivo $\mathbb{Z}$. Como os caracteres para pares $(p, d)$ distintos são não isomórficos, qualquer intertwiner equivariante entre eles deve ser zero.
3. Previsão do CRT: Para uma tarefa com módulo $N = p_1 \cdots p_k$, a teoria prevê que apenas os $k$ canais primos correspondentes aos fatores de $N$ carregarão informação relevante para a tarefa, enquanto todos os outros canais permanecerão irrelevantes.

Resultados Experimentais

Os autores validam essas previsões através de estudos de ablação sistemáticos em tarefas de adição modular $(a+b) \mod N$.

• Experimento 1 (Tarefas de Primo Único): Modelos foram treinados em tarefas com módulos primos variados. Quando a linha prima específica correspondente ao módulo da tarefa era ablada (zerada), a precisão de teste caía significativamente. Em contraste, ablar linhas primas irrelevantes causava quedas de precisão próximas de zero.
  • Resultado: As razões de especialização (queda diagonal vs. queda off-diagonal) excederam 500$\times$ na maioria das configurações, confirmando que a rede utiliza apenas o canal primo relevante.
• Experimento 2 (Módulos Compostos): Modelos foram treinados em módulos compostos livres de quadrados (ex: $N=15, 21, 105$).
  • Resultado: Ablar fatores primos de $N$ causou grandes quedas de precisão, enquanto ablar não-fatores causou quedas negligenciáveis. A razão fator/não-fator também excedeu 500$\times$.
  • Precisão: Os modelos alcançaram 100% de precisão de teste in-distribution em todos os módulos compostos livres de quadrados testados e intervalos de entrada (desde que o intervalo de entrada fosse suficientemente grande em relação a $N$).

Significância e Alegações

O artigo afirma que o PFE transforma o aprendizado de aritmética modular de um problema de descoberta representacional para um problema de seleção.

• Principiado vs. Arbitrário: Ao contrário de escolhas de frequência heurísticas, a base prima é matematicamente fundamentada na decomposição de Fourier adélica de $\mathbb{Q}$. Isso garante que a estrutura algébrica necessária esteja presente desde a inicialização.
• Interpretabilidade Mecanística: O trabalho demonstra que o design de embedding pode ser guiado pela matemática dos objetos que ele representa. A estrutura interna da rede é caracterizada antes do treinamento (via teoria da representação) e verificada após o treinamento (via ablações direcionadas).
• Teoria Falseável: O Teorema de Decomposição em Bloco-Diagonal fornece uma restrição estrutural provável sobre o comportamento da rede. A confirmação empírica dessa restrição oferece um caminho para um raciocínio numérico mais confiável e transparente em redes neurais, indo além da observação post-hoc para uma teoria matematicamente fundamentada.

Os autores observam que, embora a estrutura block-diagonal seja provada, o mecanismo específico pelo qual o gradiente descendente converge para uma solução equivariante (em vez de quebrar a simetria) permanece uma questão aberta para trabalhos futuros.