Resource	Info
Paper	https://arxiv.org/abs/2501.12599
Code & Data	/
Public	arXiv
Date	2025.03.19

Summary Overview

Kimi k1.5 技术报告，主要聚焦于其中使用的 Reinforcement Learning 技术。

Main Content

扩展强化学习（RL）为人工智能的持续改进开辟了一条新的轴线，大型语言模型（LLM）有望通过学习探索奖励来扩展其训练数据。

Using RL with LLMs, the models learns to explore with rewards and thus is not limited to a pre-existing static dataset.

Key ingredients:

• Long context scaling: 观察到持续的性能提升伴随着增长的 context length.
• Improved policy optimization: We derive a formulation of RL with long-CoT and employ a variant of online mirror descent for robust policy optimization. This algorithm is further improved by our effective sampling strategy, length penalty, and optimization of the data recipe.

Moreover, we present effective long2short methods that use long-CoT techniques to improve short-CoT models. Specifically, our approaches include applying length penalty with long-CoT activations and model merging.

Stages: pretraining, vanilla supervised fine-tuning (SFT), long-CoT supervised fine-tuning, and reinforcement learning (RL).

Three key properties define a high-quality RL prompt set:

• Diverse Coverage
• Balanced Difficulty: The prompt set should include a well-distributed range of easy, moderate, and difficult questions to facilitate gradual learning and prevent overfitting to specific complexity levels.
• Accurate Evaluability: Prompts should allow objective and reliable assessment by verifiers, ensuring that model performance is measured based on correct reasoning rather than superficial patterns or random guess.

经验观察表明，一些复杂的推理问题可能有相对简单和容易猜测的答案，从而导致假阳性验证--即模型通过错误的推理过程得出正确答案。为解决这一问题，我们排除了容易出现此类错误的问题，如多项选择题、真/假题和基于证明的问题。此外，对于一般的答题任务，我们提出了一种简单而有效的方法来识别和删除易被破解的提示。具体来说，我们会提示模型猜测可能的答案，而无需任何 CoT 推理步骤。如果模型在 N 次尝试内预测出正确答案，则认为该提示太容易被破解，并将其删除。我们发现，设置 N = 8 可以删除大部分易被破解的提示。

Human-like reasoning:

• planning: the model systematically outlines steps before execution;
• evaluation: critical assessment of intermediate steps;
• reflection: enable the modle to reconsider and refine its approach;
• exploration: encourage consideration of alternative solutions;

Reinforcement Learning

When solvinig problem $x$, thoughts $z_t\sim\pi_\theta(\cdot|x,z_1,\cdots,z_{t-1})$ are auto-regressively sampled, followed by the final answer $y\sim\pi_\theta(\cdot|x,z_1,\cdots,z_{t-1})$ are auto-regressively sampled, followed by the final answer $y\sim\pi_\theta$ to denote this sampling procedure.

Rather than explicitly constructing a search tree and implementing a planning algorithm, we could potentially train a model to approximate this process.

Reward

• For verifiable problems, the reward is directly determined by predefined criteria or rules.
• For problems with free-form ground truth, we train a reward model $r(x,y,y^*)$ that predicts if the answer matches the ground truth.

Optimize the policy:

Policy Optimization

surrogate loss:

gradient

Length Penalty

This length-based reward is then added to the original reward with a weighting parameter.

In our preliminary experiments, length penalty may slow down training during the initial phases. To alleviate this issue, we propose to gradually warm up the length penalty during training. Sepcifically, we employ standard policy optimization without length penalty,followed by a constant lenght penalty for the rest of training.

Sampling Strategies

First, the RL training data we collect naturally come with different difficulty labels. Second, because the RL training process samples the same problem multiple times, we can also track the success rate for each individual problem as a metric of difficulty.

• Curriculum Sampling: We start by training on easier tasks and gradually progress to more challenging ones.
• Prioritized Sampling: We track the success rate $s_i$ for each problem $i$ and sample problems proportional to $1-s_i$, so that problems with lower success rates receive higher sampling probabilities. This directs the model's efforts toward its weakest areas, leading to faster learning and better overall performance.

Long2short: Context Compression for Short-CoT Models

Though long-CoT models achieve strong performance, it consumes more test-time tokens compared to standard short-CoT LLMs. However, it is possible to transfer the thinking priors from long-CoT models to short-CoT models so that performance can be improved even with limited test-time token budgets.

• Model Merging: we merge the two models by simply averaging their weights.
• Shortest Rejection Sampling
• DPO
• Long2short RL

RL Infrastructure

Rollout Phase & Training Phase: During the rollout phase, rollout workers, coordinated by a central master, generate rollout trajectories by interacting with the model, producing sequences of responses to various inputs. In the subsequent training phase, trainer workers access these experiences to update the model's weights. This cyclical process allows the model to continuously learn from its actions, adjusting its strategies over time to enhance performance.

Partial Rollouts for Long CoT RL: Partial rollouts is a key technique that effectively addresses the challenge of handling long-CoT features by managing the rollouts of both long and short trajectories. It ensures that no single lengthy trajectory monopolizes the system's resources. Moreover, since the rollout workers operate asynchronously, when some are engaged with long trajectories, others can independently process new, shorter rollout tasks. The implementation of partial rollouts also offers repeat detection. The system identifies repeated sequences in the generated content and terminates them early, reducing unnecessary computation while maintaining output quality. Detected repetitions can be assigned additional penalties, effectively discouraging redundant content generation in the prompt set.

Hybird Deployment of Training and Inference

• Training Phase: At the outset, Megatron and vLLM are executed within separate containers, encapsulated by a shim process known as checkpoint-engine. Megatron commences the training procedure. After the training is completed, Megatron offloads the GPU memory and prepares to transfer current weights to vLLM.
• Inference Phase: Following Megatron's offloading, vLLM starts with dummy model weights and updates them with the latest ones transferred from Megatron via Mooncake. Upon completion of the rollout, the checkpoint-engine halts all vLLM processes.
• Subsequent Training Phase: Once the memory allocated to vLLM is released, Megatron onloads the memory and initiates another round of training.

Results

🤖

Others