以下内容转载自本人公众号《Road to AI》，如果内容对您有帮助，欢迎关注以便能接收最新的文章。
今天介绍的技术来自论文《Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking》，论文的作者Eric Zelikman来自Stanford，目前在XAI任职。之前有传闻说这篇论文是Q-Star背后的技术，甚至有人在youtube做了视频进行分析。抛开这些传闻，就该技术本身具备的创新性，作者在这个方向上的探索深度，已经具备成为一篇有影响力工作的基础。
为什么要用CoT
transformer架构下的LLM一直存在的一个问题就是，对于每一个token生成的过程使用的算力是相同的。因为LLM的工作方式就是这样，基于上下文给出下一个token预测的概率分布。这就造成模型回答“天空是什么颜色”和“证明费马大定理”2个问题时，背后的运作方式是完全相同的，显然这是不太符合认知的。因此就衍生出来2个方向，第一个方向是降低简单的token预测使用的算力，比如Mixture of Depths；第二种就是为复杂的问题在推理时分配更多的算力，比如各种CoT的方法。显然，第二个方向更加符合人们目前对于大模型的期望，因此相关的工作会更多一点，进而能诞生出像GPT-o1这样强大的模型。
Quiet-STaR的目的
论文提出了一种能在复杂问题上使用更多算力的方法，但坏消息是它并不具备根据问题难易自动调整的能力，意味着就算是最简单的问题也会使用同样多的算力，这就给未来的改进方法留下了非常多的空间。具体来说论文使用到的技术是模型在预测每一个token之前就行固定长度的思考输出，然后再基于思考的内容进行预测，注意是每一个token。
听起来和普通的CoT好像没有太大的区别，其不同之处在与论文中的方法是将CoT的能力变成模型内在的模式，训练不依赖于人工数据标注，并且方法具备普适性。这里的普适性是指，在文本数据上进行训练也能提升math测试集上面的能力。这些特性决定了这个方法更容易进行scale，论文中训练模型只用了1台8卡H100，基座模型只有7B的规模，理论上如果用更大的模型更长的训练时间，最后取得的效果也会更好。
论文内容
概括
论文是在预训练阶段验证该方法，理论上再SFT阶段一样成立。下面简单介绍一下该方法的主要流程：
-
在输入的每个token处生成对应的thought，因为方法没有标注数据，这里完全利用了大模型本身具备的输出能力。
具体来说，在每个前缀最后添加一个"<|startthought|>"字符，该字符的embedding初始化用的是em dash"---"的embedding，一般表示句子中间提顿思考的含义。
-
使用输入+对应位置的thought作为输入预测下一个单词，以此作为训练的主要loss。
-
使用with thought与without thought之间的预测概率提升作为监督信号，训练模型生成更好的thought，作为第二部分loss。
关键技术
Parallel Generation
在每个位置生成对应的thought所需要的算力无疑是巨大的，为了降低这部分消耗论文使用了Parallel Generation的方法。这个方法利用了transformer模型输出的特性，因为模型本身会输出每个位置的logits，只是在正常推理的时候，我们只需要最后一个token的logits。但此时模型需要每个位置token都输出thought，因此只需要设计好对应的mask方式，保证thought之间互相不可见即可。
“Mixing” (Residual) Heads
因为模型生成thought的方式决定了，刚开始生成的thought质量不会太好，因此将thought直接加入模型的训练loss会大幅提升语言模型对应的loss。因此，论文中使用一种加权平滑的方式权衡 with thought和without thought对应的logits。权重是一个标量值，由一个多层MLP输出，输入是thought最后一个字符对应和正常输入对应的hidden states。
Non-myopic Scoring and Teacher-forcing
模型训练的主要loss是带thouht的语言模型loss，目的是使模型习惯使用thought进行预测的方式。但这就决定了thought的好坏评价全部取决于它对应的1个token，这其实有台不符合常理，因为thought生效的位置并不一定是下一个token，因此在评价thought的好坏应该采用更长远的目标。如果基于当前的thought采样生成一定长度的字符串，并和真实的输出字符串进行对比，那么会引入太多的变量导致监督效果下降。因此论文中采用Teacher-forcing的方法，即假设thought接下来的字符为真实的输出。简单来说就是评估当前thought在未来token预测的时候的效果，评价指标还是预测概率提升。
REINFORCE loss
简单来说，REINFORCE方法类似于importance sampling，对每个位置生成多个thought，并根据每个thought对应的reward（预测概率提升）挑出其中比较好的部分给模型学习。REINFORCE发布于1992年，论文中通过理论证明了该方法的梯度于RL的梯度是正相关的。
下图是算法的完整流程：
总结
这篇论文的思路非常巧妙，但仍有不足之处。比如：
- 缺少简单问题上的适应能力，thought长度是一个超参数； 当然也可以由一个单独的模型在推理的时候给出thought建议长度，thought的长度也会影响训练时的效果。
- thought的好坏评价是由预测概率提升作为评价指标，本质上这是一种后验证方法，只有当thought产生出来才能判断其好坏。如果能在thought生成之前预测出thought的好坏，那么就可以使的模型在适当的位置插入thought。 本质上1和2方法都是Q函数可以干的事情，即需要一个价值函数预测thought的价值
- thought的初始化生成是分布外的，虽然后续使用REINFORCE对生成效果进行提升，但仍让人担心thought的质量。
整体来讲仍然是一篇非常高质量的论文，并且由于它和gpt-o1表现出特性具备的种种相似之处，不仅让人浮想联翩。总之如果想实现具备强大推理能力的模型，这个方法是值得尝试的，因为它具备了规模化的能力。
