Self-Taught Optimizer (STOP): Recursively Self-Improving Code Generation

Self-Taught Optimizer (STOP)

一句话理解

STOP 先写一个会借助固定语言模型改进任意程序的初始 Improver,然后把这个 Improver 自己的源代码当成待改进程序,让它改写自己;改写后的 Improver 再进入下一轮。

STOP 改进的是调用语言模型的脚手架程序,而不是语言模型参数。

论文中的对象

语言模型

论文把语言模型记为:

$$
L:\Sigma^\rightarrow\Sigma^
$$

  • $\Sigma^*$:所有有限字符串的集合。
  • $L$:固定的、随机性的黑盒语言模型。
  • 输入和输出都是字符串。

Utility

论文把一个 Utility 写成:

$$
u=(u_{\mathrm{func}},u_{\mathrm{str}})
$$

  • $u_{\mathrm{func}}:\Sigma^*\rightarrow\mathbb R$:真正执行的评分函数。
  • $u_{\mathrm{str}}$:给语言模型阅读的自然语言或代码形式的评分说明。
  • 论文简写为 $u(s)$,表示程序 $s$ 的得分。

下游任务

一个下游任务是:

$$
\tau=(u,s)
$$

  • $u$:任务的评分标准。
  • $s$:该任务当前已有的初始程序。

例如,$s$ 可以是一个解背包问题的程序,$u(s)$ 是该程序在测试实例上的平均表现。

Improver

Improver $I$ 本身也是一个程序:

$$
s’=I(u,s,L)
$$

它接收评分标准 $u$、待改进程序 $s$ 和语言模型 $L$,输出一个新程序 $s’$。希望满足:

$$
u(s’)>u(s)
$$

Seed Improver $I_0$

论文先由人提供一个简单的初始 Improver $I_0$。其基本行为是:

  1. 把评分目标和当前程序放入 prompt。
  2. 多次调用语言模型 $L$,生成多个候选程序。
  3. 从回复中提取候选代码。
  4. 用 $u$ 实际运行并评分每个候选。
  5. 返回得分最高的候选。

因此 $I_0$ 已经能够改进普通下游程序,但它采用的搜索策略很简单。

如何评价一个 Improver

设下游任务来自分布 $\mathcal D$,理想目标是评价 Improver 在该分布上的期望表现:

$$
\bar u(I)
=
\mathbb E_{(u,s)\sim\mathcal D}
\left[u\left(I(u,s,L)\right)\right]
$$

真实任务分布无法穷举,所以论文取训练任务集:

$$
D={(u_i,s_i)}_{i=1}^{n}
$$

并定义经验 Meta-Utility:

$$
\hat u(I)
=
\frac{1}{|D|}
\sum_{(u,s)\in D}
u\left(I(u,s,L)\right)
$$

计算 $\hat u(I)$ 的顺序是:

  1. 对每个训练任务 $(u,s)$,让 $I$ 生成改进程序 $s’=I(u,s,L)$。
  2. 用该任务自己的 $u$ 评价 $s’$。
  3. 对所有任务得分取平均。

所以 $u$ 评价一个下游程序,$\hat u$ 评价一个 Improver。

记号说明:Algorithm 1 的辅助函数标题排成了 $\tilde u(I)$,但算法更新行和正文公式 (2) 使用 $\hat u(I)$。根据两处函数定义,它们都表示上述下游训练任务平均分,并不是两个不同目标。

STOP 的核心更新式

论文的核心定义是:

$$
I_t
:=
I_{t-1}(\hat u,I_{t-1},L)
$$

这三个实参分别对应 Improver 原本的三个输入:

形参 本轮传入的实参 含义
Utility $u$ Meta-Utility $\hat u$ 新程序按照“作为 Improver 的能力”评分
待改进程序 $s$ $I_{t-1}$ 的源代码 当前 Improver 自己成为修改对象
语言模型 固定的 $L$ 继续负责提出代码修改

这里能够自我应用,关键在于类型对齐:$I_{t-1}$ 原本能改进任何以 Utility 衡量的程序,而它自己也恰好是程序;$\hat u$ 又恰好能评价 Improver。

Algorithm 1 逐行解释

输入:

  • Seed Improver $I_0$。
  • 固定语言模型 $L$。
  • 迭代轮数 $T$。
  • 下游训练任务集 $D$。

循环第 $t$ 轮:

$$
I_t\leftarrow I_{t-1}(\hat u,I_{t-1},L)
$$

展开后发生的是:

  1. 当前 Improver $I_{t-1}$ 阅读自己的源代码。
  2. 它调用 $L$ 产生若干个候选 Improver。
  3. 对候选 Improver $J$,调用 $\hat u(J)$。
  4. $\hat u(J)$ 又让 $J$ 在所有训练任务上改进程序并计算平均得分。
  5. 当前 Improver 返回其中得分最高的候选,成为 $I_t$。

循环完成后返回 $I_T$。

前两轮的具体展开

第一轮:

$$
I_1=I_0(\hat u,I_0,L)
$$

$I_0$ 不再修改普通任务程序,而是修改 $I_0$ 的源代码。候选代码只有在下游训练任务上能取得更高平均分,才会被 $\hat u$ 认为更好。

第二轮:

$$
I_2=I_1(\hat u,I_1,L)
$$

此时由第一轮产生的 $I_1$ 改写自己的源代码。例如,$I_1$ 可能已经学会维护候选池、分轮搜索或接受较差解来增加探索;这些能力会参与第二轮自我改写。

一般地:

$$
I_0\rightarrow I_1\rightarrow I_2\rightarrow\cdots\rightarrow I_T
$$

最容易忽略的嵌套

外层正在寻找更好的 Improver,但每次评价一个候选 Improver,都要运行一次完整的下游优化过程:

$$
\text{候选 Improver }J
\longrightarrow
\left{
J(u_i,s_i,L)
\right}{i=1}^{|D|}
\longrightarrow
\left{
u_i(J(u_i,s_i,L))
\right}
{i=1}^{|D|}
\longrightarrow
\hat u(J)
$$

而 $J$ 在改进每个下游程序时还会多次调用 $L$。因此 STOP 的计算代价来自“候选 Improver 数量 × 下游任务数量 × 每个 Improver 的 LM 调用数”。

“Recursive” 应当怎样理解

这里的 Recursive 主要指反复自我应用:

$$
I_{t}=I_{t-1}(\hat u,I_{t-1},L)
$$

它不一定是程序运行时函数调用栈意义上的递归。更准确地说,是“上一轮得到的新 Improver,下一轮继续改写自己”。

它也不是完整的 Recursive Self-Improvement:

  • 语言模型 $L$ 的权重没有更新。
  • $I_0$、$\hat u$、训练任务集 $D$、预算和沙箱由人设定。
  • 改进范围主要是 LM 调用、候选搜索、筛选和执行策略组成的脚手架代码。

STOP 也没有划分出 MCE 那样的 Meta Agent 与 Base Agent。当前 Improver 在普通调用时改进下游程序;把自身源码作为输入时,同一个 Improver 又改进自己。外层循环、Meta-Utility、预算和执行沙箱仍是人工固定的 harness。

为什么它可能有效

初始 $I_0$ 只会多次采样并选最好候选。语言模型在改写这个搜索程序时,可以提出更复杂的优化策略,例如:

  • Beam search 或树搜索。
  • 维护 top-$k$ 候选的种群搜索。
  • 遗传算法式的变异和组合。
  • 模拟退火式的探索。
  • 任务分解以及探索/利用调度。

关键不是语言模型突然获得了新参数,而是相同模型被新的控制程序以更有效的方式调用。

论文结果与限制

  • 论文报告 GPT-4 驱动的 STOP 在其实验设置中能逐轮提高平均下游表现,并让得到的 Improver 迁移到未参与自我改进的新任务。
  • 较弱模型并不稳定;GPT-3.5 和 Mixtral 在论文的平均结果中可能退化。
  • 单轮更新没有单调改进保证。一个 Improver 擅长下游优化,并不必然擅长改写自己,语言模型采样和程序运行也具有随机性。
  • 评分只覆盖有限训练任务,候选 Improver 可能过拟合 $\hat u$。
  • 自动生成并执行代码带来安全风险,因此论文专门在沙箱中测试越界行为。

当前结论

STOP 的真正核心不是“模型训练自己”,而是下面这个类型闭环:

$$
\underbrace{I:(u,s,L)\mapsto s’}{\text{能够改进程序}}
\quad+
\underbrace{I\text{ 自己也是程序}}
{\text{可以作为 }s}
\quad+
\underbrace{\hat u\text{ 能评价 Improver}}_{\text{可以作为 }u}
$$

因此同一个改进接口可以作用到自身:

$$
(\hat u,I_{t-1},L)\xrightarrow{I_{t-1}}I_t
$$

这就是 Algorithm 1 最核心的数学和程序设计结构。

阅读问答

Q1:怎么理解 STOP 算法?

STOP 先用简单的 Seed Improver 改进普通程序,再把 Meta-Utility $\hat u$ 当成目标函数、把 Improver 自身源代码当成待改进程序,让当前 Improver 生成并筛选下一代 Improver。重复 $T$ 轮后,返回 $I_T$。

每一轮真正优化的对象都是 Improver 的代码;训练任务集的作用是判断这段新代码是否更会改进其他程序。

Q2:$\hat u(I)=\frac{1}{|D|}\sum_{(u,s)\in D}u(I(u,s,L))$ 中每个符号是什么?

这一定义的类型是:

$$
\hat u:\mathcal I\rightarrow\mathbb R
$$

即 $\hat u$ 接收一个候选 Improver $I$,返回一个实数分数。

  • $I$:待评价的 Improver 程序,而不是某个下游任务的最终答案。
  • $D$:有限的下游训练任务集合。
  • $(u,s)\in D$:遍历 $D$ 中的每个任务;每项由该任务的评分函数 $u$ 和初始解程序 $s$ 组成。
  • $|D|$:训练任务数量。
  • $L$:固定的黑盒语言模型。
  • $I(u,s,L)$:让 Improver 利用 $L$,按照 $u$ 改进初始程序 $s$;结果是新程序 $s’$。
  • $u(I(u,s,L))$:用当前任务自己的 $u$ 给新程序 $s’$ 评分,得到一个实数。
  • $\sum_{(u,s)\in D}$:把候选 $I$ 在所有训练任务上的得分相加。
  • $1/|D|$:除以任务数,得到平均分。
  • $\hat u$:经验 Meta-Utility;帽子表示它是用有限训练集 $D$ 对真实任务分布上的期望表现 $\bar u(I)$ 所作的经验估计。

计算顺序可写成:

$$
s_i’=I(u_i,s_i,L)
$$

$$
r_i=u_i(s_i’)
$$

$$
\hat u(I)=\frac{1}{n}\sum_{i=1}^{n}r_i
$$

如果:

$$
D={(u_1,s_1),(u_2,s_2),(u_3,s_3)}
$$

那么原式就是:

$$
\hat u(I)
=
\frac{1}{3}
\left[
u_1(I(u_1,s_1,L))
+u_2(I(u_2,s_2,L))
+u_3(I(u_3,s_3,L))
\right]
$$

需要区分两个层次:每个 $u_i$ 评价一个下游程序;$\hat u$ 汇总这些分数,评价 Improver $I$ 本身。原公式平均的是改进后程序的最终得分,并没有显式计算 $u(s’)-u(s)$。