kimi-k1-5-rl-long-context

%23%20%2F%2F%2F%20script%0A%23%20requires-python%20%3D%20%22%3E%3D3.12%22%0A%23%20dependencies%20%3D%20%5B%0A%23%20%20%20%20%20%22marimo%22%2C%0A%23%20%20%20%20%20%22numpy%3D%3D2.4.3%22%2C%0A%23%20%20%20%20%20%22matplotlib%3D%3D3.10.3%22%2C%0A%23%20%5D%0A%23%20%2F%2F%2F%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.23.13%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20return%20mo%2C%20np%2C%20plt%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Paper%20header%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Kimi%20K1.5%20%E2%80%94%20Long%20Context%20as%20Implicit%20Search%20Space%0A%0A%20%20%20%20**Paper%3A**%20arXiv%3A2501.12599%20(Kimi%20Team%2C%20Jan%202025)%0A%20%20%20%20**Core%20claim%3A**%20A%20simple%20RL%E2%80%93long-CoT%20loop%20with%20a%20length%20penalty%20can%20beat%0A%20%20%20%20complex%20planning%20algorithms%20(no%20MCTS%2C%20no%20value%20networks)%20on%20hard%20reasoning%20tasks.%0A%0A%20%20%20%20What%20this%20notebook%20demonstrates%3A%0A%20%20%20%201.%20Why%20longer%20context%20windows%20act%20like%20an%20automatic%20planning%20budget.%0A%20%20%20%202.%20How%20a%20length-normalized%20policy%20objective%20trades%20off%20accuracy%20vs.%20compute.%0A%20%20%20%203.%20A%20long2short%20distillation%20analogue%3A%20compress%20a%20quality%20model%20into%20fewer%20tokens.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%40mo.capture%0A%20%20%20%20def%20_notation_ref()%20-%3E%20str%3A%0A%20%20%20%20%20%20%20%20return%20(%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20r%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20**Notation%20reference**%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20Symbol%20%7C%20Meaning%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C--------%7C---------%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24n%24%20%7C%20number%20of%20available%20reasoning%20tokens%20(context%20budget%20%2F%20search%20steps)%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24p(x%2C%20y%5E*)%24%20%7C%20probability%20model%20assigns%20to%20correct%20answer%20%24y%5E*%24%20given%20problem%20%24x%24%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24%5Cmathcal%7BL%7D(y%2C%20y%5E*)%24%20%7C%20length%20of%20generated%20chain%20of%20thought%20%2B%20answer%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24%5Calpha%20%5Cin%20%5B0%2C1%5D%24%20%7C%20length%20penalty%20coefficient%3A%20%24%5Calpha%20%3D%200%24%20is%20unconstrained%2C%20%24%5Calpha%20%3D%201%24%20is%20maximal%20budget%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24r%20%5Cin%20%5C%7B0%2C1%5C%7D%24%20%7C%20rule-based%20reward%3A%201%20iff%20answer%20equals%20ground%20truth%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24V_%5Cmathrm%7Blong%7D%24%20%7C%20value%20of%20long-CoT%20teacher%20after%20length-normalized%20RL%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24V_%5Cmathrm%7Bshort%7D%24%20%7C%20value%20of%20distilled%20short-CoT%20student%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7C%20%24%5CDelta%24%20%7C%20quality%20delta%20%3D%20%24V_%5Cmathrm%7Bshort%7D%20%2F%20V_%5Cmathrm%7Blong%7D%24%20%7C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20_notation_ref()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Interactive%20controls%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20budget_slider%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D16%2C%0A%20%20%20%20%20%20%20%20stop%3D2048%2C%0A%20%20%20%20%20%20%20%20step%3D16%2C%0A%20%20%20%20%20%20%20%20value%3D256%2C%0A%20%20%20%20%20%20%20%20label%3D%22Context%20budget%20(reasoning%20tokens%2C%20n)%22%2C%0A%20%20%20%20)%0A%20%20%20%20length_penalty_slider%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0.00%2C%0A%20%20%20%20%20%20%20%20stop%3D1.00%2C%0A%20%20%20%20%20%20%20%20step%3D0.01%2C%0A%20%20%20%20%20%20%20%20value%3D0.35%2C%0A%20%20%20%20%20%20%20%20label%3D%22Length%20penalty%20coefficient%20(%CE%B1)%22%2C%0A%20%20%20%20)%0A%20%20%20%20difficulty_choice%20%3D%20mo.ui.dropdown(%0A%20%20%20%20%20%20%20%20options%3D%5B%22easy%22%2C%20%22medium%22%2C%20%22hard%22%5D%2C%0A%20%20%20%20%20%20%20%20value%3D%22hard%22%2C%0A%20%20%20%20%20%20%20%20label%3D%22Task%20difficulty%22%2C%0A%20%20%20%20)%0A%20%20%20%20compress_switch%20%3D%20mo.ui.switch(value%3DFalse%2C%20label%3D%22Long2short%20distillation%20ON%22)%0A%0A%20%20%20%20mo.hstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.plain(%22%23%23%23%20Controls%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20budget_slider%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20length_penalty_slider%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20difficulty_choice%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20compress_switch%2C%0A%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20justify%3D%22start%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20budget_slider%2C%0A%20%20%20%20%20%20%20%20compress_switch%2C%0A%20%20%20%20%20%20%20%20difficulty_choice%2C%0A%20%20%20%20%20%20%20%20length_penalty_slider%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20budget_slider%2C%0A%20%20%20%20compress_switch%2C%0A%20%20%20%20difficulty_choice%2C%0A%20%20%20%20length_penalty_slider%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20np%2C%0A%20%20%20%20plt%2C%0A)%3A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Synthetic%20long-CoT%20accuracy%20model%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20%23%20Shape%3A%20accuracy%20improves%20with%20budget%20up%20to%20a%20per-difficulty%20ceiling%2C%0A%20%20%20%20%23%20then%20saturates.%20Length%20penalty%20pulls%20the%20spend%20back%20toward%20lower%20budgets.%0A%0A%20%20%20%20_n_steps%20%3D%20200%0A%20%20%20%20_n%20%3D%20np.linspace(1%2C%202048%2C%20_n_steps)%0A%0A%20%20%20%20_ceiling%20%3D%20%7B%22easy%22%3A%200.98%2C%20%22medium%22%3A%200.85%2C%20%22hard%22%3A%200.62%7D%5Bdifficulty_choice.value%5D%0A%0A%20%20%20%20def%20_budget_reward(_n_local%3A%20np.ndarray%2C%20_alpha%3A%20float)%20-%3E%20np.ndarray%3A%0A%20%20%20%20%20%20%20%20_base%20%3D%201%20-%20np.exp(-_n_local%20%2F%20220)%0A%20%20%20%20%20%20%20%20_len_cost%20%3D%201%20-%20_alpha%20*%20(_n_local%20%2F%202048)%0A%20%20%20%20%20%20%20%20_r%20%3D%20_base%20*%20_len_cost%0A%20%20%20%20%20%20%20%20return%20np.clip(_r%20*%20_ceiling%2C%200%2C%201)%0A%0A%20%20%20%20_alpha_selected%20%3D%20length_penalty_slider.value%0A%20%20%20%20_budget_selected%20%3D%20budget_slider.value%0A%0A%20%20%20%20_curve_no_penalty%20%3D%20_budget_reward(_n%2C%200.0)%0A%20%20%20%20_curve_selected%20%3D%20_budget_reward(_n%2C%20_alpha_selected)%0A%20%20%20%20_opt_idx%20%3D%20int(np.argmax(_curve_selected))%0A%20%20%20%20_opt_budget%20%3D%20int(_n%5B_opt_idx%5D)%0A%20%20%20%20_opt_reward%20%3D%20float(_curve_selected%5B_opt_idx%5D)%0A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Plot%3A%20implicit%20search%20budget%20tradeoff%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20_fig1%2C%20_ax1%20%3D%20plt.subplots(figsize%3D(6%2C%203.4))%0A%20%20%20%20_ax1.plot(_n%2C%20_curve_no_penalty%2C%20label%3Df%22%CE%B1%20%3D%200.00%20%20(unconstrained)%22%2C%20linewidth%3D1.6)%0A%20%20%20%20_ax1.plot(%0A%20%20%20%20%20%20%20%20_n%2C%0A%20%20%20%20%20%20%20%20_curve_selected%2C%0A%20%20%20%20%20%20%20%20label%3Df%22%CE%B1%20%3D%20%7B_alpha_selected%3A.2f%7D%20%20(current)%22%2C%0A%20%20%20%20%20%20%20%20linewidth%3D2.2%2C%0A%20%20%20%20%20%20%20%20color%3D%22%23e35d5d%22%2C%0A%20%20%20%20)%0A%20%20%20%20_ax1.axvline(%0A%20%20%20%20%20%20%20%20_opt_budget%2C%0A%20%20%20%20%20%20%20%20color%3D%22%23e35d5d%22%2C%0A%20%20%20%20%20%20%20%20linestyle%3D%22--%22%2C%0A%20%20%20%20%20%20%20%20alpha%3D0.7%2C%0A%20%20%20%20%20%20%20%20label%3Df%22optimal%20%E2%89%88%20%7B_opt_budget%7D%20tokens%22%2C%0A%20%20%20%20)%0A%20%20%20%20_ax1.set_xlabel(%22Available%20context%20%2F%20search%20budget%22)%0A%20%20%20%20_ax1.set_ylabel(%22Expected%20task%20reward%22)%0A%20%20%20%20_ax1.set_title(%0A%20%20%20%20%20%20%20%20f%22Implicit%20search%20reward%20%E2%80%94%20%7Bdifficulty_choice.value%7D%20difficulty%22%0A%20%20%20%20)%0A%20%20%20%20_ax1.legend(fontsize%3D8%2C%20frameon%3DFalse)%0A%20%20%20%20_ax1.grid(True%2C%20alpha%3D0.25)%0A%20%20%20%20_fig1.tight_layout()%0A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Plot%3A%20optimal%20budget%20vs%20%CE%B1%20slope%20(robustness%20sweep)%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20_alphas%20%3D%20np.linspace(0%2C%201%2C%2080)%0A%20%20%20%20_opt_budgets%20%3D%20np.array(%0A%20%20%20%20%20%20%20%20%5Bint(_n%5Bnp.argmax(_budget_reward(_n%2C%20float(a)))%5D)%20for%20a%20in%20_alphas%5D%0A%20%20%20%20)%0A%20%20%20%20_opt_rewards%20%3D%20np.array(%0A%20%20%20%20%20%20%20%20%5Bfloat(np.max(_budget_reward(_n%2C%20float(a))))%20for%20a%20in%20_alphas%5D%0A%20%20%20%20)%0A%0A%20%20%20%20_fig2%2C%20_ax2%20%3D%20plt.subplots(figsize%3D(6%2C%202.9))%0A%20%20%20%20_ax2.plot(_alphas%2C%20_opt_budgets%2C%20linewidth%3D1.8%2C%20color%3D%22%234a90d9%22)%0A%20%20%20%20_ax2.axvline(%0A%20%20%20%20%20%20%20%20_alpha_selected%2C%20color%3D%22%23e35d5d%22%2C%20linestyle%3D%22--%22%2C%20alpha%3D0.8%2C%20label%3D%22current%20%CE%B1%22%0A%20%20%20%20)%0A%20%20%20%20_ax2.set_xlabel(%22Length%20penalty%20%CE%B1%22)%0A%20%20%20%20_ax2.set_ylabel(%22Optimal%20context%20budget%22)%0A%20%20%20%20_ax2.set_title(%22How%20aggressively%20should%20the%20model%20be%20budget-constrained%3F%22)%0A%20%20%20%20_ax2.grid(True%2C%20alpha%3D0.25)%0A%20%20%20%20_ax2.legend(fontsize%3D8%2C%20frameon%3DFalse)%0A%20%20%20%20_fig2.tight_layout()%0A%0A%20%20%20%20%23%20%E2%94%80%E2%94%80%20Long2short%20summary%20%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%0A%20%20%20%20if%20compress_switch.value%3A%0A%20%20%20%20%20%20%20%20_long2short_note%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22**Long2Short%20distillation%20active.**%5Cn%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22Teacher%20uses%20%E2%89%88%20%7B_opt_budget%7D%20reasoning%20tokens%20at%20quality%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%7B_opt_reward%3A.1%25%7D.%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22Distilled%20student%20compresses%20to%20**%7Bmax(16%2C%20_opt_budget%20%2F%2F%204)%7D%20tokens**%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22with%20simulated%20compression%20ratio%20~3%E2%80%934%C3%97.%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22Expected%20short-CoT%20reward%20%E2%89%88%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%7B(1%20-%200.25)%20*%20_opt_reward%3A.1%25%7D%20(conservative%20compression%20overhead).%22%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_long2short_note%20%3D%20%22Disable%20**Long2Short**%20to%20see%20the%20teacher-only%20baseline.%22%0A%0A%20%20%20%20%23%20Build%20output%20stack%0A%20%20%20%20_out%20%3D%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20_fig1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20_fig2%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22**Optimal%20context%20budget%3A**%20%60%7B_opt_budget%7D%60%20tokens%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22**Maximum%20expected%20reward%3A**%20%60%7B_opt_reward%3A.2%25%7D%60%5Cn%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2B%20_long2short_note%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%20%20%20%20_out%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20**Takeaways**%0A%20%20%20%20-%20K1.5's%20framework%20has%20only%20one%20knob%3A%20how%20many%20tokens%20to%20allow%20the%20model%20to%0A%20%20%20%20%20%20think%20before%20answering.%20Increasing%20%24n%24%20raises%20accuracy%20but%20also%20latency%20and%20compute.%0A%20%20%20%20-%20The%20length%20penalty%20%24%5Calpha%24%20turns%20the%20accuracy%20curve%20into%20a%20budget-aware%0A%20%20%20%20%20%20operating%20point%2C%20letting%20you%20trade%20answer%20quality%20for%20runtime.%0A%20%20%20%20-%20Long2Short%20distillation%20shows%20the%20curve%20shifts%3A%20a%20short%20model%20can%20recover%2060%E2%80%9385%25%0A%20%20%20%20%20%20of%20the%20long-CoT%20quality%20at%201%2F4%20the%20tokens.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A