LatentLoop: Mở rộng mô hình Looped MoE với phân bổ tính toán bất đối xứng

Các mô hình Looped Mixture-of-Experts (MoE) kết hợp việc tái sử dụng tham số với tính toán thưa để cải thiện hiệu suất trong cùng một mức ngân sách. Các nghiên cứu trước đây cho thấy các lượt lặp khác nhau có thể đóng các vai trò khác nhau ngay cả khi chúng dùng chung tham số. Để tìm hiểu đóng góp của các lượt lặp khác nhau, chúng tôi học các hệ số residual riêng biệt trong một mô hình áp dụng một block dùng chung hai lần, qua đó quan sát thấy quy mô nhỏ hơn ở lượt lặp đầu tiên. Các thử nghiệm với hệ số cố định cho thấy việc giảm quy mô lượt lặp đầu tiên giúp duy trì hoặc cải thiện hiệu suất, trong khi việc giảm tương tự ở lượt lặp thứ hai lại gây hại. Sự bất đối xứng này thúc đẩy chúng tôi khám phá việc phân bổ tính toán bất đối xứng. Chúng tôi đề xuất LatentLoop, mô hình nén chuỗi của lượt lặp đầu tiên, giữ nguyên độ phân giải đầy đủ ở lượt lặp thứ hai và tái phân bổ phần tính toán tiết kiệm được để tăng hidden size. Chúng tôi cũng tìm hiểu mức độ mạnh mẽ mà lượt lặp đầu tiên cập nhật toàn bộ chuỗi, đồng thời chia sẻ trọng số MoE giữa các lượt lặp trong khi dùng các tham số attention riêng biệt để thích ứng với độ phân giải chuỗi khác nhau. Qua 15 cấu hình với quy mô từ 1,3B đến 7B tham số, trong đó ngân sách lớn nhất vượt quá 10^{21} FLOPs, LatentLoop liên tục làm giảm NLL trên tập validation.