KTO Trainer

Tổng quan

TRL hỗ trợ bộ huấn luyện Tối ưu hóa Kahneman-Tversky (KTO) để huấn luyện các mô hình ngôn ngữ, như được mô tả trong bài báo KTO: Model Alignment as Prospect Theoretic Optimization của Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, Douwe Kiela.

Tóm tắt từ bài báo như sau:

Lý thuyết triển vọng của Kahneman & Tversky cho thấy con người nhận thức các biến ngẫu nhiên theo một cách có thiên kiến nhưng được xác định rõ ràng; chẳng hạn, con người nổi tiếng là tránh thua lỗ. Chúng tôi chỉ ra rằng các hàm mục tiêu để căn chỉnh LLM với phản hồi của con người ngầm tích hợp nhiều thiên kiến này — thành công của các hàm mục tiêu này (ví dụ: DPO) so với việc tối thiểu hóa entropy chéo có thể một phần được quy cho việc chúng là các hàm mất mát nhận thức con người (HALOs). Tuy nhiên, các hàm tiện ích mà các phương pháp này gán cho con người vẫn khác so với những hàm trong tài liệu về lý thuyết triển vọng. Sử dụng mô hình tiện ích con người của Kahneman-Tversky, chúng tôi đề xuất một HALO trực tiếp tối đa hóa tiện ích của các bản sinh thay vì tối đa hóa log-likelihood của các sở thích, như các phương pháp hiện tại đang làm. Chúng tôi gọi cách tiếp cận này là Tối ưu hóa Kahneman-Tversky (KTO), và nó đạt hoặc vượt hiệu suất của các phương pháp dựa trên sở thích ở quy mô từ 1B đến 30B. Quan trọng là, KTO không cần dữ liệu sở thích — chỉ cần một tín hiệu nhị phân về việc một đầu ra có mong muốn hay không mong muốn cho một đầu vào nhất định. Điều này khiến nó dễ sử dụng hơn nhiều trong thực tế, nơi dữ liệu sở thích khan hiếm và đắt đỏ.

Mã nguồn chính thức có thể được tìm thấy tại ContextualAI/HALOs.

Phương pháp hậu huấn luyện này được đóng góp bởi Kashif Rasul, Younes Belkada, Lewis Tunstall, Pablo Vicente, và sau đó được tái cấu trúc bởi Albert Villanova del Moral.

Bắt đầu nhanh

Ví dụ này minh họa cách huấn luyện một mô hình bằng phương pháp KTO. Chúng ta sử dụng mô hình Qwen 0.5B làm mô hình cơ sở. Chúng ta sử dụng dữ liệu ưu tiên từ KTO Mix 14k. Bạn có thể xem dữ liệu trong bộ dữ liệu tại đây:

Dưới đây là script để huấn luyện mô hình:

# train_kto.py
from datasets import load_dataset
from trl import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

training_args = KTOConfig(output_dir="Qwen2-0.5B-KTO")
trainer = KTOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

Chạy script bằng lệnh sau:

accelerate launch train_kto.py

Được phân phối trên 8 x GPU H100, quá trình huấn luyện mất khoảng 30 phút. Bạn có thể kiểm tra tiến độ huấn luyện bằng cách xem biểu đồ phần thưởng. Xu hướng tăng của biên độ phần thưởng cho thấy mô hình đang cải thiện và tạo ra các phản hồi tốt hơn theo thời gian.

Để xem mô hình đã huấn luyện hoạt động như thế nào, bạn có thể sử dụng Transformers Chat CLI.

$ transformers chat trl-lib/Qwen2-0.5B-KTO
<quentin_gallouedec>:
What is the best programming language?

<trl-lib/Qwen2-0.5B-KTO>:
The best programming language can vary depending on individual preferences, industry-specific requirements, technical skills, and familiarity with the specific use case or task. Here are some widely-used programming languages that have been noted as popular and widely used:

Here are some other factors to consider when choosing a programming language for a project:

 1 JavaScript: JavaScript is at the heart of the web and can be used for building web applications, APIs, and interactive front-end applications like frameworks like React and Angular. It's similar to C, C++, and F# in syntax structure and is accessible and easy to learn, making it a popular choice for beginners and professionals alike.
 2 Java: Known for its object-oriented programming (OOP) and support for Java 8 and .NET, Java is used for developing enterprise-level software applications, high-performance games, as well as mobile apps, game development, and desktop applications.
 3 C++: Known for its flexibility and scalability, C++ offers comprehensive object-oriented programming and is a popular choice for high-performance computing and other technical fields. It's a powerful platform for building real-world applications and games at scale.
 4 Python: Developed by Guido van Rossum in 1991, Python is a high-level, interpreted, and dynamically typed language known for its simplicity, readability, and versatility.

Loại và định dạng bộ dữ liệu mong đợi

KTO yêu cầu một tập dữ liệu thích thích không ghép cặp. Ngoài ra, bạn có thể cung cấp một tập dữ liệu thích thích ghép cặp (còn được gọi đơn giản là tập dữ liệu thích thích). Trong trường hợp này, trình huấn luyện sẽ tự động chuyển đổi nó sang định dạng không ghép cặp bằng cách tách các phản hồi được chọn và bị từ chối, gán label = True cho các phần hoàn thành được chọn và label = False cho các phần bị từ chối.

KTOTrainer tương thích với cả định dạng tập dữ liệu chuẩn và đối thoại. Khi được cung cấp một tập dữ liệu đối thoại, trình huấn luyện sẽ tự động áp dụng mẫu trò chuyện (chat template) lên tập dữ liệu đó.

# Standard format
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

# Conversational format
unpaired_preference_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                               "completion": [{"role": "assistant", "content": "It is blue."}],
                               "label": True}

Về mặt lý thuyết, tập dữ liệu nên chứa ít nhất một phần hoàn thành được chọn và một phần hoàn thành bị từ chối. Tuy nhiên, một số người dùng đã chạy thành công KTO chỉ sử dụng chỉ dữ liệu được chọn hoặc chỉ dữ liệu bị từ chối. Nếu chỉ sử dụng dữ liệu bị từ chối, nên áp dụng tốc độ học bảo thủ.

Tìm hiểu sâu hơn về phương pháp KTO

Kahneman-Tversky Optimization (KTO) là một phương pháp huấn luyện được thiết kế để điều chỉnh một mô hình ngôn ngữ bằng cách chỉ sử dụng tín hiệu nhị phân cho biết một đầu ra là mong muốn hay không mong muốn đối với một đầu vào cụ thể, thay vì các cặp hoàn thành được ưu thích/không được ưu thích. Dựa trên lý thuyết triển vọng của Kahneman và Tversky, phương pháp này xác định một mất mát nhận thức con người (HALO) trực tiếp tối đa hóa tiện ích của các chuỗi sinh ra, trong đó các ví dụ mong muốn và không mong muốn được gán trọng số bất đối xứng để phản ánh sự né tránh tổn thất của con người.

Phần này phân tích cách KTO hoạt động trên thực tế, bao gồm các bước chính: xử lý trước và tính toán hàm mất mát.

Tiền xử lý và phân từ

Trong quá trình huấn luyện, mỗi ví dụ được kỳ vọng chứa một prompt, một completion và một label boolean cho biết completion có mong muốn (True) hay không mong muốn (False). Để biết thêm chi tiết về các định dạng được kỳ vọng, xem Dataset formats. KTOTrainer sẽ tokenize mỗi đầu vào bằng tokenizer của mô hình.

Tính toán tổn thất

Số hạng phân kỳ KL được sử dụng bởi tổn thất KTO được ước lượng bằng cách ghép mỗi prompt với một phần hoàn thành không khớp được lấy từ nơi khác trong batch. Vì lý do này, các loại tổn thất ước lượng số hạng KL (tất cả ngoại trừ apo_zero_unpaired) yêu cầu kích thước batch huấn luyện trên mỗi thiết bị lớn hơn 1 và chiến lược lấy mẫu tuần tự, để các cặp không khớp ổn định trong suốt một batch.

Mất được sử dụng trong KTO (Phương trình 7 của bài báo) được định nghĩa như sau: $$\mathcal{L}_{\mathrm{KTO}}(\theta) = \mathbb{E}_{(x,y)}\!\left[ w(y)\Big(1 - v(x, y)\Big) \right]$$

trong đó giá trị $v(x, y)$ là $$v(x, y) = \begin{cases} \sigma\!\left(\beta\big(\log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)} - \mathrm{KL}\big)\right) & \text{if } y \text{ is desirable} \\ \sigma\!\left(\beta\big(\mathrm{KL} - \log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}\big)\right) & \text{if } y \text{ is undesirable} \end{cases}$$

Ở đây $x$ là prompt, $y$ là completion, $\pi_{\theta}$ là mô hình policy đang được huấn luyện, $\pi_{\mathrm{ref}}$ là mô hình tham chiếu, $\sigma$ là hàm sigmoid, $\beta > 0$ kiểm soát mức độ sai lệch so với mô hình tham chiếu, và $\mathrm{KL}$ là số hạng phân kỳ KL ước tính. Trọng số $w(y)$ là desirable_weight cho các completion mong muốn và undesirable_weight cho các completion không mong muốn, được sử dụng để khắc phục sự mất cân bằng giữa số lượng ví dụ mong muốn và không mong muốn.

Các loại hàm mất mát

Các chỉ số được ghi log

Trong quá trình huấn luyện và đánh giá, chúng ta ghi lại các chỉ số sau:

• global_step: Tổng số bước tối ưu hóa đã thực hiện cho đến nay.
• epoch: Số epoch hiện tại, dựa trên việc lặp qua tập dữ liệu.
• num_tokens: Tổng số token đã được xử lý cho đến hiện tại.
• loss: Mất KTO trung bình trong khoảng thời gian ghi log hiện tại.
• entropy: Entropy trung bình của phân phối token dự đoán của mô hình trên các token không bị mask.
• aux_loss: Tổn phụ trợ cân bằng tải của mô hình Mixture-of-Experts, trước khi được nhân với router_aux_loss_coef và cộng vào tổn. Chỉ được ghi log khi mô hình là mô hình MoE và router_aux_loss_coef khác không.
• kl: Trung bình của độ phân kỳ KL ước tính giữa mô hình chính sách và mô hình tham chiếu, được dùng làm điểm tham chiếu trong tổn thất KTO.
• learning_rate: Tốc độ học hiện tại, có thể thay đổi động nếu sử dụng trình lập lịch.
• grad_norm: Chuẩn L2 của các gradient, được tính trước khi cắt gradient.
• logits/chosen: Giá trị logit trung bình mà mô hình gán cho các token trong phần hoàn thành được chọn (mong muốn).
• logits/rejected: Giá trị logit trung bình mà mô hình gán cho các token trong phần hoàn thành bị từ chối (không mong muốn).
• logps/chosen: Xác suất log trung bình mà mô hình gán cho các token trong phần hoàn thành được chọn (mong muốn).
• logps/rejected: Xác suất log trung bình mà mô hình gán cho các token trong phần hoàn thành bị từ chối (không mong muốn).
• rewards/chosen: Phần thưởng ẩn trung bình được tính cho phần hoàn thành được chọn (mong muốn), được tính bằng $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/rejected: Phần thưởng trung bình ẩn được tính cho phần hoàn thành bị từ chối (không mong muốn), được tính bằng $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/margins: Biên độ phần thưởng ẩn trung bình giữa các câu trả lời được chọn và bị từ chối.

Tùy chỉnh

Tương thích và ràng buộc

Một số tổ hợp tham số bị hạn chế có chủ đích trong triển khai KTOTrainer hiện tại:

• PEFT adapters nhắm vào lm_head không được hỗ trợ (hãy dùng modules_to_save=["lm_head"] thay thế).
• sync_ref_model=True không được hỗ trợ khi huấn luyện với các mô hình PEFT không giữ một ref_model độc lập.
• sync_ref_model=True không thể kết hợp với precompute_ref_log_probs=True.
• precompute_ref_log_probs=True không được hỗ trợ với IterableDataset (train hoặc eval) hoặc với các bộ dữ liệu thị giác.
• Các loại loss ước lượng thành phần KL divergence (tất cả ngoại trừ "apo_zero_unpaired") yêu cầu train_sampling_strategy="sequential" và kích thước train batch trên mỗi thiết bị lớn hơn 1.

Khởi tạo mô hình

Bạn có thể truyền trực tiếp các kwargs của phương thức from_pretrained() vào KTOConfig. Ví dụ, nếu bạn muốn tải một mô hình ở độ chính xác khác, tương tự như

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct", dtype=torch.bfloat16)

Bạn có thể thực hiện điều này bằng cách truyền đối số model_init_kwargs={"dtype": torch.bfloat16} vào KTOConfig.

from trl import KTOConfig

training_args = KTOConfig(
    model_init_kwargs={"dtype": torch.bfloat16},
)

Lưu ý rằng tất cả các đối số từ khóa của from_pretrained() đều được hỗ trợ.

Huấn luyện adapter với PEFT

Chúng tôi hỗ trợ tích hợp chặt chẽ với thư viện 🤗 PEFT, cho phép bất kỳ người dùng nào thuận tiện huấn luyện các adapter và chia sẻ chúng trên Hub, thay vì phải huấn luyện toàn bộ mô hình.

from datasets import load_dataset
from trl import KTOTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

trainer = KTOTrainer(
    "Qwen/Qwen2-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)

trainer.train()

Bạn cũng có thể tiếp tục huấn luyện PeftModel của mình. Để làm điều đó, trước tiên hãy tải một PeftModel bên ngoài KTOTrainer và truyền nó trực tiếp cho trainer mà không truyền đối số peft_config.

Khi huấn luyện adapter, bạn thường sử dụng tốc độ học cao hơn so với fine-tuning toàn bộ vì chỉ có các tham số mới được học.

Gọi công cụ với KTO

KTOTrainer hỗ trợ đầy đủ việc tinh chỉnh mô hình có khả năng tool calling. Trong trường hợp này, mỗi ví dụ trong bộ dữ liệu nên bao gồm:

• Các tin nhắn hội thoại (prompt và completion), bao gồm các cuộc gọi công cụ (tool_calls) và phản hồi công cụ (tin nhắn có vai trò tool)
• Danh sách các công cụ khả dụng trong cột tools, thường được cung cấp dưới dạng schema JSON

Để biết chi tiết về cấu trúc bộ dữ liệu mong đợi, hãy xem phần Định dạng Dataset — Tool Calling.

Huấn luyện Mô hình Ngôn ngữ Thị giác

KTOTrainer hỗ trợ đầy đủ việc huấn luyện các Mô hình Ngôn ngữ - Thị giác (VLM). Để huấn luyện một VLM, hãy cung cấp một tập dữ liệu có cột image (một hình ảnh cho mỗi mẫu) hoặc cột images (danh sách các hình ảnh cho mỗi mẫu). Để biết thêm thông tin về cấu trúc tập dữ liệu mong đợi, hãy xem phần Định dạng tập dữ liệu — Tập dữ liệu thị giác.

from trl import KTOConfig, KTOTrainer
from datasets import load_dataset

trainer = KTOTrainer(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    args=KTOConfig(max_length=None),
    train_dataset=load_dataset("trl-internal-testing/zen-image", "conversational_unpaired_preference", split="train"),
)
trainer.train()

Đối với VLM, việc cắt ngắn có thể loại bỏ các token hình ảnh, dẫn đến lỗi trong quá trình huấn luyện. Để tránh điều này, hãy đặt max_length=None trong KTOConfig. Điều này cho phép mô hình xử lý độ dài chuỗi đầy đủ mà không cắt ngắn các token hình ảnh.

Đã sao chép

KTOConfig(max_length=None, ...)

Chỉ sử dụng max_length khi bạn đã xác minh rằng việc cắt bớt sẽ không loại bỏ các token hình ảnh cho toàn bộ bộ dữ liệu.

Mẫu script

Chúng tôi cung cấp một script mẫu để huấn luyện mô hình bằng phương pháp KTO. Script này có sẵn tại trl/scripts/kto.py

Để kiểm tra kịch bản KTO với mô hình Qwen2 0.5B trên dataset UltraFeedback, hãy chạy lệnh sau:

accelerate launch trl/scripts/kto.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/kto-mix-14k \
    --num_train_epochs 1 \
    --output_dir Qwen2-0.5B-KTO

Mẹo sử dụng

Đối với các Mô hình Hỗn hợp Chuyên gia: Bật tổn thất bổ trợ

MOE đạt hiệu quả cao nhất khi tải được phân bổ gần như đồng đều giữa các chuyên gia.
Để đảm bảo việc huấn luyện MOE tương tự trong quá trình tinh chỉnh theo sở thích, việc bổ sung tổn thất phụ từ bộ cân bằng tải vào tổn thất cuối cùng là rất có lợi.

Tùy chọn này được bật bằng cách đặt output_router_logits=True trong cấu hình mô hình (ví dụ MixtralConfig).
Để điều chỉnh mức độ đóng góp của tổn thất phụ vào tổng tổn thất, hãy sử dụng siêu tham số router_aux_loss_coef=... (mặc định: 0.001) trong cấu hình mô hình.

Gợi ý về kích thước lô

Hãy sử dụng kích thước lô (batch size) cho mỗi bước ít nhất là 4, và kích thước lô hiệu dụng nằm trong khoảng từ 16 đến 128. Ngay cả khi kích thước lô hiệu dụng của bạn lớn, nếu kích thước lô cho mỗi bước không tốt, thì ước lượng KL trong KTO sẽ không chính xác.

Gợi ý về tốc độ học

Mỗi lựa chọn của beta có một tốc độ học tối đa mà nó có thể chịu đựng trước khi hiệu suất học suy giảm. Với cài đặt mặc định là beta = 0.1, tốc độ học thường không nên vượt quá 1e-6 đối với hầu hết các mô hình. Khi beta giảm, tốc độ học cũng cần được giảm tương ứng. Nhìn chung, chúng tôi khuyến nghị mạnh mẽ giữ tốc độ học trong khoảng từ 5e-7 đến 5e-6. Ngay cả với các bộ dữ liệu nhỏ, chúng tôi cũng không khuyến khích sử dụng tốc độ học nằm ngoài khoảng này. Thay vào đó, hãy chọn nhiều epoch hơn để đạt được kết quả tốt hơn.

Dữ liệu mất cân bằng

Các tham số desirable_weight và undesirable_weight của KTOConfig đề cập đến trọng số được gán cho các tổn thất của các ví dụ mong muốn/dương và không mong muốn/âm. Mặc định, cả hai đều bằng 1. Tuy nhiên, nếu bạn có nhiều hơn một trong hai loại, bạn nên tăng trọng số cho loại ít phổ biến hơn sao cho tỷ lệ giữa (desirable_weight $\times$ số lượng ví dụ dương) và (undesirable_weight $\times$ số lượng ví dụ âm) nằm trong khoảng từ 1:1 đến 4:3.

KTOTrainer

>>> from trl import KTOTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

>>> trainer = KTOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

KTOConfig