Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TACO-RL: Task Aware Prompt Compression Optimization with Reinforcement Learning

Shivam Shandilya, Menglin Xia, Supriyo Ghosh, Huiqiang Jiang, Jue Zhang, Qianhui Wu, Victor Rühle

TL;DR

TACO-RL tackles the computational bottleneck of large prompts by introducing a task-aware prompt compression method trained with on-policy reinforcement learning. It uses a latency-aware encoder to perform binary token pruning, guided by task-specific rewards computed from outputs of a lightweight LLM (GPT-3.5-Turbo) and optimized via REINFORCE with entropy regularization. The approach yields substantial gains (8%–189%) across text summarization, question answering, and code summarization at the same compression rate and latency, demonstrating the value of task-aligned compression signals. This work offers a practical, generalizable framework for efficient prompt compression and provides open-source code to foster wider adoption.

Abstract

The increasing prevalence of large language models (LLMs) such as GPT-4 in various applications has led to a surge in the size of prompts required for optimal performance, leading to challenges in computational efficiency. Prompt compression aims to reduce the inference cost by minimizing input tokens without compromising on the task performance. However, existing prompt compression techniques either rely on sub-optimal metrics such as information entropy or model it as a task-agnostic token classification problem that fails to capture task-specific information. To address these issues, we propose a novel and efficient reinforcement learning (RL) based task-aware prompt compression method. To ensure low latency requirements, we leverage existing Transformer encoder-based token classification model while guiding the learning process with task-specific reward signals using lightweight REINFORCE algorithm. We evaluate the performance of our method on three diverse and challenging tasks including text summarization, question answering and code summarization. We demonstrate that our RL-guided compression method improves the task performance by 8% - 189% across these three scenarios over state-of-the-art compression techniques while satisfying the same compression rate and latency requirements.

TACO-RL: Task Aware Prompt Compression Optimization with Reinforcement Learning

TL;DR

TACO-RL tackles the computational bottleneck of large prompts by introducing a task-aware prompt compression method trained with on-policy reinforcement learning. It uses a latency-aware encoder to perform binary token pruning, guided by task-specific rewards computed from outputs of a lightweight LLM (GPT-3.5-Turbo) and optimized via REINFORCE with entropy regularization. The approach yields substantial gains (8%–189%) across text summarization, question answering, and code summarization at the same compression rate and latency, demonstrating the value of task-aligned compression signals. This work offers a practical, generalizable framework for efficient prompt compression and provides open-source code to foster wider adoption.

Abstract

The increasing prevalence of large language models (LLMs) such as GPT-4 in various applications has led to a surge in the size of prompts required for optimal performance, leading to challenges in computational efficiency. Prompt compression aims to reduce the inference cost by minimizing input tokens without compromising on the task performance. However, existing prompt compression techniques either rely on sub-optimal metrics such as information entropy or model it as a task-agnostic token classification problem that fails to capture task-specific information. To address these issues, we propose a novel and efficient reinforcement learning (RL) based task-aware prompt compression method. To ensure low latency requirements, we leverage existing Transformer encoder-based token classification model while guiding the learning process with task-specific reward signals using lightweight REINFORCE algorithm. We evaluate the performance of our method on three diverse and challenging tasks including text summarization, question answering and code summarization. We demonstrate that our RL-guided compression method improves the task performance by 8% - 189% across these three scenarios over state-of-the-art compression techniques while satisfying the same compression rate and latency requirements.
Paper Structure (38 sections, 16 equations, 4 figures, 8 tables, 1 algorithm)

This paper contains 38 sections, 16 equations, 4 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Prompt compression for LLMs.
  4. RL-based prompt compression.
  5. Background
  6. Prompt Compression
  7. REINFORCE Algorithm
  8. Methodology
  9. TACO-RL
  10. Encoder Model and Action Sampling.
  11. Reward Calculation.
  12. Policy Optimization.
  13. Overall Approach.
  14. Task-specific Rewards
  15. Reward Formulation
  16. ...and 23 more sections

Figures (4)

  • Figure 1: Encoder fine-tuning with RL using task specific reward signals on a Q/A task. The RL-guided compression model is able to understand the specificity of the question and retains the relevant context in the compressed prompt.
  • Figure 2: Comparison of QA F1 Scores and EM Counts across different compression rates for various models on the Squad Dataset. The bars represent QA F1 Scores, and the dashed lines represent EM Scores. The numbers on top of the bars represent the EM Counts. The two lines on top represent the scores with Original Context.
  • Figure 3: Comparison of BLEU score across different compression rates for various models on the CodeSearchNet dataset.
  • Figure 4: Impact on downstream task performance by different c, L values during training on the MeetingBank dataset.