Hiểu Mô Hình Khuếch Tán: Một Cái Nhìn Sâu Vào AI Sáng Tạo

Mô hình khuếch tán đã nổi lên như một phương pháp mạnh mẽ trong AI sáng tạo, tạo ra những kết quả đỉnh cao trong việc tạo hình ảnh, âm thanh và video. Trong bài viết kỹ thuật chuyên sâu này, chúng ta sẽ khám phá cách hoạt động của mô hình khuếch tán, những đổi mới chính của chúng và lý do tại sao chúng trở nên thành công đến vậy. Chúng ta sẽ đề cập đến nền tảng toán học, quá trình huấn luyện, thuật toán lấy mẫu và các ứng dụng tiên tiến của công nghệ mới thú vị này.

Giới Thiệu Về Mô Hình Khuếch Tán

Mô hình khuếch tán là một lớp mô hình sáng tạo học cách khử nhiễu dữ liệu dần dần bằng cách đảo ngược một quá trình khuếch tán. Ý tưởng cốt lõi là bắt đầu với nhiễu thuần túy và lặp đi lặp lại tinh chỉnh nó thành một mẫu chất lượng cao từ phân phối mục tiêu.

Cách tiếp cận này được lấy cảm hứng từ nhiệt động lực học không cân bằng – cụ thể là quá trình đảo ngược sự khuếch tán để khôi phục cấu trúc. Trong bối cảnh học máy, chúng ta có thể nghĩ về nó như việc học cách đảo ngược việc thêm nhiễu dần dần vào dữ liệu.

Một số lợi thế chính của mô hình khuếch tán bao gồm:

• Chất lượng hình ảnh đỉnh cao, vượt trội GANs trong nhiều trường hợp
• Huấn luyện ổn định mà không có động lực đối kháng
• Có khả năng song song hóa cao
• Kiến trúc linh hoạt – bất kỳ mô hình nào ánh xạ đầu vào đến đầu ra có cùng chiều đều có thể được sử dụng
• Nền tảng lý thuyết vững chắc

Hãy cùng đi sâu hơn vào cách hoạt động của mô hình khuếch tán.

Source: Song et al.

Phương trình vi phân ngẫu nhiên chi phối các quá trình thuận và nghịch trong mô hình khuếch tán. SDE thuận thêm nhiễu vào dữ liệu, dần dần biến đổi nó thành một phân phối nhiễu. SDE nghịch, được hướng dẫn bởi một hàm điểm số đã học, loại bỏ nhiễu một cách tiến triển, dẫn đến việc tạo ra hình ảnh thực tế từ nhiễu ngẫu nhiên. Cách tiếp cận này là chìa khóa để đạt được hiệu suất sáng tạo chất lượng cao trong các không gian trạng thái liên tục

Quá Trình Khuếch Tán Thuận

Quá trình khuếch tán thuận bắt đầu với một điểm dữ liệu x₀ được lấy mẫu từ phân phối dữ liệu thực, và dần dần thêm nhiễu Gaussian qua T bước thời gian để tạo ra các phiên bản nhiễu ngày càng tăng x₁, x₂, …, xT.

Tại mỗi bước thời gian t, chúng ta thêm một lượng nhỏ nhiễu theo:

x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε

Trong đó:

• β_t là một lịch trình phương sai kiểm soát lượng nhiễu được thêm vào ở mỗi bước
• ε là nhiễu Gaussian ngẫu nhiên

Quá trình này tiếp tục cho đến khi xT gần như là nhiễu Gaussian thuần túy.

Về mặt toán học, chúng ta có thể mô tả điều này như một chuỗi Markov:

q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)

Trong đó N biểu thị một phân phối Gaussian.

Lịch trình β_t thường được chọn là nhỏ cho các bước thời gian đầu và tăng dần theo thời gian. Các lựa chọn phổ biến bao gồm lịch trình tuyến tính, cosine hoặc sigmoid.

Quá Trình Khuếch Tán Nghịch

Mục tiêu của một mô hình khuếch tán là học đảo ngược của quá trình này – để bắt đầu với nhiễu thuần túy xT và dần dần khử nhiễu để khôi phục một mẫu sạch x₀.

Chúng ta mô hình hóa quá trình nghịch này như:

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))

Trong đó μ_θ và σ_θ^2 là các hàm đã học (thường là mạng nơ-ron) được tham số hóa bởi θ.

Đổi mới then chốt là chúng ta không cần phải mô hình hóa rõ ràng toàn bộ phân phối nghịch. Thay vào đó, chúng ta có thể tham số hóa nó theo quá trình thuận, mà chúng ta đã biết.

Cụ thể, chúng ta có thể chỉ ra rằng giá trị trung bình tối ưu của quá trình nghịch μ* là:

μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))

Trong đó:

• α_t = 1 – β_t
• ε_θ là một mạng dự đoán nhiễu đã học

Điều này cho chúng ta một mục tiêu đơn giản – huấn luyện một mạng nơ-ron ε_θ để dự đoán nhiễu đã được thêm vào ở mỗi bước.

Mục Tiêu Huấn Luyện

Mục tiêu huấn luyện cho mô hình khuếch tán có thể được suy ra từ suy luận biến phân. Sau một số đơn giản hóa, chúng ta đi đến một tổn thất L2 đơn giản:

L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]

Trong đó:

• t được lấy mẫu đồng đều từ 1 đến T
• x₀ được lấy mẫu từ dữ liệu huấn luyện
• ε là nhiễu Gaussian được lấy mẫu
• x_t được xây dựng bằng cách thêm nhiễu vào x₀ theo quá trình thuận

Nói cách khác, chúng ta đang huấn luyện mô hình để dự đoán nhiễu đã được thêm vào ở mỗi bước thời gian.

Kiến Trúc Mô Hình

Source: Ronneberger et al.

Kiến trúc U-Net là trung tâm của bước khử nhiễu trong mô hình khuếch tán. Nó có cấu trúc mã hóa-giải mã với các kết nối bỏ qua giúp bảo tồn các chi tiết tinh tế trong quá trình tái tạo. Bộ mã hóa giảm kích thước đầu vào hình ảnh một cách tiến triển trong khi nắm bắt các đặc trưng cấp cao, và bộ giải mã tăng kích thước các đặc trưng đã mã hóa để tái tạo hình ảnh. Kiến trúc này đặc biệt hiệu quả trong các nhiệm vụ yêu cầu định vị chính xác, chẳng hạn như phân đoạn hình ảnh.

Mạng dự đoán nhiễu ε_θ có thể sử dụng bất kỳ kiến trúc nào ánh xạ đầu vào đến đầu ra có cùng chiều. Các kiến trúc kiểu U-Net là một lựa chọn phổ biến, đặc biệt cho các nhiệm vụ tạo hình ảnh.

Một kiến trúc điển hình có thể trông như:

 class DiffusionUNet(nn.Module): def __init__(self): super().__init__() # Downsampling self.down1 = UNetBlock(3, 64) self.down2 = UNetBlock(64, 128) self.down3 = UNetBlock(128, 256) # Bottleneck self.bottleneck = UNetBlock(256, 512) # Upsampling self.up3 = UNetBlock(512, 256) self.up2 = UNetBlock(256, 128) self.up1 = UNetBlock(128, 64) # Output self.out = nn.Conv2d(64, 3, 1) def forward(self, x, t): # Embed timestep t_emb = self.time_embedding(t) # Downsample d1 = self.down1(x, t_emb) d2 = self.down2(d1, t_emb) d3 = self.down3(d2, t_emb) # Bottleneck bottleneck = self.bottleneck(d3, t_emb) # Upsample u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb) u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb) u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb) # Output return self.out(u1)