Prism: Dynamic Sparse Attention cho Huấn luyện Mô hình Tạo Video-Âm thanh Đồng thời Độ phân giải Gốc 2K

Hướng dẫn đọc

Vấn đề và động lực

Huấn luyện nguyên bản các mô hình diffusion tạo video-âm thanh đồng thời ở độ phân giải 2K cho phép tổng hợp các chi tiết không gian phong phú, động lực chuyển động sắc nét và âm thanh đồng bộ. Tuy nhiên, sau khi nén bằng VAE, một clip 2K tạo ra từ hàng trăm nghìn đến hơn một triệu token. Chuẩn full attention gặp phải hai vấn đề nghiêm trọng:

1. Độ phức tạp bậc hai: Chi phí tính toán và bộ nhớ khiến việc huấn luyện trở nên đắt đỏ đến mức bất khả thi.
2. Sự pha loãng trọng số attention (Attention weight dilution): Lượng lớn token nền lặp đi lặp lại chiếm ưu thế trong chuỗi token, làm loãng ngân sách attention phân bổ cho nội dung cục bộ giàu thông tin (chẳng hạn như khuôn mặt cử động, bàn tay và các vật thể phát ra âm thanh) và phá vỡ các prior đa phương thức đã học trước.

Các phương pháp sparse attention trước đây hoặc áp dụng xấp xỉ không cần huấn luyện (training-free), hoặc chỉ tập trung vào video không có tiếng ở độ phân giải thấp hơn (480p/720p), hoặc áp đặt các hình dạng khối cứng nhắc, đồng nhất. Vì sự đồng bộ âm thanh-thị giác tập trung nghiêm ngặt vào các vùng phát âm thanh nhỏ gọn, các hình dạng khối cố định sẽ trộn lẫn các token khác biệt lại với nhau, làm hỏng các biểu diễn khối được mean-pool và không thích ứng được với các mẫu tương tác đa phương thức.

Các ý tưởng cốt lõi

• Phân vùng Macro-Zone: Thể tích token không-thời gian được chia thành các macro-zone kích thước $Z_T\times Z_H\times Z_W=8\times 8\times 8$ để tương thích với các cấu hình khối được căn chỉnh cho Triton.
• Chỉ số biến thiên dẫn dắt kép (Dual-Guidance Variation Indicator): Mỗi zone đo lường động sự biến thiên nội dung theo hướng bằng cách sử dụng phương sai theo kênh video cùng với các chuẩn audio-to-video cross-attention.
• Hình dạng khối 3D thích ứng nội dung: Các zone chọn từ một tập ứng viên gồm 16 hình dạng với tổng thể tích $B\in \{64,128,256\}$, phân vùng mịn dọc theo các trục có sự thay đổi nhanh và liên kết âm thanh mạnh mẽ.
• Lựa chọn lai Top-$k$ / Top-$p$: Sparse attention cấp khối sử dụng Top-$k$ để đảm bảo ngữ cảnh cơ sở chắc chắn, được mở rộng thích ứng thông qua Top-$p$ cho các query có phân phối chú ý phân tán.

Phương pháp

Prism hoạt động trên một Diffusion Transformer (DiT) hai nhánh được khởi tạo từ MOVA-720p, thay thế video self-attention đầy đủ bằng dynamic sparse attention trong khi vẫn giữ nguyên cross-attention và audio self-attention.

1. Xây dựng chỉ số biến thiên

Đối với mỗi macro-zone $m$, tỷ lệ biến thiên video theo hướng dọc theo trục $d\in \{T,H,W\}$ được lấy từ các đặc trưng value đã được mean-pool ${\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)}$$

Từ các chuẩn audio-to-video cross-attention $a_{t,h,w}^m=\Vert {\hat{\mathbf{h}}}_{t,h,w}^m{\Vert }_2$ được lưu trong cache từ lớp trước đó, cường độ liên kết âm thanh $\bar{a}(m)\in [0,1]$ và phương sai theo hướng chuẩn hóa ${\hat{v}}_{a,d}(m)$ được tính toán. Chỉ số biến thiên tổng hợp là: $$g_d(m)=r_{v,d}(m)+\bar{a}(m)\cdot {\hat{v}}_{a,d}(m)+ϵ$$

2. Định lượng thể tích và hình dạng khối

Tổng thể tích khối $B(m)\in \{64,128,256\}$ được gán bằng cách đánh giá mật độ thông tin không-thời gian chuẩn hóa $\rho (m)$ dựa trên các ngưỡng ${\tau }_{128}=0.5$ và ${\tau }_{256}=0.25$. Việc giảm thiểu tổn thất thông tin nội khối ${\sum }_d{g}_d(b_d^2-1)$ thỏa mãn ${\prod }_d{b}_d=B(m)$ mang lại chiều dài cạnh liên tục tối ưu $b_d^{\ast }\propto g_d^{-1/2}$. Trong không gian ${\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)$$ Hình dạng khối $s^{\ast }=(b_T,b_H,b_W)$ được chọn từ tập ứng viên ${\mathcal{C}}_{B(m)}$ bằng cách giảm thiểu khoảng cách Euclid tới ${\ell }_d^{\ast }$.

3. Cơ chế chú ý lai Top-$k$ / Top-$p$

Các khối key được mean-pool thành các đại diện ${\bar{\mathbf{K}}}_j$, được chấm điểm với các khối query ${\mathbf{Q}}_i$ để tạo ra xác suất khối ${\mathbf{P}}_i$. Tập hợp được chú ý ${\mathcal{S}}_i$ kết hợp ngưỡng xác suất cố định và tích lũy: $${\mathcal{S}}_i=\text{Top-}k({\mathbf{P}}_i,k)\cup \text{Top-}p({\mathbf{P}}_i,p)$$ Một Triton kernel online softmax sẽ thực thi sparse attention chỉ trên ${\mathcal{S}}_i$.

Kết quả

• Vượt trội trên Benchmark: Trên 2K-Bench (300 video thử nghiệm 2K gốc), Prism đạt $0.61$ Aesthetic Quality (AQ), $0.982$ Temporal Flickering (TF), $0.63$ DeSync, và $0.89$ MotionQ, vượt trội đáng kể so với MOVA (Full Attn: $0.42$ AQ, $1.05$ DeSync, $0.53$ MotionQ) và LTX-2.3 có siêu phân giải ($0.48$ AQ, $0.95$ DeSync, $0.66$ MotionQ).
• Hiệu quả tính toán: Tốc độ huấn luyện nhanh hơn $2.5\times$ so với full attention ở 2K ($10.6\text{ phuˊt}$ so với $26.5\text{ phuˊt}$ mỗi step). Khi suy luận trên 4 GPU, Prism chạy trong $331\text{ giaˆy}$ sử dụng $38.0\text{ GB}$ VRAM mỗi GPU, so với $1091\text{ giaˆy}$ và $72.8\text{ GB}$ của MOVA.
• Mở rộng theo độ phân giải: Prism là khung duy nhất được thử nghiệm cho thấy sự cải thiện đơn điệu trên tất cả các chỉ số chất lượng và đồng bộ đa phương thức khi mở rộng từ 480p lên 2K, trong khi full attention và các baseline thưa thớt khối cố định bị suy giảm chất lượng ở 2K do bão hòa token nền.
• Đánh giá từ người dùng: Trong nghiên cứu cảm nhận 2AFC với 200 người tham gia trên 30 clip 2K, Prism đạt tỷ lệ ưa chuộng $>87\%$ so với tất cả các baseline mã nguồn mở xét trên các khía cạnh: độ khớp văn bản, chất lượng âm thanh, chi tiết thị giác, độ chân thực của chuyển động và khớp khẩu hình.