TextGrad: Improving Prompting Using AutoGrad
Last year researchers from Stanford released DSPy, a framework for automatic self prompting, this framework replaces the tedious task of writing human prompts that are often sub-optimal. DSPy was the biggest breakthrough in the LLM space after RAG (Retrieval Augmented Generation). Now these amazing researchers are back with TextGrad. It is a powerful framework performing automatic “differentiation” via text.
TextGrad backpropagates textual feedback provided by LLMs to improve individual components of a compound AI system. In this framework, LLMs provide rich, general, natural language suggestions to optimize variables in computation graphs, ranging from code snippets to molecular structures. TextGrad showed effectiveness and generality across a diverse range of applications, from question-answering and molecule optimization to radiotherapy treatment planning. So, without further ado, let's go deeper into this awesome paper.
I would highly recommend checking out our DSPy blog before we go further into TextGrad.
Other prompting blogs: Self-Rewarding Language Model, Promptbreeder: Prompting LLMs in a Better Way, and Giving self-reflection capabilities to LLMs
Topics Covered
- A Brief Intro To DSPy
- Understanding AutoGrad In PyTorch
- What Is TextGrad
- Prompt Optimization
- How Does TextGrad Works?
- Further Discussions
A Brief Intro To DSPy
DSPy is a framework for algorithmically optimizing LM prompts and weights, especially when LMs are used one or more times within a pipeline. To use LMs to build a complex system without DSPy, you generally have to: (1) break the problem down into steps, (2) prompt your LM well until each step works well in isolation, (3) tweak the steps to work well together, (4) generate synthetic examples to tune each step, and (5) use these examples to finetune smaller LMs to cut costs. Currently, this is hard and messy: every time you change your pipeline, your LM, or your data, all prompts (or finetuning steps) may need to change.
To make this more systematic and much more powerful, DSPy does two things. First, it separates the flow of your program (modules) from the parameters (LM prompts and weights) of each step. Second, DSPy introduces new optimizers, which are LM-driven algorithms that can tune the prompts and/or the weights of your LM calls, given a metric you want to maximize.
DSPy can routinely teach powerful models like GPT-3.5 or GPT-4 and local models like T5-base or Llama2-13b to be much more reliable at tasks, i.e. having higher quality and/or avoiding specific failure patterns. DSPy optimizers will "compile" the same program into different instructions, few-shot prompts, and/or weight updates (finetunes) for each LM. This is a new paradigm in which LMs and their prompts fade into the background as optimizable pieces of a larger system that can learn from data. tldr; less prompting, higher scores, and a more systematic approach to solving hard tasks with LMs.
Understanding AutoGrad In PyTorch
It is important to understand the AutoGrad of PyTorch because that’s where paper draws its analogy.
How Autograd Works
1. Forward Pass: In the forward pass, your neural network makes predictions by processing input data through its layers. This pass computes the loss function, which measures the difference between the network’s prediction and the actual target values.
2. Recording Operations: During the forward pass, autograd systems typically record all operations performed on tensors (data structures used to hold input, output, and weights) in a directed acyclic graph (DAG). Each node in the graph represents a tensor, and each edge represents a function that transforms one tensor into another.
3. Backward Pass: Once the forward pass is completed and the loss is calculated, the backward pass begins. In this phase, autograd computes the gradients of the loss function with respect to each tensor (parameter) that contributed to the loss. I)
This is done by traversing the DAG from the loss tensor back to the input tensors. Using the chain rule of calculus, autograd accumulates the partial derivatives of the loss with respect to each tensor.
4. Gradient Update: After the gradients are computed, they are used to update the parameters of the neural network. This is typically done using an optimization algorithm like stochastic gradient descent (SGD) or its variants (e.g., Adam, RMSprop). The updates are intended to minimize the loss, thus improving the model’s predictions.
AutoGrad
The core feature of autograd is its ability to perform automatic differentiation. This means it can automatically compute derivatives of complex functions, which is essential for training deep neural networks where manual calculation of derivatives is practically infeasible.
Dynamic Computation Graphs
Some autograd systems (like PyTorch) use dynamic computation graphs, which are built on-the-fly during the forward pass. This is in contrast to static graphs (used in TensorFlow v1.x), where the graph is defined once and then executed. Dynamic graphs are particularly useful for models where the computations depend on the input data, such as with varying input sizes or conditional operations.
What is TextGrad?
An autograd engine — for textual gradients!
TextGrad is a powerful framework building automatic differentiation via text. TextGrad implements backpropagation through text feedback provided by LLMs, strongly building on the gradient metaphor
TextGrad provides a simple and intuitive API that allows us to define our own loss functions and optimize them using text feedback. This API is similar to the Pytorch API, making it simple to adapt to our use cases.
So, let’s break it down further and see what is happening here.
Basically what we have here are two LLMs, one teacher LLM, which is supposedly more smart than the student LLM. Here teacher LLM is used as critique and helps to improve the prompts for student model.
Note: Keep in mind if a task is already tough for teacher LLM, then the entire pipeline doesn’t work and falls apart.
Prompt Optimization
Under prompt optimization, there are two main ideas.
First, DSPy pioneered the idea of viewing complex LLM-based systems as programs with potentially many layers and proposes ways to build and optimize them in a programmatic fashion. The framework is extensive, with results improving LLM performance in various question-answering, reasoning, and prompt optimization tasks.
Secondly, Prompt Optimization with Textual Gradients (ProTeGi) defines the Textual Gradients in the context of prompt optimization, where gradients are natural language feedback from LLMs given to the mistakes made during the task. While ProTeGi is built on the textual gradient analogy, TextGrad expands this analogy more broadly to automatic differentiation and goes substantially beyond prompt optimization tasks.
In particular, both DSPy and ProTeGi focused on prompt optimization, while a significant advance of TextGrad, as demonstrated through its diverse applications, is in instance optimization.
TextGrad takes a different perspective that backpropagation and its extensions can be a general and powerful framework to optimize the new generation of Al systems and perform multiple tasks outside of prompt optimization. In particular, we treat not only instructions or demonstrations as variables to optimize, but also the instances we care about themselves — such as molecules, treatment plans, code snippets, and so on.
How does ProteGei work?
Natural Language Gradients: Act like numerical gradients but in semantic space, pointing out deficiencies in the current prompt based on its perfoffnance on a given batch of data. Generated through a feedback loop where the LLM critiques the current prompt by identifying and describing its shortcomings.
Prompt Editing: Adjusts the prompt by addressing the issues highlighted by the natural language gradients. Uses another set of LLM instructions to modify the prompt in a way that it moves semantically opposite to the problems described by the gradients.
Beam Search and Bandit Selection: Aids in navigating through the space of potential prompt improvements by evaluating multiple candidate prompts and selecting the most promising ones.
- Expansion Step: Generates new prompt candidates based on current prompt evaluations.
- Selection Step: Employ strategies like UCB Bandits and Successive Rejects to identify the most effective new prompts with minimal computational expense.
ProTeGi integrates learning from mistakes (through feedback on existing prompts) with a strategic exploration of the new prompt configurations (through beam search and selection algorithms), aiming to find the optimal configuration in fewer steps.
The limitation of such a system is a large number of API calls.
Generation LLM: This model generates initial responses based on the current prompt. You might consider this the primary LLM prompt that you are optimizing. It operates by processing input data according to the existing prompt and generates outputs that are then assessed for their quality or accuracy.
Critique LLM: This model critiques the outputs from the Generation LLM. It analyses how well the prompt guided the generation of responses and identifies specific weaknesses in the prompt. The Critique LLM generates what the method refers to as “natural language gradients,” which are essentially detailed feedback on how the prompt could be altered to yield better responses from the Generation LLM. This feedback mimics the function of gradients in traditional numerical optimization by pointing out the direction in which the prompt should be modified to improve performance.
Prompt Modification
The concept of modifying a prompt to move “semantically opposite” to the problems described by the gradients is a key aspect of the ProTeGi method, which seeks to improve the performance of Large Language Models (LLMs) through prompt optimization. This process mimics numerical gradient descent but is here now applied in a semantic, or language-based, context.
Understanding Semantic Movement Opposite to Gradients
Identifying Shortcomings: The first step involves identifying specific shortcomings or errors in the current prompt's performance. This is done by evaluating the LLM’s responses against a set of criteria or expected outputs, leading to the generation of what is called “natural language gradients”. These gradients are essentially detailed, human-readable feedback that describes how and why the current prompt is failing. But it all depends on the cleverness of the LLM!
Semantic Direction: In numerical gradient descent, the gradient points in the direction of the steepest increase of a loss function; hence, moving against the gradient leads towards a local minimum of the loss. In the textual domain, moving “semantically opposite” to a gradient means adjusting the prompt in ways that address and correct the noted deficiencies. For example, if a gradient indicates that a prompt is too vague, making it more specific would be a move in the opposite direction.
How Does TextGrad Works?
Example computation graph. In traditional automatic differentiation, we compute gradients that provide a direction that would improve the variable with respect to a downstream loss with the chain rule. For a simple analog, let us look at the simple system:
System Overview
- Prediction: Generated using a prompt and a question:
- Evaluation: Generated by evaluating the prediction with evaluation instructions:
This process is represented as:
Gradient Computation
To improve the prompt based on the evaluation, gradients are computed as follows:
where we use ∇LLM for the gradient operator when the forward function is an LLM call. In particular, this function returns natural language feedback such as:
‘This prediction can be improved by. . . ’ where the feedback describes how to modify the variable to improve the downstream objective, analogous to gradients in optimization.
We first collect feedback to the Prediction variable using the evaluation. Then, given this feedback and the (Prompt LLM −−→ Prediction) call, we collect the feedback on the Prompt.
Textual Gradient Descent (TGD)
Objective functions
In numerical optimization and automatic differentiation, the objective function is typically a differentiable function, such as mean squared error or cross-entropy. In TextGrad, the objective can be a complex and potentially nondifferentiable function, where the domain and codomain of the function can be unstructured data. This choice adds important generality and flexibility to the framework. For instance, we demonstrate that the objective can be specified in natural language text and computed by prompting a language model, an output of a code interpreter running unit tests, or outputs of molecular simulation engines. For instance, a simple loss function for a code snippet can be the following:
here we can use this evaluation signal to optimize the code snippet, powered by the well-documented ability of LLMs to simulate human feedback, self-evaluate, and self-improve.
Further Discussions
The TextGrad draws its analogies from PyTorch AutoGrad. The backpropagation of gradients is the driving force of deep learning. We do not have gradients for compound systems of black box AI systems, but we can construct analogous backpropagation for text-based feedback, forming the basis of TextGrad.
Here are a few results of TextGrad, for starters, it looks pretty interesting idea. But I’ve my own doubts about its feasibility, especially across different types of tasks.
Example of how a prompt is optimized:
I would like to once again reiterate the system is only as good as our teacher model. If the teacher/critique model is bad, we can’t do much with this framework.
Click here for more articles on prompting-related research:
- DSPy: New Paradigm For Prompting
- Self-Rewarding Language Model
- Promptbreeder: Prompting LLM in a better way
- Giving self-reflection capabilities to LLMs
Thank you for reading! If you enjoyed this article and want more insights on AI, consider subscribing to my newsletter. Get exclusive content, updates, and resources directly to your inbox: https://medium.com/aiguys/newsletter
Writing such articles requires considerable effort and time. I would highly appreciate your support through claps and shares. Your engagement motivates me to write more beyond the hype on SOTA AI topics with the utmost clarity and simplicity. Don’t miss out on future insights — make sure to follow me on 𝕏 as well. Happy learning ❤
