EmbedGrad: Gradient-Based Prompt Optimization in Embedding Space for Large Language Models

TL;DR

EmbedGrad introduces a gradient-based method to optimize prompt embeddings, enabling precise, interpretable, and effective task adaptation for large language models without altering their architecture.

Contribution

This work presents EmbedGrad, a novel framework that refines prompt embeddings via gradient-based optimization, bridging the gap between prompt engineering and parameter tuning.

Findings

Significant accuracy improvements on mathematical reasoning tasks.

Consistent performance gains across various model sizes and tasks.

Enhanced reasoning capabilities through embedding refinement.

Abstract

Effectively adapting powerful pretrained foundation models to diverse tasks remains a key challenge in AI deployment. Current approaches primarily follow two paradigms:discrete optimization of text prompts through prompt engineering, or continuous adaptation via additional trainable parameters. Both exhibit limitations-discrete methods lack refinement precision while parameter-based techniques increase complexity and reduce interpretability. To address these constraints, we propose EmbedGrad, a novel framework that optimizes text prompt embeddings through gradient-based refinement. Our approach uniquely decouples training from deployment:during optimization,labeled examples guide precise embedding adjustments while preserving semantic meaning; during inference, only optimized embeddings integrate with user queries. This enables fine-grained calibration impossible in text space, such as enhancing the reasoning capability of prompts like please reason step by step. Comprehensive evaluations across mathematical reasoning, sentiment analysis, and causal judgment tasks demonstrate EmbedGrad's effectiveness:optimizing this reasoning prompt for Qwen2.5-Math-1.5B increased accuracy from 14.74\% to 58.96\% on mathematical problems. Consistent improvements were observed across model scales (0.5B-14B) and all tasks, with particularly significant gains for smaller models on complex problems like causal judgment. By bridging prompt engineering and parameter efficiency without architectural changes, our work establishes embedding refinement as a powerful new paradigm for task adaptation.