La geometria conta nell’attenzione 3D su sequenze impacchettate

Abstract degli autori

Compattare un volume 3D in una sequenza di token non è semplicemente un problema di modellazione di sequenze lunghe. Una volta compattato il volume, l’ordine della sequenza non codifica più la geometria del campione: un meccanismo posizionale indicizzato in base alla posizione nella sequenza può quindi apprendere la disposizione in memoria anziché l’oggetto. È una modalità di errore generale dell’attenzione su dati geometrici compattati. Introduciamo PACK3D, un operatore di attenzione locale-globale che separa la disposizione in memoria dalla geometria dell’attenzione. PACK3D mantiene esatta l’attenzione sul vicinato locale, comprime soltanto il contesto distante in riepiloghi di blocchi 3D contigui e indicizza ogni interazione query-key con RoPE assiale sulle coordinate effettive della griglia dei token, anziché sulle posizioni nella sequenza compattata. Poiché la rotazione viene applicata prima che un blocco sia riassunto, il logit tra query e riepilogo rimane una somma ponderata degli spostamenti 3D relativi alla query; il contesto globale compresso resta quindi interpretabile spazialmente dopo la compattazione, invece di ridursi a contenuto aggregato anonimo. Misuriamo direttamente questa modalità di errore. In un compito sintetico controllato — addestrare su una modalità di compattazione, cambiare soltanto l’ordine in memoria dei token e mantenere invariata la regola di etichettatura fisica — gli embedding posizionali appresi e RoPE applicata alla sequenza compattata risolvono perfettamente la disposizione di training, ma crollano quando si ricompattano i token con ordinamento Morton, in blocchi contigui o casuale; RoPE assiale basata sulle coordinate della griglia, invece, resta esatta in tutti i casi. Lo stesso sistema di coordinate si applica ai dati reali: nell’analisi al microscopio di hiPSC Allen WTC-11, PACK3D supera una U-Net 3D con un numero di parametri corrispondente nel compito congiunto sulle 15 strutture e una baseline locale-globale con posizione appresa e caratteristiche corrispondenti; i controlli isolano nel sistema di coordinate, e non nella capacità dell’encoder, la causa del miglioramento. Un’implementazione Triton fusa eguaglia la latenza dell’attenzione a finestre con 65k token, fornendo al contempo il contesto a lungo raggio che le finestre non possono offrire: l’operatore non è quindi soltanto corretto, ma anche addestrabile su larga scala.