Hình ảnh có thể được tạo ra trên điện thoại di động trong 0,2 giây, là tốc độ nhanh nhất hiện tại, Google tạo ra một mô hình phân tán siêu nhanh MobileDiffusion

Hình ảnh có thể được tạo ra trên điện thoại di động trong 0,2 giây, là tốc độ nhanh nhất hiện tại, Google tạo ra một mô hình phân tán siêu nhanh MobileDiffusion

Chạy các mô hình trí tuệ nhân tạo sinh lớn như Stable Diffusion trên các thiết bị di động như điện thoại di động đã trở thành một trong những điểm nóng được theo đuổi bởi ngành công nghiệp, trong đó tốc độ tạo ra là ràng buộc chính.

Gần đây, một bài báo từ Google, “MobileDiffusion: Tạo hình ảnh từ văn bản trong vòng một giây trên các thiết bị di động,” đề xuất Tạo hình ảnh từ văn bản nhanh nhất trên các thiết bị di động, chỉ mất 0,2 giây trên iPhone 15 Pro. Bài báo này đến từ cùng một nhóm với UFOGen. Trong khi tạo ra một mô hình phân tán siêu nhỏ, nó cũng áp dụng con đường công nghệ Diffusion GAN phổ biến hiện tại để tăng tốc quá trình lấy mẫu.

Địa chỉ bài báo: https://arxiv.org/abs/2311.16567

Dưới đây là các kết quả được tạo ra bởi MobileDiffusion trong một bước.

Vậy, MobileDiffusion được tối ưu hóa như thế nào?

Hãy bắt đầu với câu hỏi, tại sao cần tối ưu hóa.

Tạo hình ảnh từ văn bản phổ biến nhất hiện nay dựa trên mô hình phân tán. Dựa trên khả năng tạo hình ảnh cơ bản mạnh mẽ của mô hình được huấn luyện trước và tính ổn định của nó trong các nhiệm vụ điều chỉnh sau cùng, chúng ta đã thấy hiệu suất phi thường của mô hình phân tán trong các ứng dụng như chỉnh sửa ảnh, tạo hình ảnh có thể điều chỉnh, tạo hình ảnh cá nhân và tạo hình ảnh video.

Tuy nhiên, như một Mô hình Cơ bản, nhược điểm của nó cũng rõ ràng, chủ yếu bao gồm hai khía cạnh: thứ nhất, số lượng lớn tham số của mô hình phân tán dẫn đến tính toán chậm, đặc biệt là khi tài nguyên bị hạn chế; thứ hai, mô hình phân tán yêu cầu nhiều bước lấy mẫu, điều này dẫn đến tốc độ suy luận rất chậm. Lấy mô hình Stable Diffusion1.5 (SD) phổ biến nhất làm ví dụ. Mô hình cơ bản của nó chứa gần 1 tỷ tham số. Chúng tôi đã định lượng mô hình trên iPhone 15 Pro để suy luận, và 50 bước lấy mẫu mất gần 80 giây. Yêu cầu tài nguyên đắt đỏ và trải nghiệm người dùng chậm chạp này giới hạn rất nhiều các kịch bản ứng dụng trên di động.

Để giải quyết các vấn đề trên, MobileDiffusion tối ưu hóa từng điểm. (1) Đối với vấn đề kích thước mô hình lớn, chúng tôi chủ yếu tiến hành nhiều thí nghiệm và tối ưu hóa trên thành phần cốt lõi UNet của nó, bao gồm đặt các hoạt động tính toán tốn kém và hoạt động chú ý trên các tầng thấp hơn, và tối ưu hóa hoạt động trên các thiết bị di động, chẳng hạn như các chức năng kích hoạt, v.v. (2) Đối với vấn đề mô hình phân tán yêu cầu lấy mẫu nhiều bước, MobileDiffusion khám phá và thực hiện các công nghệ suy luận một bước như Progressive Distillation và UFOGen hiện đại nhất hiện tại.

Tối ưu hóa mô hình

MobileDiffusion được tối ưu hóa dựa trên SD1.5UNet, hiện đang là mô hình phổ biến nhất trong cộng đồng mã nguồn mở. Sau mỗi hoạt động tối ưu hóa, mất hiệu suất so với mô hình UNet gốc sẽ được đo đạc cùng một lúc. Các chỉ số đo lường bao gồm hai chỉ số thông dụng: FID và CLIP.

Thiết kế tổng thể

Phía trái của hình ảnh trên là sơ đồ thiết kế của UNet gốc. Có thể thấy rằng nó bao gồm chủ yếu là Convolution và Transformer, và Transformer cũng bao gồm Self-Attention và Cross-Attention.

Ý tưởng cốt lõi của việc tối ưu hóa UNet của MobileDiffusion được chia thành hai điểm: 1) Convolution được tinh gọn. Như chúng ta đã biết, Convolution trên không gian đặc trưng có độ phân giải cao rất tốn thời gian và có số lượng tham số lớn. Điều này liên quan đến Convolution đầy đủ; 2) Cải thiện hiệu suất Attention. Giống như Convolution, Attention cao yêu cầu tính toán độ dài của toàn bộ không gian đặc trưng. Độ phức tạp của Self-Attention liên quan trực tiếp đến độ dài được làm phẳng của không gian đặc trưng, và Cross-Attention cũng tỷ lệ thuận với độ dài của không gian.

Các thí nghiệm cho thấy việc di chuyển 16 Transformers của UNet vào lớp trong cùng với độ phân giải đặc trưng thấp nhất và cắt bỏ một Convolution trong mỗi lớp sẽ không ảnh hưởng đáng kể đến hiệu suất. Hiệu quả đạt được là: MobileDiffusion có thể giảm 22 Convolution và 16 Transformers gốc xuống còn 11 Convolution và khoảng 12 Transformers, và tất cả các Attention này đều được thực hiện trên bản đồ đặc trưng có độ phân giải thấp, vì hiệu suất sẽ cải thiện rất nhiều. Đây là một cải tiến lớn, mang lại cải thiện hiệu suất khoảng 40% và giảm tham số khoảng 40%. Mô hình cuối cùng được hiển thị ở phía bên phải của hình ảnh trên. So sánh với nhiều mô hình khác nhau như sau:

Thiết kế chi tiết

Chỉ có một số thiết kế mới lạ sẽ được giới thiệu ở đây. Độc giả quan tâm có thể đọc văn bản để có một sự giới thiệu chi tiết hơn.

Tách riêng Self-Attention và Cross-Attention

Transformer trong UNet truyền thống bao gồm cả Self-Attention và Cross-Attention. MobileDiffusion đặt tất cả Self-Attention trong bản đồ đặc trưng có độ phân giải thấp nhất, nhưng vẫn giữ lại một Cross-Attention ở lớp trung gian. Đã được phát hiện rằng thiết kế này không chỉ cải thiện hiệu suất tính toán mà còn đảm bảo chất lượng tạo hình ảnh của mô hình được cải thiện.

Tinh chỉnh softmax thành relu

Softmax nổi tiếng khó để tối ưu hóa song song trong hầu hết các trường hợp chưa được tối ưu hóa và do đó không hiệu quả. MobileDiffusion đề xuất tinh chỉnh trực tiếp hàm softmax thành relu, vì relu là hoạt động của mỗi điểm, hiệu quả hơn. Ngạc nhiên thay, chỉ với khoảng 10.000 bước tinh chỉnh, chỉ số của mô hình đã được cải thiện và chất lượng của các hình ảnh được đảm bảo. Do đó, ưu điểm của relu so với softmax là rõ ràng.

Convolution phân tách (separable convolution).

Khóa cho việc giảm tham số của MobileDiffusion là việc sử dụng Convolution Separable. Công nghệ này đã được chứng minh là rất hiệu quả bởi công trình như MobileNet, đặc biệt là trên di động, nhưng nó thường ít được sử dụng trong các mô hình sinh. Các thí nghiệm của MobileDiffusion đã phát hiện ra rằng Convolution Separable rất hiệu quả trong việc giảm tham số, đặc biệt khi đặt nó ở lớp nội bộ nhất của UNet. Phân tích chứng minh rằng không có sự mất mát về chất lượng mô hình.

Tối ưu hóa mẫu

Các phương pháp tối ưu hóa mẫu phổ biến nhất hiện nay bao gồm Progressive Distillation và UFOGen, có thể đạt được 8 bước và 1 bước tương ứng. Để chứng minh rằng những mẫu này vẫn áp dụng được sau khi mô hình được đơn giản hóa đáng kể, MobileDiffusion đã tiến hành xác minh thực nghiệm trên cả hai.

So sánh giữa trước và sau tối ưu hóa mẫu và mô hình tham chiếu như sau. Có thể thấy các chỉ số của các mô hình 8 bước và 1 bước sau khi tối ưu hóa mẫu rất xuất sắc.

Thực nghiệm và Ứng dụng

Đánh giá di động

MobileDiffusion có thể đạt được tốc độ hiển thị hình ảnh nhanh nhất trên iPhone15Pro, chỉ 0,2 giây!

Kiểm tra nhiệm vụ phụ

MobileDiffusion khám phá các nhiệm vụ phụ bao gồm ControlNet/Plugin và LoRA Finetune. Như có thể thấy từ hình dưới đây, sau khi tối ưu hóa mô hình và mẫu, MobileDiffusion vẫn duy trì khả năng điều chỉnh mô hình xuất sắc.

Tóm tắt

MobileDiffusion khám phá nhiều mô hình và phương pháp tối ưu hóa mẫu, và cuối cùng có thể đạt được khả năng hiển thị hình ảnh trong vòng một giây trên thiết bị di động, và ứng dụng điều chỉnh phụ xuống vẫn được đảm bảo. Chúng tôi tin rằng điều này sẽ ảnh hưởng đến thiết kế mô hình diffusion hiệu quả trong tương lai và mở rộng các ví dụ ứng dụng di động.