Transformer più snelli possono imparare facilmente a fare clustering

Abstract degli autori

I transformer possiedono capacità di apprendimento nel contesto: alcuni algoritmi di apprendimento noti possono essere eseguiti durante il passaggio in avanti del modello. Studi recenti mostrano che i transformer possono eseguire esattamente l’algoritmo di Lloyd per il clustering $k$-means di $n$ punti in $d$ dimensioni, con una dimensione dell’embedding $d_{\textsf{emb}} = d+k$ (che richiede quindi matrici di proiezione dell’attenzione di dimensione $(d+k)^2$). In questo lavoro sviluppiamo tale risultato in diversi modi. Innanzitutto, presentiamo un transformer altrettanto espressivo ma più piccolo, che esegue l’algoritmo di Lloyd con una dimensione dell’embedding $d_{\textsf{emb}} = (d + \lceil \log_2 k \rceil)$. Addestriamo poi questi transformer ad apprendere gli algoritmi di clustering a partire da una distribuzione di compiti di clustering, e caratterizziamo teoricamente e convalidiamo empiricamente i fattori che influenzano la convergenza e la generalizzazione all’interno della distribuzione degli algoritmi di apprendimento basati su gradienti stocastici. Infine, mettiamo alla prova le capacità generali di clustering di questi algoritmi appresi (sotto forma di transformer) e cerchiamo di capire in quali situazioni funzionano e in quali falliscono.