Amazon 與 Booking.com 結合知識蒸餾和 Medusa 加速旅遊服務的大語言模型
Speeding up language models for travel applications
Amazon 團隊與 Booking.com 結合知識蒸餾(KD)和 Medusa,優化旅遊服務使用的大語言模型,最終模型在維持 70 billion 參數教師模型品質的情況下,推理速度提升 10 至 20 倍。
Travel platforms serve millions of daily users by answering questions about destinations, offering hotel recommendations, and helping users plan their trips. In recent years, many platforms have turned to large language models (LLMs) to improve these services. The LLMs give accurate, contextually relevant travel advice thanks to their size: they typically have tens of billions of parameters. Unfortunately, such large models are also computationally expensive, and that expense scales with the number of customer queries. What’s more, such large models can be too slow for web and mobile applications, where a delay of even 100 milliseconds can cost a company customers and revenue. To address these scaling and latency challenges, we explored two complementary optimization techniques. One is a compression technique known as knowledge distillation (KD), in which a slow but powerful model produces training targets for a faster, lightweight model. The other, called Medusa, is a framework that transforms sequential next-token predictions — a common bottleneck in LLMs — into parallel computations. Working with leading online travel agency Booking.com, we applied KD and Medusa techniques to the LLMs that power various aspects of its user experience, looking to improve either latency or cost effectiveness. We focused on three features in particular: The AI Trip Planner: A conversational assistant that must extract travel preferences from dialogue in less than a second (latency); Smart Filters: A feature that converts natural-language search queries into structured information filters in real time (latency); Accommodation Description Generation: A feature that produces property descriptions on the scale of millions of listings (cost effectiveness). We demonstrated that combining both KD and Medusa maintains the quality of the best open-source LLMs at less than one-third the cost while running up to 20 times as fast. Our framework ensures that travel platforms or other applications that rely on LLMs can deliver fast, cost-effective, accurate responses in real-world scenarios. We presented these techniques at last year’s meeting of the Association for Computational Linguistics (ACL), where it received the Industry Track best-paper award. Our approach: Medusa plus knowledge distillation Medusa is a framework designed to reduce LLM latency in real-time applications such as chat and code generation. It improves computational efficiency at run time — inference time — rather than during training. Ordinarily, an LLM generates text one token — word or word part — at a time. Medusa accelerates this process by predicting multiple tokens in parallel. Medusa does this by adding prediction modules, called Medusa heads, to the base LLM, one head for each token decoded in parallel. GPU hardware that LLMs run on is optimized for parallelism, so it can handle parallel token decoding with little additional overhead. Medusa offers significant speedups, typically 2-4-fold without compromising generation quality. Knowledge distillation addresses the deployment challenges of large, parameter-heavy models by transferring knowledge from larger “teacher” models to smaller “student” models. Typically, a teacher model is first trained or fine-tuned on high-quality labeled data, then used to generate predictions on additional unlabeled data. The manually labeled and machine-labeled data are then combined to train a smaller student model, which learns to mimic the teacher’s behavior. The student model retains much of the teacher’s performance but is significantly more efficient, making it better suited for real-world scenarios such as high-volume consumer-facing websites. In our work, we integrate KD and Medusa in a unified pipeline starting with pretrained base models. As in standard KD, we fine-tune the teacher on the labeled dataset and use it to generate predictions on unlabeled, domain-relevant text. During the student-training stage, however, we attach Medusa heads to the student model and train it on the combined corpus of teacher predictions and labeled data. This produces a single compact model that inherits the teacher's knowledge while also gaining multitoken-prediction capabilities. The experiment We chose TinyLlama — a compact model with 1.1 billion parameters — as our student and Llama-3.1 with 70 billion parameters as the teacher. We ran four progressive experiments, each building on the last. We first fine-tuned TinyLlama on task-specific data for each of our three use cases — AI trip planning, smart filtering, and accommodation description generation. Even without further optimization, the fine-tuned models were significantly more accurate than few-shot GPT-4o, confirming that domain-specific fine tuning compensates for their compact size. However, their performance still fell short of the much larger teacher model’s. Applying knowledge distillation — training the student on a blend of human-annotated and teacher-annotated data — closed that gap by more than 90%, bringing the small model to near parity with the teacher. The smaller model size also made inference roughly 6 times as fast, but that’s still not fast enough for real-time use. Adding Medusa heads, however, sped up inference an additional 2 to 3.6 times. We tested two Medusa variants: with Medusa-1, we first fine-tuned the base LLM, then froze its weights and fine-tuned the Medusa heads separately; with Medusa-2, we trained everything in a single stage. Both achieved comparable speedups with no meaningful quality loss, but Medusa-2's simpler pipeline made it more practical for production teams managing multiple models. Outcomes Combining KD with Medusa-2 delivered both gains simultaneously: our final model matched the quality of the 70-billion-parameter teacher while running 10 to 20 times as fast, depending on the use case. For each of the three use cases, we used the pipeline described above to fine-tune a dedicated model with domain-specific data. For AI Trip Planner, we showed a 2.9% increase in user engagement vs. GPT-3.5. For Smart Filters, we achieved responses that were 15 times as fast for 99 percent of queries. For AI-generated accommodation descriptions, we showed 1.4% more “helpful” ratings from users, at less than one-third the cost of GPT-4. With acknowledgements to our AWS colleague Aleksandra Dokic and to the Booking.com team: Emmanouil Stergiadis, Eran Fainman, Ilya Gusev, and Moran Beladev.
來源:Amazon Science · amazon.science