Trilinear Compute-in-Memory Architecture for Energy-Efficient Transformer Acceleration

Título: Arquitetura Trilinear de Computação em Memória para Aceleração Eficiente de Energia de Transformers

1. O Problema: O Gargalo da Dinâmica em Transformers

Os modelos Transformer (como BERT e ViT) revolucionaram a IA, mas sua eficiência energética é severamente limitada pelo mecanismo de Auto-atenção.

• Incompatibilidade com CIM Tradicional: As arquiteturas de Computação em Memória (CIM) convencionais baseiam-se na premissa de que os pesos (matrizes) são estáticos após o treinamento. No entanto, o mecanismo de atenção gera operandos dinâmicos ($Q$, $K$, $V$) para cada entrada de sequência.
• Ciclo Custo-Energia: Em CIM tradicional, isso força um ciclo de "Computar-Escrever-Computar". As matrizes dinâmicas devem ser reprogramadas repetidamente na memória não volátil (NVM) a cada inferência.
• Consequências:
  • Latência e Energia: Operações de escrita em NVM são ordens de magnitude mais lentas e consomem muito mais energia do que leituras.
  • Durabilidade: O desgaste cíclico (endurance) das células de memória é rapidamente excedido, limitando a vida útil do dispositivo.
  • Soluções Atuais Insuficientes: Abordagens existentes ou reduzem, mas não eliminam, as escritas (via decomposição de matriz), ou migram o cálculo de atenção para CMOS digital, sacrificando a densidade e eficiência da memória.

2. Metodologia: A Arquitetura TrilinearCIM

Os autores propõem o TrilinearCIM, uma arquitetura que realiza a computação completa da atenção dentro do núcleo de memória (NVM) sem reprogramação dinâmica.

• Dispositivo Chave: DG-FeFET (Transistor de Efeito de Campo Ferroelétrico de Dupla Porta):
  • Utiliza uma estrutura com duas portas: uma Porta Superior (Top-Gate) com camada ferroelétrica para armazenar o peso não volátil ($W$) e uma Porta Traseira (Back-Gate) dielétrica para modulação volátil.
  • A tensão na porta traseira ($V_{BG}$) modula linearmente a condutância do canal, atuando como um terceiro operando volátil.
• Primitiva Trilinear:
  • A arquitetura executa uma operação de multiplicação-acumulação (MAC) de três operandos: $Y = A \cdot B \cdot C$.
  • Mapeamento:
    • $B$ (Pesos estáticos $W_Q, W_K, W_V$): Armazenados na condutância não volátil da porta superior.
    • $A$ e $C$ (Ativações dinâmicas $X, X^T$ e Scores): Aplicados como tensões na porta traseira e nas linhas de entrada.
• Fluxo de Dados Fusionado:
  • Em vez de calcular $Q$, $K$ e $V$ separadamente e armazená-los, o TrilinearCIM funde as etapas de projeção e atenção.
  • Estágio 1 (Geração de Query Escalada): $X \cdot W_Q \cdot (1/\sqrt{d_k})$. A escala é aplicada via porta traseira estática.
  • Estágio 2 (Síntese de Scores): $R_1 \cdot W_K \cdot X^T$. O vetor $X^T$ modula a condutância dinamicamente, gerando os scores de atenção sem formar a matriz intermediária $K$.
  • Estágio 3 (Agregação de Valor): $Score \cdot X \cdot W_V$. O score modula a condutação para agregar os valores.
• Redução de Tráfego de Memória: Elimina a necessidade de armazenar e buscar tensores intermediários ($Q, K, V$) na memória DRAM externa, reduzindo o buffer global em ~3x.

3. Contribuições Principais

1. Primitiva CIM Trilinear: Introdução de uma operação de três operandos nativa baseada nas propriedades elétricas do DG-FeFET, estendendo o paradigma CIM tradicional de dois operandos.
2. Eliminação de Escrita Dinâmica em NVM: A computação de atenção ocorre inteiramente dentro da matriz de memória sem reprogramar o estado ferroelétrico durante a inferência, resolvendo os problemas de latência de escrita e desgaste de endurance.
3. Fluxo de Dados Simplificado: Fusão das etapas de projeção e atenção reduz a pressão sobre o buffer e elimina o tráfego de dados para memória externa.
4. Recuperação de Eficiência: Restaura a vantagem de "Tudo em Memória" para cargas de trabalho de Transformer, superando as limitações das arquiteturas híbridas (NVM + CMOS).

4. Resultados Experimentais

A avaliação foi realizada no framework TransCIM (baseado em NeuroSim) utilizando modelos BERT-base (GLUE) e ViT-base (ImageNet/CIFAR).

• Desempenho e Eficiência (BERT-base):
  • Energia: Redução de até 46,6% no consumo de energia por inferência em comparação com CIM FeFET convencional (bilinear).
  • Latência: Melhoria de até 20,4% na latência.
  • Throughput: Aumento de 25,5% em inferências por segundo.
  • Overhead de Área: Aumento de 37,3% na área do chip devido aos circuitos de driver da porta traseira e DACs, mas compensado pela maior eficiência energética.
• Precisão (GLUE):
  • O TrilinearCIM superou a abordagem bilinear em 7 de 9 tarefas do GLUE (ex: +3,74% em QNLI, +3,14% em MNLI).
  • A estabilidade foi maior (desvio padrão menor), pois elimina erros de conversão analógico-digital repetidos e remapeamento de tensores.
• Desempenho em Visão (ViT):
  • Observou-se uma queda de precisão em tarefas de visão (CIFAR/ImageNet) em comparação com a abordagem bilinear. Isso foi atribuído à quantização adicional introduzida pelo DAC da porta traseira, que distorce picos de atenção esparsos críticos para imagens.
• Estudos de Ablação:
  • A configuração 1-bit/célula com ADC de 6 bits (1b/6b) mostrou-se o ponto ótimo para precisão, oferecendo a maior melhoria de precisão sobre o bilinear com menor sobrecarga de área (32,4%).
  • Arrays de sub-matriz 32x32 ofereceram a maior redução de latência e utilização de memória quase perfeita.

5. Significado e Conclusão

O TrilinearCIM representa um avanço significativo na aceleração de hardware para Transformers:

• Viabilidade de "All-in-Memory": Demonstra que é possível realizar a computação completa da atenção (incluindo a parte dinâmica $Q \cdot K^T$) dentro de núcleos NVM sem sacrificar a durabilidade ou a latência.
• Solução para o Gargalo de Escrita: Resolve o problema fundamental de que as escritas dinâmicas em NVM são o principal limitador de eficiência para modelos de atenção.
• Trade-off Aceitável: O aumento de área (37,3%) é justificado pelas grandes economias de energia e ganhos de latência, especialmente para cenários de NLP e contextos longos onde o volume de escrita seria proibitivo.
• Futuro: Embora a precisão em visão ainda necessite de otimização (talvez via fine-tuning consciente de ruído), a arquitetura abre caminho para aceleradores de Transformer verdadeiramente eficientes e escaláveis, eliminando o gargalo de von Neumann para dados dinâmicos.

Em resumo, o trabalho propõe uma mudança de paradigma, utilizando a modulação de porta traseira em FeFETs para transformar a limitação de operandos dinâmicos em uma vantagem de computação trilinear in-memory.