🤖 AI

TAPR: Enhancing LLM Performance with a Task-Aware Prompt Rewriter

This paper introduces TAPR, a Task-Aware Prompt Rewriter trained via Group Relative Policy Optimization that reformulates user inputs into optimized prompts, significantly enhancing downstream LLM performance on tasks like question answering and arithmetic reasoning.

Original authors: Oliver Savolainen, Emanuele Bastianelli, Hosein Azarbonyad

Published 2026-08-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Oliver Savolainen, Emanuele Bastianelli, Hosein Azarbonyad

Original paper licensed under CC BY 4.0 (http://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 talk to a super-smart robot that knows almost everything in the world. This robot is called a Large Language Model (LLM). It's like a genius librarian who has read every book ever written, but it has a quirky personality: it only works its absolute best when you ask it questions in a very specific, perfect way. If you ask it a vague question like "Tell me about space," it might give you a boring, generic answer. But if you ask, "Explain the lifecycle of a star to a 10-year-old using a metaphor about baking," it might unlock its full potential and give you something magical.

The problem is that figuring out those perfect questions is hard. It's like trying to tune a radio to find a clear station; if you turn the dial just a tiny bit wrong, you only get static. Most people aren't experts at "prompt engineering" (the art of writing perfect instructions for AI), so they often get stuck with the static. Scientists have been trying to build a helper that can automatically fix your questions, turning your rough, messy requests into the perfect, polished instructions that the robot loves. This paper dives into that exact challenge: Can we teach a small AI to be a "prompt editor" that makes our big AI friends perform better?

Meet TAPR: The AI Editor That Hacks Your Questions

In this paper, the authors introduce a new tool called TAPR (Task-Aware Prompt Rewriter). Think of TAPR as a tiny, specialized editor sitting between you and the giant AI. When you type a question, TAPR doesn't just pass it along; it rewrites it first. It takes your simple "Answer the question" and transforms it into a detailed, super-clear instruction that tells the big AI exactly how to think, what format to use, and what to avoid.

But how does TAPR learn to be such a good editor? The authors didn't just teach it by showing examples. Instead, they used a method called Reinforcement Learning, which is like training a dog with treats. Here's how the training loop works:

The Rewrite: TAPR takes your original question and rewrites it.
The Test: The big AI (the "Task LLM") tries to answer the rewritten question.
The Judge: A special "Judge AI" looks at the answer and the rewritten question to decide: "Did this new question help the big AI get the right answer? Is the question itself clear and well-written?"
The Reward: If the Judge says "Great job!", TAPR gets a digital treat (a reward). If the answer was wrong or the question was confusing, TAPR gets no treat.

Over time, TAPR learns to write better and better questions because it wants those treats. The authors used a specific training technique called GRPO (Group Relative Policy Optimization), which they found to be much more stable and effective than older methods like PPO.

What Did They Find?

The team tested TAPR on three very different types of tasks:

Question Answering: Like asking "Who was the first US President?"
Summarization: Like asking an AI to summarize a long news article.
Math Reasoning: Like solving a tricky word problem about mixing sugar and water.

The results were promising, though not a magic wand for every situation.

The Good News: When TAPR (specifically the version based on a model called Phi-4-mini-instruct) was trained, it consistently improved the performance of the big AI. For example, on a math test called GSM8K, the accuracy jumped from a low 11.91% (when using the untrained base model to rewrite) to 83.60% after TAPR was trained. On a question-answering test called Natural Questions, accuracy went from 48.80% to 59.20%.
The "Secret Sauce": The authors discovered that TAPR learned to write prompts that were clearer, more structured, and often included "chain-of-thought" instructions (telling the AI to "think step-by-step"). This made the big AI much smarter at solving problems.
The "Judge" Factor: A key part of their success was using an LLM-as-a-Judge. Instead of just checking if the answer matched a textbook word-for-word (which can be unfair if the answer is correct but phrased differently), they used another AI to judge the quality and meaning of the answer. This helped TAPR learn to write prompts that produced genuinely better answers, not just ones that looked right on a checklist.

The Catch: It's Not Perfect Yet

The authors are careful to point out that this isn't a solved problem.

Inconsistency: Sometimes, TAPR made things better, but other times, the results were mixed or even slightly worse. For instance, on some summarization tasks, the untrained "base model" actually did a decent job, and TAPR didn't always beat it.
The "Selection" Idea: The authors tried a trick where TAPR would generate five different rewritten questions and then pick the best one. While this helped sometimes, it didn't consistently lead to better results, and it cost more computer power. They found no strong evidence that TAPR became a better "judge" of its own work just by being trained to write.
The Cost: Training TAPR takes time and computing power. The authors suggest that while it works, the extra cost might not always be worth it for every single task, especially since simple, generic prompts sometimes work just fine.

The Bottom Line

This paper suggests that we can indeed teach a small AI to be a prompt editor that helps big AIs perform better. By using reinforcement learning and a smart "Judge AI," TAPR learned to turn vague questions into clear, powerful instructions. While it's not a perfect solution that works 100% of the time, it shows a clear path forward: if we can automate the art of asking the right questions, we might be able to unlock the full potential of AI for everyone, not just the experts. The authors conclude that this is a viable and effective technique, but it still needs more refinement to be reliable in every real-world scenario.

Technical Summary: TAPR – Task-Aware Prompt Rewriter

Problem Statement

Large Language Models (LLMs) possess significant potential but often require carefully crafted input prompts to achieve optimal performance. This requirement creates a barrier for non-expert users who lack the expertise in prompt engineering to formulate instructions that fully leverage model capabilities. Existing optimization methods largely rely on manual, iterative adjustments or automated techniques that may produce non-interpretable prompts or fail to optimize for real-world utility due to reliance on surface-level metrics (e.g., exact match). There is a need for an automated system that can refine under-specified user queries into clear, robust instructions to ensure reliable downstream task performance.

Methodology

The authors propose Task-Aware Prompt Rewriter (TAPR), a framework designed to automatically refine and adapt user prompts for LLMs. The system consists of a smaller, specialized "TAPR LLM" that rewrites prompts, which are then executed by a larger, frozen "Task LLM."

Training Framework

The core of the methodology involves training the TAPR LLM using Reinforcement Learning (RL) with Group Relative Policy Optimization (GRPO). The training pipeline operates as follows:

Input: The TAPR LLM receives an initial user prompt.
Rewriting: It generates a revised prompt.
Execution: The revised prompt is fed to the frozen Task LLM to generate a task output.
Reward Calculation: A combined reward signal is computed based on two components:
- Task Performance Reward ( $T$ ): Instead of traditional metrics like exact match or ROUGE, the authors utilize an LLM-as-a-judge paradigm. A separate judge model evaluates the semantic correctness and quality of the Task LLM's response against the ground truth.
- Prompt Quality Reward ( $Q$ ): A dedicated reward (scored 0–5 by an LLM judge) evaluates the rewritten prompt itself, ensuring it adheres to prompt engineering principles (e.g., clarity, relevance) and does not alter the original intent or introduce hallucinations.
Optimization: The combined reward $R = \alpha T + \beta Q$ is used to update the TAPR policy via GRPO.

Key Architectural Choices

GRPO vs. PPO: The authors employ Group Relative Policy Optimization (GRPO) rather than Proximal Policy Optimization (PPO). GRPO updates the policy based on a group of sampled trajectories, removing the need for an explicit value function (critic) and demonstrating greater consistency in this specific application.
Prompt Selection Mechanism: An optional inference-time mechanism allows the TAPR model to sample multiple candidate prompts (at a higher temperature) and select the best one based on its own judgment, though the paper notes this does not consistently yield further gains.
Evaluation: The system is evaluated across diverse tasks including Question Answering (Natural Questions, HotpotQA), Summarization (CNN/Daily Mail, SciTLDR), and Arithmetic Reasoning (GSM8K).

Key Contributions

TAPR Framework: Introduction of a task-aware prompt rewriter trained via reinforcement learning to optimize prompts using feedback derived from both the rewritten prompt quality and the downstream task performance.
LLM-as-a-Judge Integration: The novel application of LLM-as-a-judge rewards for both task output evaluation and prompt quality assessment. The authors demonstrate that this approach improves training stability and effectiveness compared to relying solely on verifiable metrics like exact match.
GRPO for Prompt Rewriting: The first exploration of GRPO for enhancing prompt rewriting abilities, showing it achieves greater consistency and avoids the model collapse or stagnant rewards observed with certain PPO configurations in this context.
Empirical Validation: Extensive experiments showing that TAPR-generated prompts lead to higher accuracy and better summary quality compared to baseline prompts and rewrites from base models (untrained).

Experimental Results

The study utilized Phi-4-mini-instruct and LLaMA-3.2-3B-Instruct as TAPR models and LLaMA-3.1-8B-Instruct as the Task LLM.

Question Answering: On the Natural Questions (NQ) dataset, TAPR variants significantly outperformed the Base Model and Baseline prompts. For Phi-4-mini-instruct, accuracy improved from 48.80% (Base) to 59.20% (TAPR). Similar gains were observed on HotpotQA.
Summarization: On CNN/Daily Mail and SciTLDR, TAPR variants achieved the highest LLM-as-a-judge scores. Notably, Phi-4-mini-instruct TAPR variants achieved win-rates of up to 84.58% against baseline prompts on SciTLDR.
Arithmetic Reasoning: On GSM8K, the TAPR variant of Phi-4-mini-instruct achieved 83.60% accuracy, surpassing the baseline "SOLUTION" prompt (82.38%) and drastically outperforming the Base Model (11.91%), which failed to generate coherent solutions.
Ablation Studies:
- Prompt Quality Reward: Removing this reward severely limited training progress, with models often reverting to near-original prompts, indicating the reward is crucial for learning effective rewriting.
- GRPO vs. PPO: GRPO demonstrated reliable improvement, whereas PPO variants (with varying hyperparameters) resulted in stagnant rewards or degenerate outputs.
- Supervised Fine-Tuning (SFT): Pre-training with SFT improved base performance but did not translate to enhanced RL training outcomes.

Significance and Claims

The paper claims that TAPR demonstrates the viability of automated prompt rewriting as a practical technique to unlock the full potential of LLMs, particularly for smaller models. By training a smaller model to rewrite prompts, the system can enhance downstream task performance without requiring users to possess prompt engineering expertise.

The authors emphasize that their approach addresses the limitations of traditional evaluation metrics by utilizing LLM-as-a-judge to capture semantic correctness and real-world utility. They argue that while prompt engineering remains an open and complex problem, their method consistently produces clearer, more instructive, and more detailed prompts that align with best practices.

However, the paper maintains a modest tone regarding its limitations:

Inconsistency: Improvements are not uniform across all tasks and models; in some cases, simple generic prompts perform competitively against optimized rewrites.
Training Stability: RL for text generation is inherently unstable and sensitive to hyperparameters.
Cost-Benefit: The gains from prompt rewriting are not yet consistent enough to fully offset the additional computational costs of training and inference.
Scope: The study is limited to smaller, non-reasoning models and specific benchmarks, suggesting that future work is needed to validate these methods on larger models and in diverse, real-world agentic workflows.

In conclusion, the work establishes that a smaller model, when trained with GRPO and guided by semantic rewards, can effectively learn to rewrite prompts to improve the performance of larger LLMs, offering a pathway to reduce reliance on manual prompt engineering.