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, Saravan Rajmohan · 2025
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.Answer "4,404.5 people" Question "How many people per square mile lived in Fresno in 2010?" QA Task Task Aware Compressed Text 2010 Census reported Fresno population 494,665.population density was 4,404.5 people per square mile (1,700.6/km2).racial makeup Fresno 245,306 (49.6%)White, 40,960 (8.3%) African American, 8,525 (1.7%) Native American 62,528.6%)Asian 849 (0 Pacific Islander 111,984.6%24,513 Hispanic 232,055 (46.9%).42.7% Mexican 0.4% Salvadoran 0.4% Puerto Rican.Non-Hispanic Whites were 30.0% in 2010, down from 72.6% in 1970 Task Specific RL Feedback GPT 3.5 "4,404.5 people" Original ContextThe 2010 United States Census reported that Fresno had a population of 494,665.The population density was 4,404.5 people per square mile (1,700.6/km²).The racial makeup of Fresno was 245,306 (49.6%)White, 40,960 (8.3%) African American, 8,525 (1.7%) Native American, 62,528 (12.6%)Asian (3.6% Hmong, 1.7% Indian, 1.2% Filipino, 1.2% Laotian, 1.0% Thai, 0.8% Cambodian, 0.7% Chinese, 0.5% Japanese, 0.4% Vietnamese, 0.2% Korean), 849 (0.2%) Pacific Islander, 111,984 (22.6%) from other races, and 24,513 (5.0%) from two or more races.Hispanic or Latino of any race were 232,055 persons (46.9%).Among the Hispanic population, 42.7% of the total population are Mexican, 0.4% Salvadoran, and 0.4% Puerto Rican.Non-Hispanic Whites were 30.0% of the population in 2010, down from 72.6% in 1970.LLMLingua Compressed Text 2010 Census Fresno,665 4 245,306 White 40,960.3%African American 8,525.7%Native American 62,528 Asian (3.6% Hmong 1.7% Indian 1.2% Filipino 1.2% Laotian 1.0% Thai 0.8% Cambodian 0.7% Chinese 0.5% Japanese 0.4% Vietnamese 0.2% Korean), Pacific Islander 111,984 other 24,513 two or more Hispanic Latino,055 42.7% Mexican 0.4% Salvadoran 0.4% Puerto Rican Non-Hispanic 30.0%72.6% 1970 ENCODER GPT 3.5Answer not in context Original ContextThe 2010 United States Census reported that Fresno had a population of 494,665.The population density was 4,404.5 people per square mile (1,700.6/km²).The racial makeup of Fresno was 245,306 (49.6%)White, 40,960 (8.3%) African American, 8,525 (1.7%) Native American, 62,528 (12.6%)Asian (3.6% Hmong, 1.7% Indian, 1.2% Filipino, 1.2% Laotian, 1.0% Thai, 0.8% Cambodian, 0.7% Chinese, 0.5% Japanese, 0.4% Vietnamese, 0.2% Korean), 849 (0.2%) Pacific Islander, 111,984 (22.6%) from other races, and 24,513 (5.0%) from two or more races.Hispanic or Latino of any race were 232,055 persons (46.9%).Among the Hispanic population, 42.7% of the total population are Mexican, 0.4% Salvadoran, and 0.4% Puerto Rican.Non-Hispanic Whites were 30.0% of the population in 2010, down from 72.6% in 1970.