Mô hình hóa khám phá (Explorative Modeling): Bước tiến mới trong huấn luyện AI tạo sinh

Website: https://explorative-modeling.github.io/ GitHub: https://github.com/alexiglad/XM

TLDR: Chúng tôi giới thiệu Explorative Modeling, một mô hình mới cho generative modeling, đóng vai trò như một trục tiền huấn luyện (pretraining axis) thứ ba khi được thêm vào các mô hình tạo sinh hiện có, đồng thời cho phép tạo sinh end-to-end. Việc tăng cường khám phá (exploration) giúp cải thiện các mô hình hiện tại một cách đơn điệu trên các lĩnh vực hình ảnh, video và ngôn ngữ, với mức tăng trưởng tỉ lệ thuận theo quy mô (tăng từ 7% lên 36% với dữ liệu, 13% lên 23% với tham số). Cụ thể, các Explorative Models (XMs) đạt hiệu suất mẫu gấp 6,2 lần, hiệu suất FLOP gấp 4,1 lần và hiệu suất tham số tốt hơn 47%. Khám phá cũng cho phép mở rộng khả năng tổng quát hóa và mức độ end-to-end của các mô hình hiện có. Với tư cách là các mô hình tạo sinh end-to-end, XMs sánh ngang với diffusion trong các tác vụ điều khiển (control tasks) với lượng tính toán suy luận (inference compute) ít hơn tới 256 lần.

Hãy bắt đầu với một câu hỏi nghe có vẻ đơn giản. Nếu tôi yêu cầu một mô hình “tạo ra một con chó”, thì có bao nhiêu câu trả lời đúng?

Hóa ra là có rất nhiều… có lẽ hàng tỉ hoặc nhiều hơn thế nữa những hình ảnh mà chúng ta có thể coi là hình ảnh con chó.

Vậy điều gì sẽ xảy ra nếu chúng ta huấn luyện một mạng thần kinh để dự đoán trực tiếp hình ảnh con chó? Mô hình sẽ thấy hàng ngàn con chó hợp lệ khác nhau trong quá trình huấn luyện, và dự đoán duy nhất gần nhất với tất cả chúng chính là giá trị trung bình. Đó là những gì mô hình học được để xuất ra, và giá trị trung bình của hàng ngàn con chó thì chẳng giống một con chó nào cả, nó chỉ là một vệt màu nâu mờ nhạt.

Một con chó thực tế từ dữ liệu

Những gì mô hình dự đoán

Để cụ thể hóa điều này, hãy chơi một trò chơi. Tôi sẽ ném phi tiêu vào tấm bảng dưới đây, và mỗi mũi tên sẽ rơi vào một vị trí ngẫu nhiên trên các vòng tròn. Nhiệm vụ của bạn là đoán xem mũi tên tiếp theo của tôi sẽ rơi vào đâu, và bạn càng đoán sai, điểm số của bạn càng tệ.

Vậy bạn nên đoán ở đâu? Hóa ra dự đoán giúp giảm thiểu sai số của bạn chính là điểm chính giữa của tấm bảng. Chúng tôi đã huấn luyện một mô hình để chơi trò này, và quả nhiên, nó luôn đoán vào điểm chính giữa (đây là dự đoán tối ưu trong trường hợp này)!

Tuy nhiên, điều này thật tồi tệ… điểm chính giữa hầu như không bao giờ là nơi phi tiêu thực sự rơi vào. Dự đoán “tối ưu” lại là một điểm mà không có mũi tên nào từng rơi vào.

Đây là vấn đề cốt lõi của generative modeling. Khi một dự đoán có nhiều câu trả lời hợp lệ, dự đoán đơn lẻ tốt nhất chính là giá trị trung bình của chúng, và giá trị trung bình của dữ liệu thường là một câu trả lời tồi, chẳng giống gì với dữ liệu thực tế.

Và vấn đề này không chỉ xảy ra với bảng phi tiêu hay chó, nó xuất hiện với bất kỳ loại dữ liệu nào. Khi chúng tôi huấn luyện một mô hình để tạo trực tiếp ba cụm điểm 2D, nó dự đoán một dấu chấm duy nhất ở giữa chúng, và khi chúng tôi huấn luyện một mô hình trên văn bản, tất cả những gì nó có thể nói là “the”.

Dữ liệu thực tế

Những gì mô hình dự đoán

Dữ liệu thực tế

Những gì mô hình dự đoán

Nhưng khoan đã. ChatGPT viết văn bản mạch lạc, và các mô hình hình ảnh tạo ra những bức ảnh thực sự tuyệt vời. Rõ ràng vấn đề này đã được giải quyết bằng cách nào đó, phải không?

Đúng vậy, và mọi mô hình tạo sinh có khả năng mở rộng ngày nay đều giải quyết nó theo cùng một cách: chia quá trình tạo sinh thành nhiều bước nhỏ trong quá trình huấn luyện, để mỗi bước chỉ có khoảng một câu trả lời đúng. Khi một bước chỉ có một câu trả lời đúng, sẽ không còn gì để tính trung bình, và sự mờ nhạt sẽ biến mất.

Hãy xem cách thức này hoạt động. Các mô hình tự hồi quy (Autoregressive models - như LLMs) dự đoán từng phần một, điều này trong trò chơi của chúng ta có nghĩa là không bao giờ đoán vị trí chính xác của phi tiêu cùng một lúc. Thay vào đó, trước tiên bạn chỉ đoán xem phi tiêu rơi sang trái hay phải bao xa, và sau đó dựa vào đó, bạn đoán xem nó rơi lên trên hay xuống dưới bao xa. Một khi bạn biết phi tiêu đã rơi ở phía xa bên phải, thì chỉ còn vài vị trí mà nó có thể nằm ở đó.

Bước 1: chọn vị trí trái-phải

Bước 2: chọn vị trí trên-dưới, dựa trên vị trí trái-phải

Các mô hình Diffusion thực hiện điều này khác đi. Chúng bắt đầu từ nhiễu ngẫu nhiên thuần túy và thực hiện hàng trăm bước nhỏ hướng tới dữ liệu. Ban đầu, dự đoán của chúng vẫn có thể trở thành bất kỳ mũi tên nào, nhưng mỗi bước đều thu hẹp các khả năng, vì vậy không có bước đơn lẻ nào phải đối mặt với nhiều câu trả lời hợp lệ cùng một lúc.

Bắt đầu: có thể trở thành bất kỳ mũi tên nào

Giữa chừng: ít khu vực còn lại hơn

Gần cuối: đã được xác định vị trí

Hóa ra đây về cơ bản là cách mọi mô hình tạo sinh hiện đại hoạt động, bằng cách chia quá trình tạo sinh thành các phần nhỏ hơn có thể dự đoán tốt. Điều này bao gồm LLMs, các mô hình hình ảnh và video, và thậm chí cả các mô hình vài bước mới hơn như MeanFlow và consistency models. Chúng tôi gọi ý tưởng chia nhỏ quá trình tạo sinh này là factoring generation.

Cách tiếp cận factoring generation này hiệu quả, nhưng nó cũng là một "điều tồi tệ" vì một vài lý do. Thứ nhất là các mô hình được huấn luyện trên một bước đơn lẻ, nhưng lại chạy hàng trăm hoặc hàng ngàn bước khi suy luận (inference), vì vậy các đầu ra không hoàn hảo của chính chúng lại được đưa ngược lại làm đầu vào, các lỗi tích tụ, và các kết quả tạo sinh dần dần lệch xa khỏi bất cứ thứ gì mô hình từng thấy trong quá trình huấn luyện. Vấn đề này được gọi là exposure bias (tôi đã viết cả một blog về lý do tại sao nó lại tồi tệ), và đó là lý do tại sao các mô hình video tan chảy thành đống hỗn độn sau mười giây và tại sao LLMs trở nên kém mạch lạc hơn qua các đoạn tạo sinh rất dài, gây tổn hại trực tiếp đến hiệu suất và khả năng tổng quát hóa.

"Điều tồi tệ" thứ hai dựa trên điều thứ nhất, bởi vì sự không khớp giữa huấn luyện và suy luận có nghĩa là các mô hình này không bao giờ là end-to-end, trong khi một mô hình end-to-end sẽ chạy khi suy luận chính xác như cách nó được huấn luyện. Học end-to-end là thứ đã khởi đầu cuộc cách mạng deep learning với AlexNet, và bài học đó vẫn còn giá trị cho đến ngày nay… để các mô hình học mọi thứ trực tiếp từ dữ liệu sẽ tốt hơn là thiết kế thủ công các phần của pipeline, và một mô hình chạy theo cách nó được huấn luyện sẽ không bao giờ bị ép vào lãnh thổ out-of-distribution. Hầu hết deep learning hiện nay đã chuyển sang end-to-end ngoại trừ generative modeling, và factoring generation chính là thứ đang ngăn cản điều đó.

Vì vậy, lý tưởng nhất là chúng ta nên ngừng factoring generation, nhưng factoring cũng là thủ thuật duy nhất chúng ta biết để xử lý vấn đề "nhiều câu trả lời". Câu hỏi tự nhiên đặt ra là liệu chúng ta có thể factor thứ khác thay thế hay không, và hóa ra một mô hình tạo sinh chỉ có hai quá trình: cách nó tạo sinh và cách nó huấn luyện. Nếu việc tạo sinh không thể thay đổi, thì chỉ còn lại vòng lặp huấn luyện.

Vậy factoring training trông như thế nào? Để trả lời, hãy quay lại trò chơi của chúng ta, ngoại trừ lần này tôi sẽ cho bạn hai mươi lần đoán thay vì một, và chỉ lần đoán gần nhất của bạn mới được tính. Hóa ra với hai mươi lần đoán, việc đoán vào điểm chính giữa trở thành một chiến lược tồi tệ. Điều này là do bây giờ bạn có thể rải các dự đoán của mình lên những vị trí mà phi tiêu thực sự rơi vào, giảm sai số của bạn nhiều hơn so với việc chỉ đoán ở điểm chính giữa. Nói cách khác, chiến lược chiến thắng là sử dụng các dự đoán của bạn để khám phá các câu trả lời khác nhau.

Và đây chính xác là những gì xảy ra. Khi chúng tôi huấn luyện một mô hình theo cách này với hai mươi lần đoán (bảng giữa bên dưới), các dự đoán của nó lan rộng khắp tấm bảng!

2 lần đoán

20 lần đoán

200 lần đoán

Hãy dành một giây để trân trọng những gì vừa xảy ra ở đây. Các mũi tên rơi vào đúng những vị trí như trước, nhưng vì chúng ta đã thay đổi cách tính điểm cho các dự đoán, dự đoán tốt nhất có thể đã chuyển từ điểm chính giữa của tấm bảng sang các vị trí mà phi tiêu thực sự rơi vào. Điều này tiết lộ một điều quan trọng, đó là mục tiêu huấn luyện (training objective) đơn thuần kiểm soát dự đoán tốt nhất là gì (loss minimizer), và bằng cách thay đổi nó, chúng ta đã chuyển loss minimizer từ giá trị trung bình của dữ liệu sang chính dữ liệu đó.

Đây chính là Explorative Modeling. Tại mỗi bước huấn luyện, mô hình khám phá K khả năng khớp giữa những gì nó tạo ra và dữ liệu thực tế, và chỉ có khả năng khớp tốt nhất mới được huấn luyện. Chúng tôi gọi các mô hình được huấn luyện theo cách này là Explorative Models (XMs). Trong trường hợp đơn giản nhất, đây thực sự, một cách đẹp đẽ, là một vòng lặp for:

Đây là những gì xảy ra khi chúng ta tăng mức độ khám phá K trên dữ liệu thực tế:

Ground Truth (Sự thật gốc)

K=1 (không khám phá)