Hướng dẫn xây dựng mô hình ngôn ngữ khuếch tán (Diffusion Language Models)

Mục lục

Giới thiệu: Mô hình ngôn ngữ tự hồi quy và khuếch tán

Hai họ thuật toán AI tạo sinh đang được sử dụng rộng rãi hiện nay. Đối với dữ liệu liên tục như hình ảnh hoặc video, phương pháp tiên tiến nhất dựa trên các mô hình khuếch tán (diffusion models). Đối với dữ liệu rời rạc như văn bản hoặc mã nguồn, phương pháp tiêu chuẩn lại là các mô hình tự hồi quy (autoregressive models). Bài viết này khám phá một giải pháp thay thế cho dữ liệu rời rạc, được xây dựng dựa trên mô hình khuếch tán hiện đại.

Các mô hình ngôn ngữ phổ biến hiện nay đều là tự hồi quy: chúng tạo ra các token từ trái sang phải, từng cái một, mỗi token phụ thuộc vào các token đứng trước nó. Cách tiếp cận này rất mạnh mẽ nhưng cũng có những hạn chế cố hữu:

Các mô hình khuếch tán áp dụng một cách tiếp cận khác. Thay vì tạo văn bản từng token một, chúng tạo ra toàn bộ chuỗi cùng lúc, bắt đầu từ một dự đoán ban đầu và tinh chỉnh dần qua nhiều bước. Điều này mang lại một số ưu điểm: quá trình tạo có thể cân bằng giữa tốc độ và chất lượng bằng cách điều chỉnh số bước, các lỗi có thể được sửa chữa trong quá trình thực hiện, và mỗi bước đều chú trọng đến ngữ cảnh hai chiều (bidirectional context).

Việc áp dụng khuếch tán vào ngôn ngữ từ lâu đã là một bài toán mở. Năm 2024, lĩnh vực này đã đạt đến bước ngoặt khi các mô hình khuếch tán trở nên cạnh tranh với các mô hình tự hồi quy về mặt chất lượng. Đến năm 2026, các LLM khuếch tán đã trở thành hiện thực với sự ra mắt từ các phòng thí nghiệm hàng đầu trong ngành — Mercury 2 (Inception Labs), Gemma Diffusion (Google) và Nemotron Diffusion (NVIDIA). Bài viết này sẽ truy xuất các ý tưởng và công trình nghiên cứu nền tảng cho những mô hình hiện đại này.

Bối cảnh: Khuếch tán Gaussian

Trước khi giới thiệu về khuếch tán cho ngôn ngữ, chúng ta sẽ bắt đầu với cái nhìn tổng quan ngắn gọn về khuếch tán Gaussian trong tạo ảnh. Sau đó, chúng ta sẽ xây dựng mô hình khuếch tán rời rạc dựa trên sự tương đồng này.

Tạo dữ liệu bằng cách khử nhiễu lặp lại

Khái niệm trung tâm của các mô hình khuếch tán là khử nhiễu (denoising). Thay vì vẽ một bức ảnh trong một lần duy nhất, mô hình khuếch tán tạo ra hình ảnh từng bước một, bắt đầu từ nhiễu ngẫu nhiên thuần túy và loại bỏ dần nhiễu ở mỗi bước cho đến khi một hình ảnh rõ nét xuất hiện. Việc tạo ảnh qua nhiều bước nhỏ hóa ra đơn giản hơn nhiều so với việc tạo ra tất cả cùng một lúc, và đây chính là điều làm cho các mô hình khuếch tán trở nên hiệu quả.

Mô hình học cách khử nhiễu như thế nào? Bí quyết nằm ở việc dạy cho nó bằng cách hiển thị các ví dụ về nhiễu được chuyển đổi dần dần thành hình ảnh. Khuếch tán đạt được điều này thông qua hai quá trình bổ trợ. Thứ nhất, quá trình thuận (forward process) lấy một hình ảnh gốc sạch và biến nó thành nhiễu thuần túy, từng bước một. Thứ hai, quá trình ngược (reverse process) học cách đảo ngược quá trình chuyển đổi này, biến nhiễu thuần túy trở lại thành hình ảnh; nó được huấn luyện dựa trên các quỹ đạo từ ảnh sang nhiễu do quá trình thuận tạo ra.

Quá trình thuận

Quá trình thuận lấy một hình ảnh huấn luyện sạch và tạo ra một chuỗi các hình ảnh ngày càng nhiễu, vạch ra con đường từ dữ liệu sạch đến nhiễu thuần túy. Nó thực hiện điều này bằng cách trộn thêm một lượng nhiễu Gaussian ngẫu nhiên ngày càng tăng ở mỗi bước, cho đến khi hình ảnh tan biến thành nhiễu tĩnh thuần túy. Bước này không đòi hỏi bất kỳ sự học hỏi nào — chúng ta chỉ đơn giản là thêm nhiễu — nhưng nó cực kỳ hữu ích vì nó tạo ra nguồn dữ liệu huấn luyện vô tận: các ví dụ về hình ảnh được chuyển đổi thành nhiễu và ngược lại.

Quá trình ngược

Quá trình ngược là nơi diễn ra việc học thực sự. Chúng ta huấn luyện một mô hình để biến nhiễu thành hình ảnh bằng cách thực hiện ngược lại các bước do quá trình thuận tạo ra.

Cụ thể, với một hình ảnh nhiễu, chúng ta huấn luyện một mô hình học máy để tách nhiễu khỏi hình ảnh gốc hoặc, tương đương, dự đoán nhiễu đã được thêm vào hoặc chính hình ảnh sạch đó, vì khi đã có đầu vào nhiễu, việc biết một trong hai sẽ xác định cái còn lại. Một khi mô hình có thể làm được điều này, việc tạo dữ liệu trở nên đơn giản: bắt đầu từ nhiễu thuần túy, yêu cầu mô hình ước tính và loại bỏ một chút nhiễu, rồi lặp lại. Mỗi lượt thực hiện sẽ đẩy mẫu đến gần hơn với dữ liệu thực, cho đến khi chỉ còn lại hình ảnh sạch.

Công thức thuận/ngược này — làm hỏng dữ liệu bằng nhiễu, sau đó học cách đảo ngược sự hư hỏng từng bước một — chính là bản thiết kế cho mọi mô hình khuếch tán.

Các mô hình khuếch tán có mặt nạ đơn giản (Simple Masked Diffusion Models)

Trở ngại chính khi đưa khuếch tán vào ngôn ngữ là quyết định "nhiễu" nên có ý nghĩa gì đối với các token rời rạc. Ví dụ, nhiễu được sử dụng trong khuếch tán cổ điển là Gaussian, và việc thêm nhiễu Gaussian liên tục vào các biến phân loại là không được xác định rõ ràng. Dưới đây, chúng ta giới thiệu một phương pháp đơn giản nhưng hiệu quả, định nghĩa nhiễu thông qua việc che mặt nạ (masking). Nhóm của chúng tôi đã phổ biến phương pháp này, và hiện nay nó tạo thành nền tảng cho hầu hết các mô hình ngôn ngữ khuếch tán mã nguồn mở.

Tóm tắt về khuếch tán có mặt nạ

Cách dễ nhất để hiểu khuếch tán có mặt nạ là coi nó như một transformer giải mã (unmasking transformer). Chúng ta huấn luyện mô hình bằng cách lấy các chuỗi sạch, che đi một phần ngẫu nhiên các token của chúng và yêu cầu một transformer hai chiều điền vào các chỗ trống. Nếu bạn biết BERT, thì đây về cơ bản là BERT với tỷ lệ che mặt nạ ngẫu nhiên — nhưng không giống như BERT, mô hình thu được có khả năng tạo sinh. Bạn có thể coi khuếch tán có mặt nạ là một BERT tạo sinh.

Sau khi huấn luyện transformer giải mã, chúng ta có thể tạo văn bản bằng cách bắt đầu từ một chuỗi đã bị che hoàn toàn và lặp lại hai bước nhiều lần:

Mỗi vòng lặp sẽ để lại ít vị trí bị che hơn, cho đến khi chuỗi hội tụ về một mẫu sạch từ mô hình. Do đó, việc tạo dữ liệu tương đương với việc bắt đầu từ một chuỗi đầy chỗ trống và dần dần điền các từ vào theo một thứ tự tùy ý.

Hiểu về khuếch tán có mặt nạ: Góc nhìn xác suất

Chúng ta cũng có thể hiểu rõ hơn tại sao quá trình này hoạt động bằng cách coi nó như một mô hình tương tự với mô hình khuếch tán Gaussian mà chúng ta đã thấy trước đó. Giống như khuếch tán Gaussian, khuếch tán có mặt nạ có thể được mô tả như một mô hình bao gồm quá trình thuận và quá trình ngược.

Quá trình thuận

Mục tiêu của quá trình thuận là tạo dữ liệu huấn luyện cho quá trình ngược. Đầu ra của nó là một quỹ đạo bắt đầu từ một điểm dữ liệu và kết thúc tại một chuỗi nhiễu thuần túy; quá trình ngược sau đó sẽ được huấn luyện để tạo ra quỹ đạo này theo chiều ngược lại.

Thách thức chính là quyết định "nhiễu" nên có ý nghĩa gì. Trong khuếch tán Gaussian, chúng ta thêm các lượng nhiễu trắng khác nhau vào hình ảnh. Trong khuếch tán có mặt nạ, thay vào đó, chúng ta che ngẫu nhiên một phần các token trong một chuỗi rời rạc. Mức độ che mặt nạ được điều khiển bởi một lịch trình $\alpha_t$ — xác suất mà một token nhất định vẫn không bị che — đóng vai trò như tỷ lệ tín hiệu trên nhiễu (signal-to-noise ratio) trong khuếch tán Gaussian. Nó bắt đầu ở mức $1$ khi $t = 0$ (chuỗi sạch) và giảm xuống $0$ khi $t = 1$ (chuỗi bị che hoàn toàn). Biến thời gian $t$ đánh chỉ số cho con đường từ dữ liệu sạch đến dữ liệu nhiễu, và tại thời điểm $t$, một chuỗi bị che một phần $z_t$ có kỳ vọng là một phần $\alpha_t$ các token không bị che.

Chúng ta thực hiện quá trình này như một chuỗi Markov trên một chuỗi các biến $z_t$ được đánh chỉ số bởi $t$, với $z_0$ là chuỗi sạch, không bị che. Với $s < t$, chuỗi xác định $q(z_t \mid z_s)$ bằng cách che mỗi token chưa bị che của $z_s$ với xác suất $(\alpha_s - \alpha_t)/\alpha_s$. Chạy chuỗi Markov này qua một số bước sẽ tạo ra một quỹ đạo đi từ dữ liệu sạch đến nhiễu bị che hoàn toàn.

Quá trình ngược

Tiếp theo, giống như trong khuếch tán Gaussian, chúng ta huấn luyện quá trình ngược để đi ngược lại chuỗi các biến tiềm ẩn bị che ngày càng nhiều — bắt đầu từ một chuỗi bị che hoàn toàn và cuối cùng tạo ra các đầu ra tương tự như dữ liệu sạch.

Sử dụng quy tắc Bayes, chúng ta có thể suy ra quá trình ngược tối ưu về mặt toán học $q(z_s \mid z_t, x)$ khi biết chuỗi sạch $x$. Quá trình tối ưu này có hai bước: (1) với $z_t$ bị che một phần, chúng ta nhìn vào $x$ để tìm các token sạch thực sự; (2) tạo $z_s$ bằng cách thay thế mỗi vị trí bị che của $z_t$ bằng giá trị của nó trong $x$ với xác suất $(\alpha_s - \alpha_t) / (1 - \alpha_t)$, và giữ nguyên các vị trí còn lại là bị che.

Trong thực tế, đầu ra cuối cùng $x$ rõ ràng là chưa biết khi chúng ta tạo nó. Do đó, chúng ta huấn luyện một mô hình $x_\theta(z_t)$ để dự đoán chuỗi sạch cuối cùng dựa trên trạng thái hiện tại $z_t$ và áp dụng quá trình ngược lý tưởng $q(z_s \mid z_t, x)$ bằng cách sử dụng ước tính $x_\theta(z_t)$ thay cho $x$ thực. Chính thức hơn, chúng ta định nghĩa quá trình ngược là một xác suất $p(z_s \mid z_t) = q\big(z_s \mid z_t, x_\theta(z_t)\big)$. Định nghĩa này khôi phục thuật toán lấy mẫu mà chúng ta đã mô tả trước đó: tại mỗi bước, chúng ta sử dụng mô hình $x_\theta(z_t)$ để điền vào các chỗ trống của $z_t$, và chúng ta giữ lại một tập con các token đã điền này trong $z_s$.

Một mô hình xác suất cho khuếch tán có mặt nạ

Kết hợp các phần này lại với nhau, chúng ta có định nghĩa toán học của một mô hình ngôn ngữ khuếch tán có mặt nạ (MDLM). Quá trình thuận $q(z_t \mid z_s)$ tạo ra một quỹ đạo từ dữ liệu sạch đến dữ liệu bị che hoàn toàn, và quá trình ngược $p(z_s \mid z_t)$ học cách đảo ngược nó. Hơn nữa, quá trình ngược định nghĩa một mô hình biến tiềm ẩn $p(x, z_1, \dots, z_T)$ trong đó $T$ mẫu bị che một phần trung gian $z_1,...,z_T$ là các biến tiềm ẩn. Việc tạo dữ liệu từ quá trình ngược $p(z_s \mid z_t)$ cũng giống như thực hiện lấy mẫu tổ tiên (ancestral sampling) từ mô hình này.

Chúng ta cũng có thể xem xét hàm khả năng (likelihood) $\log p(x)$ của mô hình $p$ để đánh giá chất lượng của nó. Trong các mô hình biến tiềm ẩn, điều này là không thể tính toán trực tiếp, vì vậy chúng ta phải dùng đến các phương pháp xấp xỉ thông qua suy luận biến phân (variational inference). Đối với mô hình ngôn ngữ khuếch tán có mặt nạ, cận dưới bằng chứng (ELBO) được sử dụng để xấp xỉ hàm khả năng có một dạng đơn giản đáng ngạc nhiên (giả sử đơn giản hóa $\alpha_t = 1-t$):

Hãy phân tích công thức này. Thuật ngữ bên trong $\log p_\theta(x \mid z_t)$ là hàm khả năng của một chuỗi sạch $x$ dựa trên một chuỗi bị che một phần $z_t$ được lấy mẫu từ quá trình thuận. Nói cách khác, đó là hàm mất mát cross-entropy giữa các dự đoán của transformer giải mã của chúng ta và các token thực — đây chính xác là hàm mất mát của BERT!

Khác với BERT, hàm mất mát này được tính trung bình trên tất cả $t$, và do đó trên tất cả các tỷ lệ che mặt nạ có thể có, thay vì một tỷ lệ cố định duy nhất. Nó cũng được chuẩn hóa bởi $t$, phần kỳ vọng của các token bị che (vì $\alpha_t = 1-t$); hệ số này đảm bảo rằng mỗi hàm mất mát BERT được chuẩn hóa theo số lượng token mà hàm mất mát đó được tính toán.

Tóm tắt và đánh giá khuếch tán có mặt nạ

Tóm lại, MDLM rất giống với BERT, với hai khác biệt chính: