Veloce e Semplice: Attenzione 2-Simplessiale in Triton

Abstract

Recenti lavori hanno dimostrato che la perdita durante l'addestramento scala come una legge di potenza sia con la dimensione del modello che con il numero di token, e che per ottenere modelli ottimali in termini di calcolo è necessario scalare insieme la dimensione del modello e il conteggio dei token. Tuttavia, queste leggi di scala presuppongono una fornitura infinita di dati e si applicano principalmente in contesti limitati dal calcolo. Poiché i moderni modelli linguistici di grandi dimensioni si basano sempre più su enormi dataset di scala internet, l'assunzione che siano limitati dal calcolo sta diventando meno valida. Questo cambiamento evidenzia la necessità di architetture che privilegino l'efficienza dei token. In questo lavoro, indaghiamo l'uso del Transformer 2-simpliciale, un'architettura che generalizza l'attenzione a prodotto scalare standard a funzioni trilineari attraverso un'implementazione efficiente del kernel Triton. Dimostriamo che il Transformer 2-simpliciale raggiunge una migliore efficienza dei token rispetto ai Transformer standard: per un budget fisso di token, modelli di dimensioni simili superano le loro controparti a prodotto scalare in compiti che coinvolgono matematica, codifica, ragionamento e logica. Quantifichiamo questi guadagni dimostrando che l'attenzione 2-simpliciale modifica l'esponente nelle leggi di scala per compiti di conoscenza e ragionamento rispetto all'attenzione a prodotto scalare.

English

Recent work has shown that training loss scales as a power law with both model size and the number of tokens, and that achieving compute-optimal models requires scaling model size and token count together. However, these scaling laws assume an infinite supply of data and apply primarily in compute-bound settings. As modern large language models increasingly rely on massive internet-scale datasets, the assumption that they are compute-bound is becoming less valid. This shift highlights the need for architectures that prioritize token efficiency. In this work, we investigate the use of the 2-simplicial Transformer, an architecture that generalizes standard dot-product attention to trilinear functions through an efficient Triton kernel implementation. We demonstrate that the 2-simplicial Transformer achieves better token efficiency than standard Transformers: for a fixed token budget, similarly sized models outperform their dot-product counterparts on tasks involving mathematics, coding, reasoning, and logic. We quantify these gains by demonstrating that 2-simplicial attention changes the exponent in the scaling laws for knowledge and reasoning tasks compared to dot product attention.

Veloce e Semplice: Attenzione 2-Simplessiale in Triton

Fast and Simplex: 2-Simplicial Attention in Triton

Abstract

Support