Imagen - Mô hình SOTA giải quyết bài toán Text-to-Image

Imagen - mô hình mới được công bố gần đây bởi Google với khả năng generate hình ảnh với đoạn text mô tả bất kỳ, cho dù ảnh đó không có thật hoặc phi vật lý. Phía trên là một ví dụ của ảnh được sinh ra bởi mô hình, rất ấn tượng phải không ?

1. Tổng quan về các mô hình Generative

4 họ model nổi bật trong tác vụ Generation:

• GAN: Huấn luyện dựa vào cơ chế học đối kháng với 2 thành phần: bộ sinh ảnh và bộ phân loại
• VAE: Tìm cách học một mô hình có khả năng biểu diễn phân phối của training data, sau đó sample biến ẩn từ phân phối học được và đưa qua decoder để sinh ảnh
• Flow-based model: Sử dụng một chuỗi các phép biến đổi khả nghịch (invertible transform) và sử dụng negative log-likelihood loss nhằm học phân phối dữ liệu một cách tường minh.
• Diffusion Model: Từ từ phá hủy cấu trúc của ảnh bằng cách thêm dần nhiễu, mô hình được huấn luyện để học cách phục hồi nhiễu về ảnh ban đầu.

Nhờ sử dụng Diffusion Model làm cốt lõi và tận dụng sức mạnh của mô hình pretrained language lớn, đội ngũ của Google Research Team đã phát triển và tạo ra Imagen với khả năng Generation ảnh hết sức mạnh mẽ. Nào, cùng tìm hiểu sâu hơn một chút về Diffusion model và điều gì đã khiến nó trở nên rất "hot" trong thời gian gần đây thôi. Let's go !

2. Diffusion Model - a new approach for Generation Task

2.1 Ý tưởng chung

Xuất phát từ paper năm 2015 bởi các sinh viên đại học Stanford, mô hình được lấy cảm hứng từ nguyên lý khuếch tán trong lĩnh vực nhiệt động lực học, các bạn có thể đọc thêm về paper tại đây. Ý tưởng chung là phá hủy một phân phối dữ liệu một cách từ từ và có hệ thống thông qua một chuỗi cách khuếch tán thuận, việc của chúng ta lúc này là tìm cách học một mô hình có thể đảo ngược quá trình khuếch tán đó nhờ đó có thể phục hồi lại cấu trúc dữ liệu ban đầu. Mô hình sau khi học được cách đảo ngược quá trình đó có thể được sử dụng cho tác vụ genetation một cách hiệu quả.

Ở đây khi áp dụng với hình ảnh, quá trình khuếch tán thuận sẽ từ từ làm nhiễu dần hình ảnh bằng cách thêm từ từ Gaussion noise vào ảnh trong từng step cho đến khi ảnh hoàn toàn chỉ là nhiễu mà không còn đặc trưng thông tin ban đầu, sau đó chúng ta tìm cách khử nhiễu ảnh dần dần theo từng step để thu được ảnh gốc ban đầu.

2.2 Nền tảng toán học đằng sau Diffusion model

• Khuếch tán thuận (Foward diffusion process)

  Một hình ảnh hay 1 datapoint $x _ { 0 } \sim q ( x )$ với $q(x)$ là phân phối của training data sẽ được thêm nhiễu (noise) lần lượt theo từng step. Ta định nghĩa một quá trình thuận $q \left ( x _ { t } | x _ { t - 1 } \right )$ với $t$ là tham số về mặt thời gian theo từng step, $t$ càng lớn tương ứng với bức ảnh càng nhiều nhiễu.

  Cụ thể quá trình thuận:

  $$q ( x _ { t } | x _ { t - 1 } ) = N ( x _ { t } ; \sqrt { 1 - \beta _ { t } } x _ { t - 1 } , \beta _ { t } I)$$

  Ở đây ảnh tại step t sẽ được thêm nhiễu Gaussion bằng cách sample ra $x_{t}$ theo phân phối Gauss với kỳ vọng $\sqrt { 1 - \beta _ { t } } x _ { t - 1 }$ và phương sai $\beta _ { t } I$. Lưu ý rằng $\beta$ ở đây là siêu tham số được thay đổi theo thời gian và kiểm soát bởi variance schedule : $\{ \beta _ { t } \in ( 0 , 1 ) \} _ { t = 1 } ^ { T }$. Khi số step T đủ lớn, ảnh ban đầu sẽ chuyển thành một ảnh nhiễu trắng hoàn toàn (isotropic noise).

• Khuếch tán ngược (Reverse diffusion process)

  Nếu chúng ta có thể đảo ngược quá trình bên trên thành $q \left ( x _ { t -1} | x _ { t } \right )$ thì từ một ảnh nhiễu trắng $x _ { T } \sim N ( 0 , I )$ chúng ta có thể tái tạo lại ảnh tương đương với mẫu được lấy trong phân phối $q(x)$ ban đầu. Nhưng thật không may rằng chúng ta không thể tính toán trực tiếp $q \left ( x _ { t -1} | x _ { t } \right )$ do đó chúng ta cần học một model $p _ { \theta }$ có khả năng xấp xỉ quá trình đảo ngược bên trên.

  Ta định nghĩa một quá trình ngược:

  $$p( x _ { t } | x _ { t - 1 } )=N ( x _ { t -1} ; \mu _ { \theta } ( x _ { t } , t ),\Sigma _ { \theta } ( x _ { t } , t ) )$$

  trong đó $\mu _ { \theta }$ và $\Sigma _ { \theta }$ chính là hai tham số mà mô hình của chúng ta cần xấp xỉ. Toàn bộ quá trình biến đổi ngược từ $x _ { T }$ về $x _ { 0 }$ có thể được tính toán nhờ vào tính chất độc lập của chuỗi Markov:

  $$p _ { \theta } ( x _ { 0:T } ) = p ( x _ { T } ) \prod _ { t = 1 } ^ { T } p _ { \theta } ( x _ { t - 1 } | x _ { t } )$$

• Quá trình training

  Mục tiêu của ta lúc này là minimize hàm negative log-likelihood: $- \log p _ { \theta } ( x _ { 0 } )$, tuy nhiên hàm này cũng không trực tiếp tính toán được, do đó chúng ta cần sử dụng một kỹ thuật tương tự đã áp dụng với VAE đó là Variational lower bound:

  $$- \log p _ { \theta } ( x _ { 0 } ) \leq - \log p _ { \theta } ( x _ { 0 } ) + D _ { K L } ( q ( x _ { 1:T } | x _ { 0 } ) \| p _ { \theta } ( x _ { 1:T } | x _ { 0 } ) )$$

  Chúng ta sẽ gián tiếp minimize hàm $- \log p _ { \theta } ( x _ { 0 } )$ bằng cách minimize đẳng thức phía bên phải.Tiếp tục sử dụng tính chất Bayes và bất đẳng thức Jensen chúng ta có được hàm mục tiêu cần minimize (phần biến đổi này khá dài và phức tạp nên mình không trình bày cụ thể, mình recommend các bạn nên đọc thêm ở link bên dưới để hiểu rõ):

  $$D _ { K L } ( q ( x _ { T } | x _ { 0 } ) \| p _ { \theta } ( x _ { T } ) ) +\sum_{t=2} ^ { T } D _{ K L } ( q ( x _{ t - 1 } | x _{ t } , x _{ 0 } ) \| p_{ \theta } ( x _{ t - 1 } | x _{ t } ) ) - \log p _{ \theta } ( x _{ 0 } | x _{ 1 } )$$

  Thành phần KL Divergence đầu tiên có thể tính toán trực tiếp do $p _ { \theta } ( x _ { T }), x_{T}, x_{0}$ và quá trình $q$ chúng ta đã biết, lúc này chúng ta chỉ cần minimize:

  $$\sum_{t=2} ^ { T } D_{ K L } ( q ( x_{ t - 1 } | x_{ t } , x_{ 0 } ) \| p_{ \theta } (x_{ t - 1 } | x_{ t } ) ) - \log p_{ \theta } ( x_{ 0 } | x_{ 1 } )$$

  Thành phần KL Divergence đầu tiên so sánh 2 phân phối Gauss $q ( x_{ t - 1 } | x_{ t } , x_{ 0 } )$ và $p_{ \theta } (x_{ t - 1 } | x_{t})$ do đó nó nó có thể tính toán ở dạng closed form. Ở trên ta có nói rằng $q \left ( x _ { t -1} | x _ { t } \right )$ không thể trực tiếp tính toán nhưng thật may là $q ( x_{ t - 1 } | x_{ t } , x_{ 0 } )$ thì lại có thể tính toán được ^^. Cụ thể :

  $$q ( x _ { t - 1 } | x _ { t } , x _ { 0 } ) = N ( x _ { t - 1 } ; \tilde { \mu } ( x _ { t } , x _ { 0 } ) , \tilde { \beta } _ { t } \Sigma \rangle$$

  với $\tilde { \mu }$ và $\tilde { \beta }$ có thể tính toán qua Bayes rule:

  $$\tilde { \beta } = \frac { 1 - \tilde { \alpha } _ { t - 1 } } { 1 - \tilde { \alpha } _ { t } } \cdot \beta _ { t } ; \tilde { \mu } = \frac { 1 } { \sqrt { \alpha _ { t } } } ( x _ { t } - \frac { 1 - \alpha _ { t } } { \sqrt { 1 - \tilde { \alpha } _ { t } } } \epsilon _ { t } )$$

  Cần nhắc lại rằng, chúng ta đang tìm cách huấn luyện mô hình có thể xấp xỉ được xác suất có điều kiện của quá trình biến đổi ngược $p _ { \theta } ( x _ { t - 1 } | x _ { t } ) = N ( x _ { t - 1 } ; \mu _ { \theta } ( x _ { t } , t ) , \Sigma _ { \theta } ( x _ { t } , t ) )$ . Chúng ta mong muốn $\mu _ { \theta }$ có thể dự đoán được $\tilde { \mu } _ { t }$ , một phương pháp đơn giản được sử dụng để tính toán lỗi đó là MSE loss:

  $$L _ { t } = { \mathrm { E } } _ { x _ { 0 } , \epsilon } [ \frac { 1 } { 2 \| \Sigma _ { \theta } ( x _ { t } , t ) \| ^ { 2 } } \| \tilde { \mu } _ { t } ( x _ { t } , x _ { 0 } ) - \mu _ { \theta } ( x _ { t } , t ) \| ^ { 2 } ]$$

  Và vì $\tilde { \mu } _ { t }$ và $\mu _ { \theta }$ có thể tính toán thông qua noise $\epsilon _ { t }, \epsilon_{\theta}$ và input $x_{t}$ theo từng step, cuối cùng mô hình chỉ cần predict ra noise đã được thêm vào:

  $$L _ { t } ^ { \operatorname { s i m p l e } } = E _ { t \sim [ l , T ] , x _ { 0 } , \epsilon _ { t } } [ \| \epsilon _ { t } - \epsilon _ { \theta } ( x _ { t } , t ) \| ^ { 2 } ]$$

  Chi tiết phần biến đổi toán học các bạn có thể tham khảo tại đây.

  Khi đã xác định được noise đã được thêm vào trong quá trình thuận, chúng ta có thể thu được ảnh bớt nhiễu bằng cách lấy ảnh nhiễu "trừ" đi noise.

• Kiến trúc của mạng

  Chúng ta cần xác định kiến trúc của model sử dụng trong quá trình ngược. Có thể thấy, quá trình ngược ở trên yêu cầu ảnh input và ảnh ouput phải có cùng kích cỡ đồng thời chúng ta cũng phải thiết kế sao cho model có thể nắm được semantic của ảnh đầu vào ở mức độ pixel nhằm xác định noise. Tác giả của paper DDPM đã quyết định sử dụng mạng U-Net do nó có thể đáp ứng đầy đủ các yêu cầu trên. 2223

  Ảnh: Cấu trúc của mạng U-Net được sử dụng

3. Text-to-Image Generative với Diffusion Model

Diffusion model đã được huấn luyện có thể thực hiện sinh ảnh mới tương tự với ảnh trong phân phối dùng để training từ một nhiễu trắng bất kì, tuy nhiên ảnh được sinh ra hoàn toàn ngẫu nhiên mà ta không thể kiểm soát được các đặc điểm cụ thể, hơn nữa ảnh được sinh ra thường có độ phân giải thấp. Để có thể áp dụng diffusion model cho tác vụ Text-to-image chúng ta cần bổ sung một vài cải tiến.

• Classifier Guidance

  Là phương pháp làm tăng tính trung thực cũng như chất lượng của hình ảnh với trade-off là giảm sự đa dạng của ảnh được sinh ra. Nhằm bổ sung thêm thông tin về class của ảnh trong quá trình khuyếch tán , trong paper các tác giả đến từ OpenAI có đề xuất bổ sung thêm gradient $c$ của một mô hình classifier đã được training riêng biệt qua đó hướng quá trình sinh ảnh tập trung vào class được target.

• Classifier-Free Guidance

  Với phương pháp Classifier Guidance bên trên, gradient được lấy từ một pre-trained model riêng biệt . Để tránh việc phụ thuộc vào thông tin từ một pre-trained model khác, phương pháp này huấn luyện một cách đồng thời một mô hình diffusion theo cả 2 mục tiêu là unconditional và conditional bằng cách ngẫu nhiên loại bỏ gradient $c$ trong quá trình training. Cụ thể, predicted noise được định nghĩa:

  $$\tilde { \epsilon } _ { \theta } ( z _ { t } , c ) = w \epsilon _ { \theta } ( z _ { t } , c ) + ( 1 - w ) \epsilon _ { \theta } ( z _ { t } )$$

  Trong đó $w$ được gọi là guidance weight, $w$ càng lớn tương ứng với việc tăng sức ảnh hưởng của gradient $c$.

  Note: Imagen đặc biệt chú trọng vào classifier-free guidance để generate ảnh có sự tương quan nhất với text.

• Cascaded diffusion models

  Ảnh được sinh ra bởi mô hình diffusion thường sẽ cho độ phân giải thấp (VD: 32x32), lý do là bởi độ phân giải thấp thì mô hình sẽ tập trung vào việc generate ảnh với các chi tiết và đặc điểm một cách hợp lý nhất, điều này là khó khăn khi gen luôn ra ảnh với độ phân giải cao. Do đó paper này đã đề xuất một cấu trúc nhằm tăng dần độ phân giải của ảnh đươc sinh ra bằng cách sử dụng mô hình thác nước với 3 diffusion model nối tiếp nhau, trong đó 2 model phía sau được sử dụng để tăng độ phân giải.

  Ảnh: Cấu trúc của cascaded diffusion model
  Ảnh sau khi được sinh ra bởi mô hình đầu tiên sẽ làm input cho model diffusion thứ 2, tương tự đổi với diffusion model cuối cùng. Ngoài ra để việc scale up ảnh đạt hiệu quả tốt nhất, các additional signal sử dụng trong conditional diffusion cũng được bổ sung vào input của từng model.

  Các bạn có thể tham khảo thêm về cascaded diffusion models tại đây.

• Text-to-Image Model

  Sử dụng Text Encoder nhằm embedding text input thành vectơ và sử dụng vectơ này làm input cho mô hình Cascaded Diffusion, quá trình generate sẽ sử dụng thông tin từ vectơ text embedding để ảnh được sinh ra giống với mô tả của text.

  Ảnh: Flow của mô hình sinh ảnh Imagen

4. Imagen

Sau khi đi qua một loạt các khái niệm toán học và giải thích giờ hãy bàn luận về nhân vật chính của chúng ta. Điểm khác biệt lớn khiến cho Imagen trở thành mô hình text-to-image mạnh nhất có thể kể đến đó là:

• Text Encoder sử dụng các mô hình pre-trained language lớn

DALL-E 2, GLIDE hay các mô hình text-to-image trước đó sử dụng Text Encoder được training với bimomal data - tức là data đi theo cặp (image, text caption), ưu điểm là dạng data này giúp mô hình có thể học được sự liên kết giữa ảnh và text tương ứng với nó, tuy nhiên lượng data là quá ít khi so sánh với lượng data thuần text hoặc image . Text Encoder được sử dụng trong Imagen là mô hình Encoder T5-XXL , T5 là pre-trained transformer language models và được sử dụng cho rất nhiều tác vụ NLP khác nhau. Mô hình T5 được train với lượng dữ liệu chỉ bao gồm text với lượng data là rất lớn khiến cho Encoder T5 có thể nắm bắt được ngữ nghĩa ẩn của câu text rất hiệu quả.

• Threshoding - khắc phục điểm yếu của classifier-free guidance

Việc tăng guidance weight $w$ giúp mô hình có thể cải độ phù hợp giữa ảnh và text (image-text alignment), tuy nhiên khi $w$ quá lớn (VD: $w>5$ ) ảnh được generate sẽ thường trông không tự nhiên hay tệ hơn là ảnh bị bão hòa. Lí do là bởi khi $w$ lớn, ảnh được predict tại từng step có giá trị pixel vượt quá khoảng giá trị pixel của các ảnh trong training data (VD: khoảng [-1,1]). Để giải quyết vấn đề này, tác giả Imagen đã nghiên cứu 2 giải pháp:

1. Static Thresholding: Thực hiện clipped giá trị các pixel về khoảng [-1,1], cụ thể tất cả các giá trị pixel > 1 sẽ trở thành 1, các giá trị pixel < -1 sẽ trở thành -1.
2. Dynamic Thresholding: Chọn một số lượng các pixel có giá trị dương (VD: 80% lượng pixel > 0) và threshhold trên giá trị $s$ của các pixel đó. Khi giá trị $s$ đó >1 thì clipped giá trị tất cả các pixel về khoảng $[-s,s]$, sau đó chia tất cả cho $s$ để scale về khoảng [-1,1].

• Efficient U-Net

  Với cascade diffusion model được trình bày ở trên, sau khi đưa qua 2 mạng supper-resolution ảnh thu được có độ phân giải 256x256 nhưng như vậy vẫn chưa đủ. Ảnh được sinh ra bởi Imagen có độ phân giải 1024x1024 và điều này khiến cho mạng U-net trở nên rất nặng với nhiều tham số. Do đó các tác giả đã cải tiến U-Net để có thể tiết kiệm bộ nhớ, cải thiện thời gian training và inference, đó là Efficien U-Net. Một số phương pháp họ sử dụng: đảo thứ tự của các Downsampling (và Upsampling) với các lớp convolution nhằm giảm số lượng tham số, thêm các residual block khi ảnh ở độ phân giải thấp, scale các skip-connection bởi $1 / \sqrt { 2 }$,...

5. Kết luận

Có thể nói rằng generative model nói chung và bài toán text-to-image nói riêng đang là lĩnh vực đang được các ông lớn công nghệ cũng như giới khoa học đặc biệt quan tâm. Các ứng dụng AI vẽ tranh, Deepfake,... sử dụng các mô hình generative đều thu hút được sự tò mò và phấn khích của người dùng. Mô hình mới đây của Google đã đặt một bước tiến mới trong việc cải thiện chất lượng của các mô hình sinh ảnh đồng thời mở ra nhiều hướng đi mới để tiếp tục cải thiện các mô hình này.