LLM Modules: Knowledge Transfer from a Large to a Small Model using Enhanced Cross-Attention

TL;DR

The paper tackles the challenge of transferring knowledge from a large, pre-trained LLM to a smaller model under limited compute. It introduces LLM Modules that couple a frozen large model as a knowledge source with a smaller generation model through Enhanced Cross-Attention, enabling external representations to guide generation. Experimental results on Bespoke-Stratos-17k show that after 15 training epochs, the CombinedModel achieves generation quality and coherence approaching distillation-based baselines while reducing training costs. The approach offers a practical path for task-specific deployment of compact models and includes a publicly available codebase and weights to foster reproducibility and broader adoption.

Abstract

In this work, we propose an architecture of LLM Modules that enables the transfer of knowledge from a large pre-trained model to a smaller model using an Enhanced Cross-Attention mechanism. In the proposed scheme, the Qwen2-1.5B model is frozen and its representations are passed through specially designed attention layers to the GPT-Neo-125M model, which is trained on limited computational resources. Experimental results on the Bespoke-Stratos-17k dataset demonstrate that after 15 epochs of training, the combined model generates responses comparable in quality to those obtained by distillation. We discuss the advantages of the modular approach, provide examples of input queries and comparative analysis, and outline prospects for further extension of the method.