💻 computer science

Graph-Aware Reinforcement Learning for Reusable Prompt Compression in Black-Box LLMs

This paper proposes a task-aware graph reinforcement learning framework that compresses reusable reasoning contexts in black-box LLMs by training a lightweight policy to make extractive keep-or-drop decisions on graph-structured reasoning units, achieving significant input-cost savings while preserving reasoning accuracy.

Original authors: Mehrshad Eskandarpour, Parmida Haddadnejad, Mohammadjavad Jannati

Published 2026-08-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Mehrshad Eskandarpour, Parmida Haddadnejad, Mohammadjavad Jannati

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a brilliant but very expensive robot how to solve complex puzzles. You don't just give the robot a single question; you have to hand it a thick instruction manual first. This manual contains the rules of the game, a few examples of how to solve similar puzzles, and a strict list of how to write down the final answer. In the world of Artificial Intelligence, these "robots" are called Large Language Models (LLMs), and the "manual" is the prompt. The problem is that these manuals are getting huge. Every time you ask the robot a new question, you have to resend the entire manual. This is slow, costs a lot of money (because the robot charges by the word), and fills up the robot's short-term memory.

Scientists have been trying to shrink these manuals without losing the important parts. Some have tried just cutting off the end of the text, while others have tried to summarize the whole thing into a few sentences. But here's the catch: if you cut out the wrong sentence, the robot might get confused and give a wrong answer, even if the rest of the text looks fine. The goal is to find a way to keep the manual short enough to be cheap and fast, but detailed enough to keep the robot smart. This paper tackles that exact problem, specifically for situations where you use the same manual over and over again for many different questions, like a teacher using the same lesson plan for a whole class.

The Paper's Big Idea: The "Smart Librarian" Robot

The authors of this paper, from the Iran University of Science & Technology, propose a new way to shrink these reusable manuals. They call their method Graph-Aware Reinforcement Learning. That sounds like a mouthful, so let's break it down with a story.

Imagine your reusable manual is a giant, messy library of sticky notes. Some notes are just general advice, some are specific examples, some are math formulas, and some are strict rules about how to format the answer. In the past, people tried to shrink the library by just grabbing the first few notes or picking notes that sounded similar to the question. But that's like trying to pack a suitcase by only grabbing the first few items you see; you might leave your toothbrush behind!

The authors suggest a smarter approach. First, they treat the library of sticky notes not as a simple list, but as a spiderweb (or a graph). In this web, every note is a node, and the threads connecting them show how the notes relate to each other. A formula note might be connected to an example note that uses it. A rule about "no negative numbers" might be connected to a specific math problem. This "spiderweb" helps the system understand that some notes are best friends and should stay together, while others are just acquaintances.

Next, they train a Smart Librarian using a technique called Reinforcement Learning. Think of this as a video game where the Librarian's job is to pick which sticky notes to keep and which to throw away. The Librarian doesn't know the answer to the puzzles inside the robot's brain (because the robot is a "black box"—we can't see its internal gears). Instead, the Librarian learns by trial and error. It picks a set of notes, sends them to the robot, and sees if the robot gets the answer right.

If the robot gets it right and the manual is short, the Librarian gets a high score.
If the robot gets it wrong, the Librarian gets a penalty.
If the Librarian throws away a note that was actually crucial (like a hidden rule), it gets a big penalty.

Over time, the Librarian learns exactly which notes are essential for the robot to succeed and which ones are just fluff. It learns to keep the "spiderweb" of logic intact, even if it means removing many notes.

What They Found: Shorter Manuals, Same Smart Robot

The researchers tested this "Smart Librarian" on two very tough types of tasks: math problems (like those found in the GSM8K and MATH datasets) and writing computer code (using the MBPP and HumanEval datasets). They compared their method against other ways of shrinking prompts, like just cutting the text in half or picking notes based on how similar they look to the question.

The results were quite impressive. The authors found that their method could shrink the reusable manual by 52.6%—meaning they removed more than half of the text! Despite cutting out so much, the robot's ability to solve problems only dropped by a tiny 1.0 percentage point. To put that in perspective, other methods that just cut text randomly or by similarity caused the robot's accuracy to drop by much more (sometimes over 8 percentage points).

Because they removed so many words, they also saved a lot of money and time. They estimated that using their compressed manual would save about 40.3% in input costs. In the real world, this means the robot would be faster to answer and cheaper to run, especially if you are asking it thousands of questions using the same lesson plan.

Why This Matters (and What It Doesn't Do)

The paper suggests that this method is a big step forward because it doesn't try to rewrite the manual or summarize it into new words. Instead, it simply selects the best existing pieces. This is important because it keeps the instructions clear and prevents the robot from getting confused by made-up summaries.

However, the authors are careful to note that this isn't a magic wand for every situation. Their method works best when you have a reusable manual that you use for many different questions. If you are asking a one-off question with a unique context, the time it takes to train the "Smart Librarian" might not be worth the savings. Also, the method relies on the manual being broken down into clear "sticky notes" (reasoning units) first; if the notes are messy to begin with, the Librarian might struggle.

In the end, the paper suggests that the future of efficient AI isn't just about making models bigger or faster, but about being smarter about what we feed them. By treating prompts like a connected web of ideas rather than a simple list of words, we can keep our AI assistants sharp, fast, and affordable without losing the magic that makes them work.

Technical Summary: Graph-Aware Reinforcement Learning for Reusable Prompt Compression in Black-Box LLMs

Problem Statement

Large Language Models (LLMs) increasingly rely on long, reusable reasoning contexts—comprising instructions, worked examples, constraints, formulas, and output-format rules—to perform complex tasks like mathematical reasoning and code generation. While these contexts enhance performance, their repeated use across numerous user queries incurs significant costs in input tokens, inference latency, and context-window usage.

Existing compression methods face critical limitations in black-box LLM settings (e.g., commercial APIs):

Truncation and Summarization: These often remove logically necessary information or rewrite content, risking hallucination and loss of interpretability.
Internal Signal Dependence: Many methods rely on gradients, logits, or attention scores, which are inaccessible in closed-source models.
Token-Level vs. Unit-Level: Token-level compression can break equations, code syntax, or logical dependencies.
Lack of Task-Awareness: Similarity-based approaches may retain semantically relevant but logically redundant units while discarding essential constraints.

The core challenge is to formulate an extractive, task-aware compression framework that operates strictly as a black-box, preserving the functional information required for correctness while minimizing token count, without access to internal model parameters.

Methodology

The authors propose a Graph-Aware Reinforcement Learning (RL) framework utilizing Proximal Policy Optimization (PPO). The method treats the reusable reasoning context as a graph and learns a policy to make discrete "keep-or-drop" decisions on reasoning units.

1. Problem Formulation

Template-Level Compression: The system optimizes a fixed reusable context ( $C$ ) that is paired with varying task-specific queries ( $x$ ). The goal is to learn a compressed context $C'(z)$ that is fixed for deployment, rather than compressing per query.
Black-Box Setting: The target LLM is treated as a fixed environment $F$ . The compressor observes only the input-output pairs and receives task-level rewards; no internal signals are used.
Unit-Level Segmentation: The context is segmented into semantically complete units (e.g., specific examples, constraints, formulas) rather than tokens to preserve syntactic and logical integrity.

2. Graph-Structured Representation

The reusable context is modeled as a graph $G_C = (V_C, E_C, W_C)$ :

Nodes: Represent reasoning units.
Edges: Encode four types of relations: semantic similarity, positional proximity, unit-type relations, and explicit dependencies.
Node Features: Each unit is embedded using local features and task-distribution-aware features (average/max similarity to the training query distribution) to ensure the policy learns general utility rather than query-specific relevance.

3. Policy and Optimization

Policy Network: A Graph Attention Network (GAT) encoder processes the graph to generate contextualized unit representations. A policy head outputs a keep probability for each unit.
Action Space: A binary mask ( $z \in \{0, 1\}^M$ ) determining which units to retain. Mandatory units (e.g., main instructions, output delimiters) are enforced to remain.
Reward Function: The reward balances four objectives:
1. Correctness: Task performance (e.g., exact answer match, unit-test pass) on a batch of queries.
2. Compression: The ratio of tokens removed.
3. Degradation Control: A penalty for accuracy loss exceeding a threshold $\epsilon$ .
4. Structural Coherence: A penalty ( $\Omega$ ) for fragmented selections where strongly connected units are inconsistently kept/dropped.
Algorithm: Proximal Policy Optimization (PPO) is used to stabilize training against the sparse and noisy rewards typical of black-box LLM evaluations.

4. Deployment

During deployment, the stochastic policy is replaced by a deterministic greedy selection based on the learned keep probabilities and unit lengths. The resulting compressed context is cached and reused for all future queries, recovering the offline training cost through repeated inference savings.

Key Contributions

Task-Aware RL Formulation: A framework that formulates reusable reasoning-context compression as a black-box, task-aware RL problem, optimizing for downstream correctness rather than just token reduction, extending prior work by focusing specifically on reusable templates and graph-structured dependencies.
Graph-Based Representation: Introduction of a graph model capturing semantic, positional, type-based, and dependency relations among reasoning units to prevent structurally fragmented compression.
Extractive PPO Policy: A discrete, extractive compression policy that learns keep-or-drop decisions without rewriting content or requiring internal model access, ensuring interpretability and API compatibility.
Comprehensive Evaluation: Extensive benchmarking on mathematical reasoning (GSM8K, MATH) and code generation (MBPP, HumanEval) demonstrating a strong balance between efficiency and accuracy.
Component Analysis: Detailed ablation studies validating the necessity of graph modeling, correctness-based rewards, and PPO optimization.

Experimental Results

Experiments were conducted on GSM8K, MATH, MBPP, and HumanEval using GPT-4o-mini as the primary target model, with cross-model transfer tests on Claude 3.5 Sonnet and Llama-3.1-70B-Instruct.

Compression Efficiency: The proposed method achieved a 52.6% compression ratio on reusable reasoning contexts.
Cost Savings: This resulted in an estimated 40.3% reduction in input-token costs and significant latency reductions.
Accuracy Preservation: Accuracy degradation was limited to 1.0 percentage point relative to the full-context baseline, significantly outperforming baselines like static truncation (8.4 pp loss) and similarity-based selection (5.2 pp loss).
Robustness: The method demonstrated superior stability across varying compression budgets (down to 25% retention) compared to heuristic and other RL-based baselines.
Cross-Model Transfer: When applied to unseen models without fine-tuning, the compressed context maintained high performance, with accuracy losses of 1.8 pp (Claude) and 2.7 pp (Llama), indicating the learned policy captures general reasoning patterns.

Significance and Claims

The paper argues that effective prompt compression for reasoning-heavy applications must prioritize functional necessity over mere token reduction. The authors claim that:

Graph-aware modeling is essential for preserving the logical dependencies between instructions, examples, and constraints that simple similarity metrics miss.
Task-level feedback (correctness) is a superior reward signal compared to output similarity or perplexity, as it directly aligns compression with the goal of solving the task.
Black-box compatibility is achievable without sacrificing performance, making this approach practical for commercial API-based LLM systems where internal signals are unavailable.
Reusable context optimization offers a distinct economic advantage: while training requires offline black-box calls, the resulting compressed template provides cumulative cost savings when reused across many queries, making it ideal for educational tools, coding assistants, and technical support systems.

The work concludes that while compression introduces risks (e.g., removing hidden constraints), the proposed framework mitigates these through structural regularization and degradation control, offering a practical path to reducing inference costs while maintaining reasoning quality.