Source-linked AI summary

TextReg: Mitigating Prompt Distributional Overfitting via Regularized Text-Space Optimization

Lucheng Fu, Ye Yu, Yiyang Wang, Yiqiao Jin, Haibo Jin, B. Aditya Prakash, Haohan Wang

arXiv:2605.21318v2cs.CLcs.AIcs.LG

TL;DR

Prompt optimization can add narrow, sample-specific rules that fit training data but generalize poorly, while discrete text-space optimization lacks representation control. TextReg regularizes textual prompt updates using a dual-factor view of inefficiency. Across multiple reasoning benchmarks, it improves OOD generalization over existing methods.

  • Problem

    Prompt optimization can lengthen prompts and accumulate narrow, sample-specific rules that generalize poorly, while existing approaches provide little control over these coupled dimensions.

    Source: §1 · pp. 1–2
  • Method

    TextReg regularizes textual-gradient prompt updates using representational inefficiency, combining Dual-Evidence Gradient Purification, Semantic Edit Regularization, and Regularization-Guided Prompt Update.

    Source: §1 · p. 2
  • Results

    TextReg consistently outperforms existing prompt optimization methods on OOD generalization across datasets and test engines.

    Source: §1 · p. 2, §6 · p. 10
  • Takeaways & Limitations

    The evaluation supports improved OOD generalization across multiple reasoning benchmarks, while focusing primarily on single-turn reasoning with well-defined behavioral rules.

    Source: §6 · p. 10

    Semantic Edit Regularization estimates scope degradation directionally rather than numerically evaluating scope waste or representational inefficiency.

    Source: §4.3 · p. 7

Abstract

from arXiv · show

Large language models (LLMs) are highly sensitive to the prompts used to specify task objectives and behavioral constraints. Many recent prompt optimization methods iteratively rewrite prompts using LLM-generated feedback, but the resulting prompts often become longer, accumulate narrow sample-specific rules, and generalize poorly beyond the training distribution. We study this failure mode as prompt distributional overfitting and argue that it reflects a lack of representation control in discrete text-space optimization. We formalize this view through representational inefficiency, a dual-factor measure that decomposes prompt inefficiency into capacity cost and scope narrowness, attributing distributional prompt overfitting to their coupled growth during optimization. We propose TextReg, a regularization framework that realizes a soft-penalty objective through regularized textual gradients, combining Dual-Evidence Gradient Purification, Semantic Edit Regularization, and Regularization-Guided Prompt Update. Across multiple reasoning benchmarks, TextReg substantially improves out-of-distribution (OOD) generalization, with accuracy gains of up to +11.8% over TextGrad and +16.5% over REVOLVE.

1 Introduction

Prompt optimization can produce longer prompts with narrow, sample-specific rules that fit training data but generalize poorly. TextReg formalizes this inefficiency and regularizes prompt updates; it consistently outperforms existing methods on OOD generalization across datasets and test engines.

1 Introduction

  • Optimized prompts can grow longer and accumulate narrow, sample-specific instructions as they are refined, weakening generalization beyond training data.The paper calls this failure mode prompt distributional overfitting and says existing optimizers provide little control over either dimension.Source: §1 · pp. 1–2
  • TextReg regularizes textual-gradient updates with Dual-Evidence Gradient Purification, Semantic Edit Regularization, and Regularization-Guided Prompt Update.The stages distinguish broadly applicable rules from case-specific patches, estimate structural degradation from edits, and steer prompt rewriting.Source: §1 · p. 2
  • Figure 1 contrasts long prompts saturated with narrow rules against the goal of compact prompts composed of broadly applicable rules.The stated goal is stronger OOD generalization.Source: Figure 1 · p. 2
  • Representational inefficiency captures the coupled costs of prompt length and rules that apply only to narrow subsets of inputs.The paper frames distributional overfitting as a failure of representation efficiency, not merely an artifact of optimization dynamics.Source: §1 · p. 2
  • TextReg consistently outperforms existing prompt optimization methods on OOD generalization across datasets and test engines.The introduction reports results across multiple reasoning benchmarks.Source: §1 · p. 2

2 Related Work

Prior work improves prompts through evaluator feedback, robustness methods, and regularization, but the related approaches described here do not jointly control natural-language prompt capacity and rule scope. TextReg addresses that distinction by regularizing prompt representations.

Prompt optimization

  • Evaluator-feedback methods guide iterative prompt rewriting, but primarily use task feedback rather than explicitly regularizing representational inefficiency.Examples include APO, TextGrad [9], and REVOLVE [26].Source: §2 · p. 3

Prompt robustness and overfitting

  • Prompt overfitting remains distinct from robustness to wording changes and distribution shifts: feedback-based optimization can improve training performance while narrow instructions fail to transfer.The passage describes DLPO [30] as using static simplification instructions and REMO [31] as using external memory and meta-reflection.Source: §2 · p. 3

Regularization for generalization

  • TextReg regularizes natural-language prompts by jointly controlling capacity cost and scope narrowness, unlike methods that constrain parameters, continuous prompt embeddings, or prompt length alone.The cited examples of classical regularization include weight decay, dropout, LASSO and elastic net, and early stopping.Source: §2 · p. 3

3 Problem Formulation

The formulation treats prompts as representations optimized for a training distribution, while distinguishing OOD generalization from training loss. It defines representational inefficiency through prompt capacity and rule scope, then frames regularization as a trade-off whose structural signals differ in observability.

3.1 Problem Setup and Notation

  • Prompt optimization models a black-box language model mapping a prompt and input to an output, and seeks to minimize task-specific training loss.The formulation defines a prompt space and a dataset of input-output examples.Source: §3.1 · p. 3

3.2 Prompt Distributional Overfitting and Generalization

  • Prompt distributional overfitting occurs when optimization lowers training loss while increasing the generalization gap on a different target distribution.The paper notes this is often accompanied by longer prompts and more narrow, sample-specific instructions.Source: §3.2 · p. 4
  • The paper focuses on OOD generalization: performance may hold on source-distribution samples yet degrade on harder variants or related shifted tasks.The comparison is to classical overfitting, where additional capacity fits narrow patterns at the expense of generalization.Source: §3.2 · p. 4

3.3 Representational Inefficiency

  • For analysis, a prompt is decomposed into semantically distinct behavioral directives, even when its wording is not an explicit list of rules.Directives can describe reasoning procedures, behavioral constraints, or other task-relevant instructions.Source: §3.3 · p. 4
  • A directive’s scope concerns its relevance across task-preserving variation in surface form, scale, or difficulty.The paper contrasts broadly applicable guidance with instructions tied to a particular entity, such as the tomato example.Source: §3.3 · p. 4
  • Representational inefficiency is high when prompt capacity is consumed by directives that provide little broadly applicable guidance.Its multiplicative framing emphasizes that added capacity is increasingly inefficient when devoted to narrowly applicable instructions.Source: §3.3 · p. 4

3.4 Regularized Text-Space Prompt Optimization

  • The idealized regularized objective trades task performance against representational efficiency, but its structural components differ in observability.Prompt length directly reveals capacity cost, whereas scope narrowness is semantic; the formulation leaves concrete control to the optimization procedure.Source: §3.4 · pp. 4–5

4 Method

TextReg controls representational inefficiency through three stages that shape both optimization feedback and the prompt edits it produces. Its regularization guides task-faithful rewrites without suppressing needed task corrections.

4.1 Overview

  • TextReg controls representational inefficiency before and after rewriting: it purifies feedback, diagnoses realized edits, then mediates task-faithful corrections with structural guidance.The three stages regulate how optimization signals are formed and how they are expressed in the prompt.Source: §4.1 · p. 5

4.2 Dual-Evidence Gradient Purification

  • Local batch evidence suppresses sample-specific patches, while recurrence evidence favors stable task-level directives without treating recurrence as a numeric measure of scope.Repeatedly recovered directives receive stronger support as a soft prior; no numerical mapping or hard threshold is assumed.Source: §4.2 · p. 6

4.3 Semantic Edit Regularization

  • Semantic Edit Regularization measures capacity change exactly but estimates scope narrowing directionally, activating only the correction indicated by structural degradation.The capacity signal comes from token length, while scope is assessed semantically; the channels can trigger compression, generalization, both, or neither.Source: §4.3 · pp. 6–7

4.4 Regularization-Guided Prompt Update

  • Regularization-Guided Prompt Update keeps task fidelity primary, using structural guidance to shape task-faithful rewrites without suppressing beneficial expansion or specialization.The rewrite preference is conceptual rather than an explicitly evaluated numerical optimization.Source: §4.4 · pp. 7–8

5 Experiments

TextReg is evaluated for cross-dataset, cross-engine OOD generalization, with component ablations and role-wise degradation tests examining its performance and robustness. It consistently outperforms existing methods on most test pairs, and each component contributes to its results.

5.1 Experimental Setup

  • Prompts optimized on three source tasks are evaluated for cross-dataset transfer across six BBH datasets, three arithmetic datasets, and four open-source test engines.Accuracy is measured by exact match on the final answer.Source: §5.1 · p. 8
  • TextGrad, REVOLVE, and TextReg results average four independently seeded optimization runs; CoT uses the same fixed prompt under the same evaluation protocol.All resulting prompts are evaluated on the same test set with greedy decoding.Source: §5.1 · p. 9

5.2 Main Results

  • TextReg is best or second-best on nearly all test-engine and dataset pairs, outperforming TextGrad and REVOLVE in most cells.Its advantage is especially pronounced on harder Tracking Shuffled Objects variants.Source: §5.2 · p. 9

5.3 Ablation Study

  • Disabling any of the three components lowers mean OOD accuracy, while full TextReg achieves the highest accuracy on all four test engines.Figure 3 averages results over Tracking Shuffled Objects and Logical Deduction at five and seven objects.Source: §5.3 · p. 9, Figure 3 · p. 9

5.4 Resilience Study

  • Replacing any one of the Gradient, Regularization, or Optimizer roles with Qwen2.5-7B-Instruct causes only minor accuracy drops overall.On Tracking Shuffled Objects (7 objects), Weak Regularization and Weak Optimizer surpass the All-Strong baseline.Source: §5.4 · p. 10, Figure 4 · p. 10

6 Conclusion

The conclusion frames prompt distributional overfitting as inefficient representation growth and presents TextReg as a way to control it. Its evaluation supports improved OOD generalization on reasoning benchmarks, while broader prompt settings remain future work.

6 Conclusion

  • Prompt distributional overfitting is framed as longer prompts and narrow, sample-specific rules that can hurt OOD generalization.Representational inefficiency captures the interaction between capacity cost and scope narrowness.Source: §6 · p. 10
  • TextReg controls representational inefficiency through Dual-Evidence Gradient Purification, Semantic Edit Regularization, and Regularization-Guided Prompt Update.The paper reports improved OOD generalization over existing methods across multiple reasoning benchmarks.Source: §6 · p. 10
  • Evaluation focuses primarily on single-turn reasoning with well-defined behavioral rules; open-ended generation, broader generation tasks, multi-turn prompts, and agent instructions remain future directions.The conclusion notes that rule structure may be less explicit and capacity cost may be shared across turns in these settings.Source: §6 · p. 10

A.1 Dataset Details

The appendix describes reasoning datasets chosen to test transfer from source tasks to harder or related variants, then summarizes the three comparison baselines.

A.1 Dataset Details

  • The evaluation spans nine symbolic and arithmetic reasoning datasets, with BBH task families tested at three difficulty levels defined by object count.The two BBH families are Logical Deduction and Tracking Shuffled Objects.Source: §A.1 · p. 15
  • Prompts are optimized on three-object BBH variants and evaluated on five- and seven-object variants to measure transfer to harder instances of the same task structure.The source-task split is 50 / 100 / 100 for train, validation, and test.Source: §A.1 · p. 15
  • GSM8K is the arithmetic source task, while SVAMP and MultiArith provide related arithmetic evaluation datasets.SVAMP probes robustness to targeted problem variations; MultiArith requires composing elementary operations.Source: §A.1 · p. 15

A.2 Counterpart Details

  • The baselines are Zero-shot Chain-of-Thought, TextGrad, and REVOLVE.CoT cues step-by-step reasoning; TextGrad iteratively uses evaluator feedback, while REVOLVE tracks response changes across iterations.Source: §A.2 · p. 15

A.3 Implementation Details

The implementation uses shared optimization engines and a common training budget across methods, with fixed starting prompts for each task.

A.3 Implementation Details

  • Qwen2.5-7B-Instruct executes prompts, GPT-4o performs LLM-driven optimization, and four open-source engines evaluate optimized prompts.The test engines are Qwen2-7B-Instruct, Phi-3.5-Mini-Instruct, Llama-3-8B-Instruct, and Llama-3.1-8B-Instruct.Source: §A.3 · p. 15
  • All methods use batch size 3 for 12 optimization iterations, and TextReg uses a capacity-channel length-growth threshold τC = 0.2.GSM8K validation acceptance is relaxed by 1% uniformly across methods; reported TextGrad, REVOLVE, and TextReg results average four seeds.Source: §A.3 · p. 16

A.4 Initial Prompts

  • All methods and seeds start from the same fixed, task-agnostic prompt per task, using the released TextGrad prompts without method-specific tuning.The GSM8K prompt is unchanged; the BBH prompts use the same wording with the answer-format clause adapted for multiple-choice outputs.Source: §A.4 · p. 16

B Attribution of the Gains: Compute Budget and Additional Optimizers

TextReg’s OOD gains persist against matched-compute and additional-optimizer controls, while its prompts grow less and retain better OOD accuracy than competing trajectories.

B Attribution of the Gains: Compute Budget and Additional Optimizers

  • TextReg uses 174 backward calls versus 84 for TextGrad and REVOLVE, increasing GPT-4o cost from $0.63 to $0.98 and wall-clock time by about 13%.Its shorter prompts reduce solver input tokens by 36% during optimization; the matched-budget controls use 180 backward calls per run.Source: B Attribution of the Gains: Compute… · p. 16
  • TextReg has the highest mean accuracy on all four OOD evaluations; matched-budget search, length penalties, and other optimizers do not close the gap.The controls indicate that gains cannot be explained by extra LLM calls, limited baseline coverage, or length control alone.Source: Table 2 · p. 17, B Attribution of the Gains: Compute… · p. 17

C Prompt Growth and OOD Accuracy during Optimization

  • TextGrad grows from 41 to 351 tokens as OOD accuracy fluctuates and ends below its starting prompt, while TextReg stays at 135–160 tokens and retains a 5.0-point final gain.TextReg gains 5.5–6.6 points at steps 3–6; REVOLVE also stays compact but has weaker OOD accuracy.Source: C Prompt Growth and OOD Accuracy… · p. 17, Table 3 · p. 17
  • 79% of accepted RuleBank mentions match an existing entry, with runs ending at an average of 27.4 canonical rules.The table tracks prompt length and OOD accuracy on 7-object variants across optimization steps.Source: C Prompt Growth and OOD Accuracy… · p. 17, Table 3 · p. 17

D In-Distribution Performance and the ID–OOD Trade-off

TextReg’s OOD gains can accompany a modest ID-accuracy trade-off on BBH source tasks, but the trade-off is not universal.

D In-Distribution Performance and the ID–OOD Trade-off

  • TextReg is about two points below TextGrad in average ID accuracy on the two BBH source tasks, but incurs no ID loss on GSM8K.Table 4 reports held-out source-task accuracy using Qwen2.5-7B-Instruct as the test engine.Source: D In-Distribution Performance and the… · p. 17, Table 4 · p. 18
  • +7.1/+5.4 points on Logical Deduction 5/7 objects and +9.0/+6.4 on Tracking Shuffled Objects 5/7 objects are TextReg’s largest OOD gains, averaged over four test engines.These gains occur on the BBH source tasks where TextReg is about two points lower than TextGrad in average ID accuracy.Source: D In-Distribution Performance and the… · pp. 17–18
  • All three optimizers improve over the starting prompt on BBH source tasks, while removing Semantic Edit Regularization raises ID accuracy from 88.8 to 90.1 but weakens OOD robustness.Removing Gradient Purification lowers ID accuracy to 85.8 and also harms OOD performance.Source: D In-Distribution Performance and the… · p. 18

E Robustness of the Gains

Robustness tests vary the backward pipeline, starting prompt, test engine, and protocol; TextReg retains strong comparative performance across these checks.

E.1 Fully Weak Backward Pipeline

  • With Qwen2.5-7B-Instruct in every LLM-driven role, All-Weak TextReg retains 91–102% of All-Strong accuracy and ranks first in three of four matched comparisons.It remains within one point of the best in the fourth comparison and produces substantially shorter prompts.Source: §E.1 · p. 18

E.2 Dependence on the Initial Prompt

  • From richer initial prompts, TextReg has the highest mean in seven of eight cells; its final prompts are 104–140 words versus TextGrad’s 144 and 314 words.The starting prompts contain 55 and 57 words.Source: §E.2 · p. 18, Table 6 · p. 19

E.3 Test Engines and Protocol Sensitivity

  • On Qwen2.5-32B-Instruct, TextReg leads all four BBH transfers, with its clearest margin on Logical Deduction 7 objects; its ranking also persists when forward and test engines align.On Tracking Shuffled Objects 7 objects, TextReg scores 50.8 with Llama-3-8B-Instruct and 93.3 with Phi-3.5-Mini-Instruct.Source: §E.3 · p. 19
  • Sweeping τC from 0.05 to 0.5 changes average 7-object OOD accuracy by at most 2.5 points, while final prompt length remains 152–181 tokens.The authors describe performance as not highly sensitive to the exact threshold.Source: §E.3 · p. 19
  • Under strict non-degradation, each method accepts only 0–3 of 12 GSM8K updates; the SVAMP/MultiArith ranking remains unchanged, with TextReg leading.Averaged over the four test engines, SVAMP/MultiArith accuracies are 87.2/97.1 for TextReg, 86.9/96.7 for REVOLVE, and 86.8/96.4 for TextGrad.Source: §E.3 · p. 19

F Beyond Symbolic and Arithmetic Reasoning

Preliminary tests extend TextReg’s gains beyond symbolic and arithmetic reasoning. The appendix describes how its purification, rule tracking, regularization guidance, and prompt-update components operate.

F Beyond Symbolic and Arithmetic Reasoning

  • TextReg achieves the highest GPQA accuracy with less prompt growth than TextGrad, and improves over the starting prompt on CommonGen while producing the shortest final prompt.On CommonGen, TextGrad and REVOLVE fall below the starting prompt; these results are preliminary evidence of extension beyond the original reasoning benchmarks.Source: Table 7 · p. 20, F Beyond Symbolic and Arithmetic… · p. 20

G.1 Dual-Evidence Gradient Purification (Πgen)

  • Dual-Evidence Gradient Purification retains broadly applicable logic fixes, strips case-specific details, and rejects narrow patches or purely stylistic edits.RuleBank recurrence serves as a support signal when deciding whether to accept a gradient.Source: §G.1 · p. 20

G.2 RuleBank Canonicalization

  • RuleBank canonicalization extracts mid-level rules and increments an existing entry when a rule is structurally equivalent, otherwise inserting a canonical rule.Mention counts provide an empirical proxy for rule recurrence.Source: §G.2 · p. 21

G.3 Semantic Diff Analyzer (M∆)

  • The Semantic Diff Analyzer classifies rule-level changes and judges whether prompt specificity increased, decreased, or stayed neutral.It defaults to case-patch classification unless clear RuleBank or cross-context support favors a generalized rule.Source: §G.3 · p. 21

G.4 Regularization Gradient Generator (Γ)

  • The regularization gradient generator maps active directions to compression, generalization, combined regularization, or no-regularization guidance.Its directives preserve task semantics and avoid removing rules beneficial to correctness.Source: §G.4 · p. 22

G.5 Regularization-Guided Prompt Update — System Prompt

  • At each rewrite, the optimizer receives task feedback and semantic-edit regularization feedback through the prompt-update system.The system message defines REG_FEEDBACK as structural guidance for controlling prompt length and specificity.Source: §G.5 · p. 22

G.6 Regularization-Guided Prompt Update — User Message Trailing

  • When task and regularization feedback conflict, the update folds the task fix into the regularized edit when possible, prioritizing the task item only if necessary.Feedback targeting different prompt regions is applied to both regions.Source: §G.6 · p. 23

G.7 Regularization-Guided Prompt Update — User Message Skeleton

  • The user-message skeleton supplies the variable role, prompt text, task context, and task feedback, adding REG_FEEDBACK when regularization guidance is active.Without active regularization, the regularization section is empty.Source: §G.7 · p. 23
Loading 2605.21318v2…