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

1. Il Problema: Il Collo di Bottiglia dei Operandi Dinamici negli Acceleratori CIM

I modelli Transformer (come BERT e ViT) hanno rivoluzionato l'IA, ma il loro meccanismo di Self-Attention presenta una sfida fondamentale per gli acceleratori hardware basati su Compute-in-Memory (CIM).

• Limitazione degli operandi: L'architettura CIM tradizionale (basata su memorie non volatili come ReRAM o FeFET a singolo gate) supporta nativamente operazioni a due operandi (Input $\times$ Peso).
• Natura dinamica dell'Attention: Il meccanismo di Self-Attention richiede il calcolo di prodotti tra tre matrici dinamiche generate al momento dell'inferenza: Query ($Q$), Key ($K$) e Value ($V$). Nessuna di queste è un peso statico.
• Il Collo di Bottiglia "Write-Read-Write": Per eseguire l'attenzione su CIM convenzionale, le matrici dinamiche ($K^T$ e $V$) devono essere riscritte continuamente nella memoria non volatile (NVM) ad ogni inferenza. Poiché le operazioni di scrittura (programmare la NVM) sono:
  1. Lente (latenza ~50 ns contro ~10 ns per la lettura).
  2. Energeticamente costose.
  3. Limitate nella durata (endurance del dispositivo, tipicamente $10^6$-$10^{12}$ cicli).
     Questo ciclo di riscrittura continua degrada drasticamente il throughput, aumenta il consumo energetico e riduce la vita utile del dispositivo.
• Soluzioni esistenti insufficienti: Le architetture ibride attuali (es. X-Former) spostano il calcolo dell'attenzione su logica digitale CMOS per evitare le scritture NVM, ma sacrificano la densità e l'efficienza energetica promesse dal CIM "tutto in memoria".

2. Metodologia: L'Architettura TrilinearCIM

Gli autori propongono TrilinearCIM, un'architettura innovativa che esegue l'intera computazione dell'attenzione all'interno del nucleo di memoria senza riscritture dinamiche.

Dispositivo Hardware: DG-FeFET

Il cuore della soluzione è l'utilizzo di un Transistor a Effetto di Campo Ferroelettrico a Doppio Gate (DG-FeFET):

• Top-Gate (TG): Contiene lo strato ferroelettrico e memorizza i pesi statici ($W_Q, W_K, W_V$) in uno stato di conduttanza non volatile.
• Back-Gate (BG): È un terminale di controllo volatile (separato da un dielettrico standard) che modula dinamicamente la conduttanza del canale.
• Primitiva Trilineare: Sfruttando l'accoppiamento elettrostatico tra i due gate, la corrente di uscita segue la relazione:
  $$I_{DS} \approx V_{DS} \cdot G_0 \cdot (1 + \eta_{BG} \cdot V_{BG})$$
  Dove $G_0$ è il peso statico, $V_{DS}$ è l'input di attivazione e $V_{BG}$ è il segnale di modulazione dinamica. Questo permette di realizzare un'operazione a tre operandi ($A \cdot B \cdot C$) in un singolo passo fisico.

Flusso di Dati Trilineare

Invece di calcolare $Q$, $K$, $V$ separatamente e riscriverli, TrilinearCIM fonde le fasi di proiezione e calcolo dell'attenzione:

1. Generazione di Query Scalata: $X \cdot W_Q^T \cdot (1/\sqrt{d_k})$. Il fattore di scala è applicato staticamente tramite il back-gate.
2. Sintesi del Punteggio (Score): $R_1 \cdot W_K \cdot X^T$. Il vettore $X$ (o $X^T$) modula dinamicamente la conduttanza tramite il back-gate, evitando la formazione esplicita e la scrittura della matrice $K$.
3. Aggregazione dei Valori: $Score \cdot X \cdot W_V^T$. Il punteggio di attenzione (dopo Softmax digitale) modula il back-gate per aggregare i valori.

Vantaggi del Flusso:

• Zero scritture NVM runtime: I pesi rimangono statici; solo la tensione del back-gate cambia.
• Riduzione del Buffer: Le matrici intermedie dinamiche ($Q, K, V$) non devono essere memorizzate in buffer globali, riducendo la capacità richiesta di circa 3x.
• Eliminazione del traffico DRAM: Non è necessario trasferire dati tra memoria e calcolatore per le proiezioni intermedie.

3. Contributi Chiave

1. Primitiva CIM Trilineare: Introduzione di un'operazione a tre operandi ($Y = A \cdot B \cdot C$) abilitata dalle proprietà elettriche del DG-FeFET, superando il limite a due operandi degli acceleratori CIM standard.
2. Eliminazione delle Scritture NVM Dinamiche: Mappatura dei pesi statici sul gate superiore e delle attivazioni dinamiche sulla modulazione del back-gate, permettendo l'esecuzione completa dell'attenzione senza riprogrammazione della memoria.
3. Riduzione della Pressione sui Buffer: Fusione dei passaggi di proiezione e attenzione che elimina la necessità di memorizzare matrici intermedie dinamiche, riducendo significativamente i requisiti di memoria on-chip.
4. Framework di Valutazione TransCIM: Sviluppo di un framework di simulazione basato su NeuroSim che estende i modelli CIM esistenti per supportare operazioni trilineari e modelli Transformer completi.

4. Risultati Sperimentali

L'architettura è stata valutata su BERT-base (GLUE) e ViT-base (ImageNet, CIFAR) utilizzando il framework TransCIM.

• Efficienza Energetica e Latenza:

  • Rispetto a un CIM FeFET convenzionale (bilineare con riscritture), TrilinearCIM riduce il consumo energetico fino al 46.6% e migliora la latenza fino al 20.4% (per sequenze di 64 token).
  • Il guadagno di efficienza (TOPS/W) raggiunge fino al 39.2% in più rispetto alla baseline.
  • Il costo in area è del 37.3% in più (dovuto ai driver del back-gate e ai DAC per colonna), ma giustificato dai guadagni energetici.

• Accuratezza (NLP - GLUE):

  • TrilinearCIM supera la baseline bilineare in 7 su 9 task del benchmark GLUE.
  • Miglioramenti significativi su QNLI (+3.74%), MNLI (+3.14%) e STS-B (+1.74%).
  • La stabilità è maggiore: la deviazione standard dei risultati è molto più bassa rispetto alla baseline bilineare, grazie all'assenza di conversioni analogico-digitale ripetute e rimappature.

• Accuratezza (Vision - ViT):

  • Per i task di visione (CIFAR, ImageNet), TrilinearCIM mostra un leggero calo di accuratezza rispetto alla baseline bilineare.
  • Causa: La quantizzazione introdotta dal DAC del back-gate distorce i valori "outlier" critici nelle mappe di attenzione sparse dei modelli ViT, che sono più sensibili al rumore rispetto ai modelli NLP.

5. Significato e Impatto

• Risoluzione del Dilemma Endurance/Efficienza: TrilinearCIM risolve il compromesso fondamentale tra la necessità di evitare scritture NVM (per durata e velocità) e il desiderio di mantenere l'alta densità e l'efficienza del CIM. Non sacrifica la memoria non volatile per la logica digitale.
• Prima Architettura "All-in-Memory": È, a quanto ne sanno gli autori, la prima architettura in grado di eseguire l'intera computazione dell'attenzione Transformer esclusivamente nei nuclei NVM senza riscritture runtime.
• Scalabilità: L'architettura mantiene i suoi vantaggi energetici e di latenza anche all'aumentare della lunghezza della sequenza, dove il costo delle scritture dinamiche nelle architetture convenzionali diventerebbe proibitivo.
• Prospettive Future: Sebbene la simulazione sia promettente, il lavoro richiede validazione hardware per confermare l'affidabilità del back-gate e la gestione delle distorsioni di quantizzazione nei modelli di visione.

In sintesi, TrilinearCIM rappresenta un passo significativo verso acceleratori hardware efficienti per i Transformer, trasformando un limite fisico (la natura dinamica dell'attenzione) in un'opportunità attraverso l'ingegneria di dispositivi a doppio gate.