LazyLLM : Élagage dynamique des tokens pour une inférence efficace de LLM à contexte long

L'inférence des grands modèles de langage basés sur des transformateurs comporte deux étapes séquentielles : 1) une étape de préremplissage pour calculer le cache KV des prompts et générer le premier token, et 2) une étape de décodage pour générer les tokens suivants. Pour les prompts longs, le cache KV doit être calculé pour tous les tokens durant l'étape de préremplissage, ce qui peut augmenter significativement le temps nécessaire pour générer le premier token. Par conséquent, l'étape de préremplissage peut devenir un goulot d'étranglement dans le processus de génération. Une question reste ouverte : tous les tokens du prompt sont-ils essentiels pour générer le premier token ? Pour y répondre, nous introduisons une méthode nouvelle, LazyLLM, qui calcule sélectivement le KV pour les tokens importants à la prédiction du token suivant, à la fois pendant les étapes de préremplissage et de décodage. Contrairement aux approches d'élagage statiques qui suppriment le prompt d'un seul coup, LazyLLM permet aux modèles de langage de sélectionner dynamiquement différents sous-ensembles de tokens du contexte à différentes étapes de génération, même s'ils avaient été élagués aux étapes précédentes. Des expériences approfondies sur des jeux de données standards et diverses tâches démontrent que LazyLLM est une méthode générique pouvant s'intégrer aisément aux modèles de langage existants pour accélérer significativement la génération sans besoin de fine-tuning. Par exemple, dans la tâche de question-réponse multi-documents, LazyLLM accélère l'étape de préremplissage du modèle LLama 2 7B par un facteur de 2,34 tout en maintenant la précision.