Prism: Dynamic Sparse Attention untuk Pelatihan Model Pembangkitan Bersama Video-Audio Native 2K

Panduan membaca

Masalah dan motivasi

Melatih model difusi video-audio gabungan secara native pada resolusi 2K memungkinkan sintesis detail spasial yang kaya, dinamika gerakan yang tajam, dan suara yang tersinkronisasi. Namun, setelah kompresi VAE, klip 2K menghasilkan ratusan ribu hingga lebih dari satu juta token. Standar full attention mengalami dua masalah fatal:

1. Kompleksitas kuadratik: Jejak komputasi dan memori membuat biaya pelatihan menjadi terlampau mahal.
2. Pengenceran bobot atensi (attention weight dilution): Token latar belakang berulang dalam jumlah masif mendominasi urutan token, mengencerkan anggaran atensi yang seharusnya dialokasikan untuk konten informatif yang terlokalisasi (seperti wajah yang bergerak, tangan, dan objek penghasil suara) serta merusak prior multimodal prapelatihan.

Pendekatan sparse attention sebelumnya umumnya menerapkan aproksimasi tanpa pelatihan (training-free), hanya berfokus pada video bisu pada resolusi lebih rendah (480p/720p), atau memaksakan bentuk blok yang kaku dan seragam. Karena sinkronisasi audio-visual terkonsentrasi secara ketat pada area penghasil suara yang kompak, bentuk blok yang kaku mencampurkan token-token yang berbeda, merusak representasi blok hasil mean-pooling, dan gagal beradaptasi dengan pola interaksi lintas modalitas.

Ide utama

• Partisi Makro-Zona (Macro-Zone Partitioning): Volume token spatiotemporal dipartisi menjadi makro-zona berukuran $Z_T\times Z_H\times Z_W=8\times 8\times 8$ agar kompatibel dengan konfigurasi blok berbasis Triton.
• Indikator Variasi Panduan Ganda (Dual-Guidance Variation Indicator): Setiap zona mengukur variasi konten terarah secara dinamis menggunakan varians kanal video bersamaan dengan norm cross-attention audio-ke-video.
• Bentuk Blok 3D Adaptif Konten: Zona memilih dari serangkaian kandidat 16 bentuk dengan volume total $B\in \{64,128,256\}$, mempartisi secara lebih halus di sepanjang sumbu perubahan cepat dan kopling audio yang intens.
• Pemilihan Hibrida Top-$k$ / Top-$p$: Sparse attention tingkat blok menggunakan Top-$k$ untuk memastikan konteks dasar yang terjamin, yang diperluas secara adaptif via Top-$p$ untuk query dengan distribusi atensi yang menyebar.

Metode

Prism beroperasi pada Diffusion Transformer (DiT) dua cabang yang diinisialisasi dari MOVA-720p, menggantikan self-attention video penuh dengan dynamic sparse attention sembari mempertahankan cross-attention dan self-attention audio tanpa modifikasi.

1. Formulasi Indikator Variasi

Untuk setiap makro-zona $m$, rasio variasi video terarah di sepanjang sumbu $d\in \{T,H,W\}$ diturunkan dari fitur value hasil mean-pooling ${\bar{\mathbf{V}}}_{m,i}^{(d)}$: $${\sigma }_{v,d}^2(m)=\frac{1}{D}\sum _{c=1}^D\frac{1}{Z_d}\sum _{i=1}^{Z_d}{\left({\bar{V}}_{m,i,c}^{(d)}-{\mu }_{m,c}^{(d)}\right)}^2,r_{v,d}(m)=\frac{{\sigma }_{v,d}^2(m)}{{\sum }_{d^{\prime }}{\sigma }_{v,d^{\prime }}^2(m)}$$

Dari norm cross-attention audio-ke-video $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ yang di-cache dari lapisan sebelumnya, kekuatan kopling audio $\bar{a}(m)\in [0,1]$ dan varians terarah ternormalisasi ${\hat{v}}_{a,d}(m)$ dihitung. Indikator variasi komposit dirumuskan sebagai: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Kuantisasi Volume dan Bentuk Blok

Volume blok total $B(m)\in \{64,128,256\}$ ditetapkan dengan mengevaluasi densitas informasi spatiotemporal ternormalisasi $\rho (m)$ terhadap ambang batas ${\tau }_{128}=0.5$ dan ${\tau }_{256}=0.25$. Meminimalkan kehilangan informasi intra-blok ${\sum }_d{g}_d(b_d^2-1)$ dengan batasan ${\prod }_d{b}_d=B(m)$ menghasilkan panjang sisi kontinu optimal $b_d^{\ast }\propto g_d^{-1/2}$. Dalam ruang ${\log }_2$ ($S={\log }_{2}B(m)$): $${\ell }_d^{\ast }=-\frac{1}{2}{\log }_2{g}_d+\frac{1}{3}\left(S+\frac{1}{2}{\log }_2{g}_T+\frac{1}{2}{\log }_2{g}_H+\frac{1}{2}{\log }_2{g}_W\right)$$ Bentuk blok $s^{\ast }=(b_T,b_H,b_W)$ dipilih dari rangkaian kandidat ${\mathcal{C}}_{B(m)}$ dengan meminimalkan jarak Euclidean ke ${\ell }_d^{\ast }$.

3. Atensi Hibrida Top-$k$ / Top-$p$

Blok key dirata-ratakan melalui mean-pooling menjadi representasi ${\bar{\mathbf{K}}}_j$, lalu diskor terhadap blok query ${\mathbf{Q}}_i$ untuk menghasilkan probabilitas blok ${\mathbf{P}}_i$. Himpunan yang diatensi ${\mathcal{S}}_i$ menggabungkan ambang batas probabilitas tetap dan kumulatif: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Kernel Triton online softmax menjalankan sparse attention hanya pada himpunan ${\mathcal{S}}_i$.

Hasil

• Keunggulan Tolok Ukur: Pada 2K-Bench (300 video uji native 2K), Prism mencapai $0.61$ Aesthetic Quality (AQ), $0.982$ Temporal Flickering (TF), $0.63$ DeSync, dan $0.89$ MotionQ, melampaui MOVA secara substansial (Full Attn: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) dan LTX-2.3 dengan super-resolution ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Peningkatan Efisiensi: Kecepatan pelatihan $2,5\times$ lebih cepat dibandingkan full attention pada resolusi 2K ($10,6\text{ menit}$ vs. $26,5\text{ menit}$ per langkah). Dalam inferensi 4-GPU, Prism berjalan dalam $331\text{ detik}$ menggunakan VRAM $38,0\text{ GB}$ per GPU, dibandingkan MOVA yang membutuhkan $1091\text{ detik}$ dan $72,8\text{ GB}$.
• Skalabilitas Resolusi: Prism adalah satu-satunya kerangka kerja yang diuji yang menunjukkan peningkatan monotonik di seluruh metrik kualitas dan sinkronisasi lintas modalitas saat diskalakan dari 480p ke 2K, sedangkan full attention dan baseline sparse blok tetap mengalami penurunan performa pada 2K akibat saturasi token latar belakang.
• Studi Pengguna: Dalam studi persepsi 2AFC terhadap 200 partisipan pada 30 klip 2K, Prism memperoleh preferensi $>87\%$ di atas semua baseline open-source dalam hal keselarasan teks, kualitas audio, detail visual, realisme gerakan, dan sinkronisasi bibir.

Keterbatasan dan pertanyaan terbuka

• Makro-Zona Temporal Tetap: Makro-zona memiliki jendela temporal statis $Z_T=8$. Pergantian adegan cepat (shot cuts) atau transisi zoom dalam jendela ini memicu disparitas fitur mendadak yang tidak dapat direpresentasikan oleh satu bentuk blok saja, menyebabkan artefak distorsi sementara.
• Penskalaan Resolusi dan Model: Prism baru dievaluasi hingga resolusi 2K pada arsitektur backbone 16B parameter; penskalaan ke native 4K (yang melipatgandakan jumlah token sebanyak empat kali) dan eksplorasi arsitektur DiT 30B+ yang lebih besar masih belum diteliti.
• Urutan Ketergantungan Audio-Visual: Karena self-attention mendahului cross-attention di dalam blok DiT, panduan cross-attention audio bergantung pada norm yang di-cache dari lapisan sebelumnya, yang memicu keterlambatan propagasi sinyal panduan sebesar satu lapisan.