Resumo Técnico: PowerSoftmax: Rumo à Inferência Segura de LLMs sobre Dados Criptografados

Declaração do Problema

Modelos de Linguagem de Grande Porte (LLMs) modernos que preservam a privacidade frequentemente dependem de Criptografia Homomórfica (HE), especificamente o esquema CKKS, que permite computações sobre dados criptografados sem a necessidade de descriptografia. No entanto, o CKKS suporta apenas operações polinomiais. Arquiteturas Transformer padrão dependem fortemente de componentes não polinomiais, mais notavelmente a função Softmax na autoatenção, que envolve exponenciais, divisão e operações logarítmicas (por exemplo, para estabilidade numérica via o truque log-sum-exp).

Abordagens existentes para preencher essa lacuna enfrentam limitações significativas:

1. Aproximação Pós-Treinamento (PTA): Aproximação direta de funções não polinomiais em modelos pré-treinados usando polinômios de alto grau. Isso frequentemente resulta em inferência ineficiente devido à alta profundidade multiplicativa e ruído acumulado.
2. Treinamento Consciente de Aproximação (AAT): Substituição de primitivas não polinomiais por alternativas mais fáceis de aproximar (por exemplo, atenção ponto a ponto ou ReLU escalada) antes do treinamento. Embora eficazes para modelos menores, esses métodos lutaram com estabilidade e escalabilidade, falhando em suportar LLMs com bilhões de parâmetros ou arquiteturas profundas (por exemplo, 32 camadas).

O desafio central é projetar uma arquitetura Transformer que seja inerentemente compatível com as restrições de HE (polinômios de baixo grau, intervalos numéricos estáveis) enquanto mantém o poder expressivo, as capacidades de raciocínio e a aprendizagem em contexto (ICL) dos Transformers padrão baseados em Softmax na escala de bilhões de parâmetros.

Metodologia

Os autores propõem o PowerSoftmax, uma variante nova e amigável à HE da autoatenção projetada para replicar os princípios centrais do Softmax enquanto permanece adequada à aproximação polinomial. A metodologia envolve três estágios distintos:

1. Design de Atenção Amigável à HE

Os autores destilam as propriedades essenciais do Softmax: normalização para $[0, 1]$, amplificação das diferenças de pontuação e preservação da ordem. Eles substituem a função exponencial $\exp(x)$ no Softmax por uma função potência $x^p$ (onde $p$ é um inteiro par positivo).
O mecanismo de atenção proposto é definido como:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
Esta formulação evita exponenciais e problemas de divisão por zero inerentes a aproximações polinomiais diretas do Softmax.

2. Estabilidade de Treinamento e Robustez Numérica

Para abordar a instabilidade de treinamento causada pela função potência (onde entradas $>1$ explodem e entradas $<1$ desaparecem), os autores introduzem duas modificações-chave:

• Variante Estável: Inspirados no truque log-sum-exp, eles normalizam as entradas por uma constante $c = \|x\|_\infty + \delta$ antes de aplicar a função potência. Isso garante que as entradas permaneçam dentro de um intervalo gerenciável $(0, 1)$, prevenindo overflow/underflow e preservando a precisão de ponto flutuante.
• Divisão Limitada por $\epsilon$: Para lidar com a operação de divisão, que é não polinomial e sensível perto de zero, eles modificam o denominador para $\epsilon + \sum x_i^p$. Isso cria uma função Lipschitz contínua de $1/\epsilon^2$, tornando a subsequente aproximação polinomial da divisão significativamente mais tratável e reduzindo a profundidade de multiplicação necessária.

3. Aproximação Agnóstica ao Comprimento

A aproximação de divisão padrão (por exemplo, algoritmo de Goldschmidt) luta com a soma não limitada das pontuações de atenção, que escala linearmente com o comprimento da sequência $L$. Os autores propõem um PowerSoftmax Agnóstico ao Comprimento:
$$\text{Agnóstico ao Comprimento}(x)_j = \frac{x_j^p}{L \cdot \text{Média}(x^p)}$$
Como $L$ (comprimento da sequência) é público, o termo $1/L$ pode ser computado diretamente sem aproximação. O modelo aproxima a divisão da média em vez da soma, garantindo que o domínio de aproximação permaneça limitado independentemente do comprimento da sequência.

4. Pipeline de Construção (Algoritmo 1)

O pipeline completo para criar um LLM polinomial envolve:

1. Modificação Arquitetural: Substituição do Softmax pela variante PowerSoftmax Estável e treinamento do zero (ou via pré-treinamento contínuo) com objetivos de minimização de intervalo.
2. Minimização de Intervalo: Aplicação de uma função de perda especializada para restringir os intervalos de entrada das camadas não polinomiais (PowerSoftmax, GELU, LayerNorm) para garantir que aproximações polinomiais de baixo grau sejam precisas.
3. Substituição Polinomial: Pós-treinamento, substituindo as operações não polinomiais restantes (divisão, raiz quadrada inversa) por aproximações polinomiais de baixo grau (por exemplo, método de Goldschmidt) para produzir um modelo totalmente polinomial.

Contribuições Principais

1. Mecanismo PowerSoftmax: Uma nova variante de autoatenção que substitui a escala exponencial por escala polinomial, especificamente adaptada para ambientes HE enquanto mantém os princípios centrais de design da atenção.
2. LLMs Polinomiais Escaláveis: O desenvolvimento dos primeiros LLMs polinomiais que excedem um bilhão de parâmetros (especificamente 1,4 bilhão de parâmetros com 32 camadas). Esses modelos demonstram capacidades de raciocínio e aprendizagem em contexto (ICL) comparáveis a transformers padrão do mesmo tamanho.
3. Estratégias de Treinamento e Inferência Estáveis: Introdução de variantes de treinamento numericamente estáveis e estratégias de inferência agnósticas ao comprimento que superam os problemas de instabilidade e escalabilidade de abordagens anteriores de Transformer polinomial.
4. Avaliação Abrangente: Uma análise detalhada de latência para computações criptografadas e estudos de ablação validando as escolhas de design, incluindo o impacto do parâmetro de potência $p$ e da divisão limitada por $\epsilon$.

Resultados Experimentais

Os autores avaliaram seu método em arquiteturas causais (GPT/Pythia) e bidirecionais (RoBERTa):

• Desempenho: O modelo polinomial de 1,4 bilhão de parâmetros alcançou desempenho comparável aos modelos Pythia originais baseados em Softmax em benchmarks zero-shot e few-shot (por exemplo, Lambada, PIQA, WinoGrande, ARC). Notavelmente, o modelo demonstrou desempenho competitivo em benchmarks de raciocínio como ARC-Challenge, uma capacidade previamente não observada em LLMs polinomiais.
• Modelos Bidirecionais: Uma variante polinomial do RoBERTa-Base mostrou apenas uma degradação de ~1% em comparação com o RoBERTa original baseado em Softmax em tarefas GLUE (SST-2, QNLI, MNLI), superando abordagens polinomiais anteriores como Nexus.
• Inferência HE: Experimentos no modelo Pythia 70M usando o framework HElayers confirmaram que a inferência criptografada produz resultados quase idênticos à inferência em texto puro (perda MSE máxima de 0,005).
• Eficiência: O método proposto reduz significativamente a profundidade de multiplicação em comparação com alternativas state-of-the-art (Nexus, Cho et al.). Para uma sequência de 128 tokens, o PowerSoftmax reduziu a latência por um fator de ~9,7 em comparação com abordagens Taylor/Goldschmidt e ~20 em comparação com abordagens de estratégia quadrática.

Significado e Alegações

O artigo alega um avanço no campo da aprendizagem de máquina que preserva a privacidade, demonstrando que LLMs totalmente polinomiais podem escalar para o regime de bilhões de parâmetros sem sacrificar capacidades de raciocínio ou aprendizagem em contexto.

• Escalabilidade: Diferentemente de trabalhos anteriores limitados a modelos pequenos ou tarefas simples de classificação, esta abordagem treina com sucesso transformers profundos (32 camadas) com 1,4 bilhão de parâmetros.
• Segurança Não Interativa: Ao utilizar criptografia homomórfica totalmente (FHE) com um modelo totalmente polinomial, a solução permite inferência não interativa, evitando a sobrecarga de comunicação e vulnerabilidades potenciais associadas à computação multipartidária (MPC) ou protocolos interativos.
• Praticidade: A análise detalhada de latência e o sucesso do despliegue de um modelo de 1,4 bilhão de parâmetros sugerem que a inferência segura e privada de LLMs está se tornando praticamente viável, abrindo caminho para aplicações em setores sensíveis a dados, como saúde e finanças.

Os autores reconhecem que, embora os modelos performem bem, uma avaliação completa das capacidades generativas auto-regressivas na decodificação sequencial sobre dados criptografados permanece uma direção para trabalho futuro.