Dự đoán đa token (MTP, cho phép mô hình dự đoán thêm một số token tiếp theo) đã trở thành một công thức tiền huấn luyện phổ biến, vừa cải thiện hiệu suất hạ nguồn, vừa có thể đóng vai trò là đầu nháp cho giải mã suy đoán. Các tác giả đã tách chính xác gradient cross-entropy của MTP thành hai phần: một thành phần tự chưng cất KL kéo đầu MTP về phía phân bố từ tiếp theo của chính mô hình (hữu ích, giúp cải thiện giải mã suy đoán), và một phần dư có tỷ lệ tín hiệu trên nhiễu rất thấp, gây nhiễu quá trình học và trở nên nghiêm trọng hơn khi quy mô tăng. Soft-MTP sử dụng “định tuyến gradient thẳng” (straight-through gradient routing) để chỉ giữ lại thành phần KL, tái sử dụng phép tính cross-entropy hiện có, không cần ma trận logits cho toàn bộ từ vựng và tương thích với nhân cross-entropy hợp nhất.
Thích 9 · Số sao mã nguồn 6
Lý do
- Thay đổi cực kỳ nhỏ: chỉ một dòng mã đối với MTP tiêu chuẩn, không thêm tham số, các tác giả tuyên bố không có chi phí bổ sung đo được; đồng thời hỗ trợ cả MTP kiểu DeepSeek và kiểu Meta.
- Các tác giả tuyên bố rằng ở mọi quy mô từ 167M đến 1.2B, độ perplexity của từ tiếp theo thấp hơn, giải mã suy đoán nhanh hơn, độ chính xác zero-shot cao hơn, khả năng suy luận toán học và mã hóa sau khi fine-tune ở quy mô 1.2B mạnh hơn, và không thua kém việc tối thiểu hóa trực tiếp KL trên hầu hết các chỉ số.
- Có thể tái tạo: bốn quy mô × ba mục tiêu (NTP, MTP, Soft-MTP) với tổng cộng 12 điểm kiểm tra (checkpoint) được đặt trên Hugging Face, cấu hình huấn luyện và tập lệnh đánh giá giải mã suy đoán đều có sẵn, trong thư mục tests/ có thể xác minh gradient khớp với KL tường minh ở độ chính xác float64; giấy phép MIT.
Cách sử dụng
- Mã nguồn: git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5 (FineWeb-Edu đã được token hóa trước, mỗi cấu hình huấn luyện 31.5B token, cần ít nhất 4 phân đoạn, quá trình lắp ráp cần khoảng 120 GB dung lượng ổ đĩa); trình tối ưu hóa Muon yêu cầu PyTorch phiên bản 2.9 trở lên.
- Huấn luyện: torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp (thay đổi thành hard_mtp / direct_kl / ntp để làm đối chứng).
- Xem mô hình có sẵn: AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True), sử dụng tokenizer gốc của gpt2; đánh giá giải mã suy đoán bằng lệnh python -m src.eval_spec --ckpts ….
- Để tích hợp vào mã huấn luyện của riêng bạn, hãy xem hàm combined_mtp_loss() trong src/models/loss.py.
Rủi ro
- Các mô hình được công bố đều là các mô hình nghiên cứu nhỏ dưới 1.2B, chỉ được huấn luyện trên FineWeb-Edu, không phải là các mô hình hội thoại có thể sử dụng trực tiếp; hiệu quả ở quy mô lớn hơn chưa được xác minh.
- Giá trị chủ yếu nằm ở giai đoạn tiền huấn luyện, cần tự huấn luyện từ đầu mới có thể hưởng lợi; đánh giá huấn luyện mặc định kết nối với Weights & Biases.
- Tóm tắt không cung cấp hệ số tăng tốc cụ thể của giải mã suy đoán, hiệu quả do các tác giả báo cáo; README nhắc nhở rằng tăng tốc thời gian thực (wall-clock) phụ thuộc vào GPU, việc so sánh cần được thực hiện trên cùng một máy.