TextGrad: Automatic "Differentiation" via Text
Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Zhi Huang, Carlos Guestrin, James Zou
Статья: https://arxiv.org/abs/2406.07496
Код: https://github.com/zou-group/textgrad

Мы отлично научились обучать большие сети или дифференцируемые комбинации сетей через бэкпроп. Но с LLM эволюция ушла вперёд и теперь мультиагентные системы состоят из комбинации LLM и тулов, которые не образуют дифференцируемой цепочки, узлы такого вычислительного графа (те самые LLM и тулы) соединены естественно-языковыми интерфейсами (общаются текстом) и часто живут у разных вендоров в разных датацентрах с доступом исключительно через API. То есть про бэкпроп можно забыть. Или нет?

Textgrad по сути реализует аналог бэкпропа, но через текст и с текстовыми градиентами. Как это выглядит?

Рассмотрим на простом примере. Есть два вызова LLM, мы хотим оптимизировать промпт в первом вызове:

Prediction = LLM(Prompt + Question) (1)
Evaluation = LLM(Evaluation Instruction + Prediction) (2)

Для этой цепочки можно собрать аналог бэкпропа, где используется градиентный оператор ∇LLM, если при прямом проходе вызывалась LLM. Этот оператор сам основан на LLM и в целом похож на паттерн Reflexion, он возвращает фидбек (критику, рефлексию) про то, как надо изменить переменную, чтобы улучшить итоговую objective. Он выдаёт что-то типа ‘This prediction can be improved by. . . ’

Внутри ∇LLM мы через промпт показываем “прямой проход LLM” в виде “Here is a conversation with an LLM: {x|y}”, затем вставляем критику (предыдущий градиентный оператор в цепочке) “Below are the criticisms on {y}: {∂L/∂y}” и наконец “Explain how to improve {x}.”

В примере с двумя вызовами сначала считаем

∂Evaluation/∂Prediction = ∇LLM(Prediction, Evaluation),

то есть получаем указания, как поменять переменную Prediction, чтобы улучшить Evaluation.

Затем узнаём как поменять Prompt через

∂Evaluation/∂Prompt =
∂Evaluation/∂Prediction * ∂Prediction/∂Prompt =
∇LLM(Prompt, Prediction, ∂Evaluation/∂Prediction).

На этом строится градиентный оптимизатор, который называется Textual Gradient Descent (TGD) и работает по принципу:

Prompt_new = TGD.step(Prompt, ∂Evaluation/∂Prompt).

Оптимизатор TGD.step(x, ∂L/∂x) тоже реализован через LLM и по сути задаётся промптом типа “Below are the criticisms on {x}: {∂L/∂x} Incorporate the criticisms, and produce a new variable.” выдающим новое значение переменной (в нашем случае Prompt).

В реальности промпты операторов позабористее (и наверное могли бы быть найдены методом текстового градиентного спуска, но, кажется, нет).

В общем случае вычисление может быть более сложным и задаваться произвольным вычислительным графом, где в узлах в качестве трансформаций могут быть вызовы как LLM, так и тулов и численных симуляторов. Если у узла несколько последователей, то все градиенты от них собираются и агрегируются перед тем, как идти дальше.

Ещё остаётся вопрос с objective function, которая для бэкпропа — это что-то дифференцируемое, типа L2-loss или кроссэнтропии, например. Здесь она может быть недифференцируемой и описанной человеческим языком и вычисляться через вызов LLM с промптом. Например, для кода оно могло бы выглядеть так:

Loss(code, target goal) =
LLM(“Here is a code snippet:{code}. Here is the goal for this snippet:{target goal}. Evaluate the snippet for correctness and runtime complexity.”).

Что в общем достаточно универсально и гибко. Лосс-функция человеческим языком — это прикольно.

В работе исследованы два класса задач: instance optimization (когда надо найти какое-то решение задачи, кусок кода или молекулу, например) и prompt optimization (когда надо найти промпт, улучшающий результаты на множестве запросов конкретной задачи).

В данном подходе можно также реализовывать оптимизацию батчами (с агрегацией “лоссов” по батчу), оптимизацию с ограничениями (например, требовать определённый формат ответа), аналог момента (когда оптимизатор видит и предыдущие итерации).

Проверили на нескольких классах задач: Coding, Problem Solving, Reasoning, Chemistry, Medicine.

1) В задачах на код надо надо сгенерить код, решающий задачи из LeetCode Hard. Постановка выглядит так:

Code-Refinement Objective = LLM(Problem + Code + Test-time Instruction + Local Test Results), где выделенный жирным Code оптимизируется через TextGrad.

Результат 36% Completion Rate.

2) На следующей задаче надо сделать solution optimization — улучшить решение сложной задачи из Google-proof Question Answering (GPQA), например, вопрос из квантовой механики или органической химии. Задача задана следующим образом:

Solution Refinement Objective = LLM(Question + Solution + Test-time Instruction)

TextGrad делал три итерации обновления решения с последующим мажоритарным голосованием. Получили 55%, лучший известный на момент результат. На MMLU с физикой или ML результат выше, чем у CoT.

3) Задача оптимизации промпта для reasoning на задачах из Big Bench Hard и GSM8k:

Answer = LLM(Prompt, Question)
Evaluation Metric = Evaluator(Answer, Ground Truth)

Оптимизировали промпт под более дешёвую gpt-3.5-turbo-0125 через фидбек от более сильной gpt-4o. Использовали минибатчи размера 3 и 12 итераций, то есть модель видела 36 обучающих примеров. Промпт обновлялся, если на валидации результат становился лучше.

Обошли Zero-shot Chain-of-Thought и DSPy.

4) Оптимизация молекулы, заданной SMILES нотацией с получением скоров affinity из тула Autodock Vina и druglikeness через Quantatiative Estimate of Druglikeness (QED) score из RDKit, то есть multi-objective loss:

Evaluation = LLM((Affinity(SMILES_i, target), Druglikeness(SMILES_i))
SMILES_{i+1} = TGD.step (SMILES_i, ∂Evaluation/∂SMILES_i)

Молекула инициализировалась маленьким фрагментом из функциональной группы, в качестве LLM использовалась gpt-4o. Применяли TextGrad к 58 таргетам из бенчмарка DOCKSTRING.

На каждом шаге TextGrad неплохо улучшал результаты и итоговые молекулы вполне достойны. При этом используются хемоинформатические тулы и получаются объяснимые решения.

5) Оптимизация плана радиотерапии. Задача с двумя вложенными циклами, оптимизировали внешний — гиперпараметры для внутреннего. Гиперпараметры θ задавались строкой.

Лосс для плана лечения P выглядел так: L = LLM(P(θ), g), g здесь — клинические цели.

Результат вроде как тоже осмысленный.


Интересный заход, выглядит достаточно универсально и применимо. Авторы оформили это всё в библиотеку (https://github.com/zou-group/textgrad) с API похожим на PyTorch. Вычисление лосса, градиента и шаг оптимизатора выглядят неотличимо.

Ожидаю большое и интересное будущее 🙂 Интересно было бы посмотреть и на фидбек в других модальностях, картинка там или звук, должно быть прикольно. Также напрашивается дальнейшее расширение фреймворка на тулы и RAG.

Вдруг кто ещё не видел

https://github.com/karpathy/LLM101n

LLM101n: Let's build a Storyteller

What I cannot create, I do not understand. -Richard Feynman

In this course we will build a Storyteller AI Large Language Model (LLM). Hand in hand, you'll be able create, refine and illustrate little stories with the AI. We are going to build everything end-to-end from basics to a functioning web app similar to ChatGPT, from scratch in Python, C and CUDA, and with minimal computer science prerequisits. By the end you should have a relatively deep understanding of AI, LLMs, and deep learning more generally.

Syllabus

Chapter 01 Bigram Language Model (language modeling)
Chapter 02 Micrograd (machine learning, backpropagation)
Chapter 03 N-gram model (multi-layer perceptron, matmul, gelu)
Chapter 04 Attention (attention, softmax, positional encoder)
Chapter 05 Transformer (transformer, residual, layernorm, GPT-2)
Chapter 06 Tokenization (minBPE, byte pair encoding)
Chapter 07 Optimization (initialization, optimization, AdamW)
Chapter 08 Need for Speed I: Device (device, CPU, GPU, ...)
Chapter 09 Need for Speed II: Precision (mixed precision training, fp16, bf16, fp8, ...)
Chapter 10 Need for Speed III: Distributed (distributed optimization, DDP, ZeRO)
Chapter 11 Datasets (datasets, data loading, synthetic data generation)
Chapter 12 Inference I: kv-cache (kv-cache)
Chapter 13 Inference II: Quantization (quantization)
Chapter 14 Finetuning I: SFT (supervised finetuning SFT, PEFT, LoRA, chat)
Chapter 15 Finetuning II: RL (reinforcement learning, RLHF, PPO, DPO)
Chapter 16 Deployment (API, web app)
Chapter 17 Multimodal (VQVAE, diffusion transformer)