ParaRNN: Apple destrava treino paralelo de RNNs não-lineares e escala pra 7B parâmetros

Pesquisadores da Apple publicaram o ParaRNN, framework que resolve um problema antigo de arquitetura: como treinar RNN (Recurrent Neural Network) em paralelo sem abrir mão da não-linearidade. O paper foi aceito como Oral no ICLR 2026 e o código saiu como open-source.

A ideia central: RNNs são ótimas na inferência (custo constante por token, independente do tamanho do contexto), mas historicamente impossíveis de escalar porque o treino é sequencial por natureza. Transformers dominaram justamente por permitirem paralelização no eixo da sequência. O ParaRNN ataca esse gargalo de frente e entrega 665× de speedup sobre a aplicação sequencial tradicional.

Por que RNN ainda importa

O custo do mecanismo de atenção no transformer cresce quadraticamente com o tamanho da sequência. Numa RNN, o forward pass custa o mesmo independente do quanto de contexto veio antes. Isso significa geração de token em tempo constante na inferência: atraente pra deploy em ambiente com recurso limitado (edge, on-device, celular da Apple, por exemplo).

O problema sempre foi o treino. A estrutura recorrente precisa ser desenrolada passo a passo, enquanto atenção processa tudo de uma vez.

O compromisso que SSMs fizeram

Arquiteturas modernas como Mamba (SSMs, state space models) contornaram isso simplificando a recorrência pra ser puramente linear no hidden state:

• SSMs: h_l = A_l * h_{l-1} + B_l * x_l (linear)
• RNNs clássicas: h_l = σ_l(h_{l-1}, x_l) (não-linear)

Linearidade permite usar parallel scan (operações associativas dá pra combinar em árvore), transformando O(L) em O(log₂ L). Só que linearidade limita a expressividade do modelo: certas dinâmicas de hidden state não dá pra representar bem.

A sacada: método de Newton

O ParaRNN usa o método de Newton (técnica clássica pra resolver equações não-lineares iterativamente) pra ter os dois mundos. Em vez de encarar a RNN como cadeia sequencial, reformula a sequência inteira como um sistema único de equações, onde todos os hidden states são incógnitas a resolver simultaneamente.

Newton lineariza o sistema usando os Jacobianos (derivadas locais), e o sistema linearizado tem exatamente a mesma forma de um SSM linear. Os Jacobianos fazem o papel das matrizes de estado A_l. Daí cada iteração resolve um SSM via parallel scan.

Nos experimentos, três iterações de Newton bastam pra convergir em GRU e LSTM adaptadas. Ou seja: com três parallel scans bem projetados, recupera o comportamento da RNN não-linear completa, em fração do tempo.

Engenharia pra escalar

Pra RNN genérica, os Jacobianos são densos: armazenamento cresce quadraticamente e multiplicação cubicamente com o tamanho do hidden state. Inviável em escala.

A Apple seguiu o design de SSMs modernos e introduziu ParaGRU e ParaLSTM: adaptações com matrizes diagonais na definição da célula. Resultado: Jacobianos diagonais (ParaGRU) ou bloco-diagonais 2x2 (ParaLSTM).

Além disso, escreveram kernels CUDA customizados fully-fused: Newton iterations, montagem do sistema e parallel reduction num único kernel. Competitivo com a implementação do Mamba em todos os tamanhos de sequência testados.

Resultados: RNN clássica volta ao jogo

Treinaram modelos de 400M até 7B parâmetros. Na tabela de avaliação (lm-eval-harness), em 7B:

• Mamba2 (6.96B): PPL 8.62
• ParaLSTM (6.76B): PPL 9.16
• ParaGRU (6.76B): PPL 9.19
• Transformer (6.89B): PPL 9.55

Nas tarefas downstream (Arc-C, HellaSwag, OpenBookQA, WinoGrande, PiQA, MMLU), ParaGRU e ParaLSTM ficam no mesmo patamar de transformer e Mamba2.

Ponto interessante: em tarefas sintéticas de state tracking (MQAR, k-hop, Parity), as RNNs não-lineares batem tudo com 100%, enquanto Mamba2 fica em 51% no Parity e Transformer em 53%. A não-linearidade entrega expressividade que o modelo linear simplesmente não alcança.

Como usar

O framework abstrai tudo. Você implementa só o passo de recorrência:

class MyRNNCell(BaseRNNCell):
    def recurrence_step(self, h, x, system_parameters):
        h_new = ...  # σ(h, x; params)
        return h_new

O framework cuida de Newton, montagem de Jacobiano, parallel reduction e otimizações pra estruturas específicas. Três níveis de performance: PyTorch puro (protótipo), CUDA-accelerated (Jacobiano diagonal ou bloco-diagonal) e fully-fused (kernel único, exige implementação CUDA do passo de recorrência).

Pra quem pesquisa arquitetura de sequência, isso é ferramenta nova pra explorar um espaço de design que estava congelado há anos por limitação computacional, não por limitação teórica. O trade-off entre não-linearidade e eficiência de treino deixa de ser fundamental.

Código aberto no repositório oficial do paper. ICLR 2026 (que por sinal rola no Rio de Janeiro) deve ter apresentação oral com mais detalhes.