adaptrack-constrained-decoding

%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%22%2C%0A%23%20%20%20%20%20%22matplotlib%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%0A%20%20%20%20return%20mo%2C%20np%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20AdapTrack%3A%20Constrained%20Decoding%20Without%20Distorting%20Intent%0A%0A%20%20%20%20**Paper**%20%7C%20Yongmin%20Li%2C%20Jia%20Li%2C%20Ge%20Li%2C%20Zhi%20Jin%20(PKU%20%2F%20Tsinghua)%0A%0A%20%20%20%20%23%23%20Story%0A%20%20%20%20Standard%20constrained%20decoding%20fixes%20invalid%20tokens%20by%20zeroing%20their%20probabilities.%0A%20%20%20%20That%20preserves%20validity%2C%20but%20often%20destroys%20what%20the%20model%20*actually*%20wanted%20to%20say.%0A%20%20%20%20AdapTrack%20adds%20backtracking%20so%20the%20sampler%20can%20recover%20the%20highest-probability%20legal%0A%20%20%20%20continuation%20when%20a%20greedy%20constraint%20prune%20makes%20the%20current%20prefix%20look%20hopeless.%0A%0A%20%20%20%20This%20notebook%20shows%20why%20greedy%20constraint%20pruning%20breaks%20probability%20mass%2C%20and%20how%0A%20%20%20%20adaptive%20backtracking%20plus%20rejection%20sampling%20restore%20a%20distribution%20that%20matches%20the%0A%20%20%20%20model's%20conditional%20intent.%20CPU-only.%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Probability-Collapse%20Example%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%20collapse%20%3D%20mo.ui.slider(0%2C%201%2C%200.01%2C%20value%3D0.8%2C%20label%3D%22Removed%20mass%20from%20greedy%20prune%22)%0A%20%20%20%20collapse%0A%20%20%20%20return%20(collapse%2C)%0A%0A%0A%40app.cell%0Adef%20_(collapse%2C%20np)%3A%0A%20%20%20%20_removed%20%3D%20collapse.value%0A%0A%20%20%20%20%23%20Two-stage%20API%20completion%20context%20from%20the%20paper%3A%0A%20%20%20%20%23%20%22Matrix%20rank%22%20context%20in%20PyTorch%202.0%2C%20after%20%22return%20torch.%22%0A%20%20%20%20%23%20Candidate%20API%20continuations%3A%0A%20%20%20%20%23%20%20%20%5Bmatrix_rank%20%3D%20invalid%20in%20the%20example%20scenario%2C%20linalg.matrix_rank%20%3D%20desired%2C%20matrix_power%20%3D%20legal%20alternative%5D%0A%0A%20%20%20%20%23%20Original%20model%20intent%3A%20strongly%20prefers%20%60matrix_rank%60%2C%20then%20%60linalg.matrix_rank%60%2C%20then%20%60matrix_power%60.%0A%20%20%20%20model_names%20%3D%20%5B%22matrix_rank%22%2C%20%22linalg.matrix_rank%22%2C%20%22matrix_power%22%5D%0A%20%20%20%20%23%20model_true%20%3D%20probs%20conditioned%20*without*%20constraint%0A%20%20%20%20_true%20%3D%20np.array(%5B0.5744%2C%200.1921%2C%200.2335%5D)%0A%20%20%20%20_true%20%2F%3D%20_true.sum()%0A%0A%20%20%20%20%23%20Greedy%20constrained%20decoding%20sees%20%60matrix_rank%60%20invalid%20and%20removes%20it.%0A%20%20%20%20_valid_mask%20%3D%20np.array(%5BFalse%2C%20True%2C%20True%5D)%0A%20%20%20%20%23%20Convention%3A%20remaining%20mass%20is%20renormalized%20greedily%3B%20all%20removed%20probability%20flows%20to%0A%20%20%20%20%23%20the%20*highest*%20remaining%20probability%20token%20because%20the%20model%20%22sees%22%20all%20others%20as%20equally%0A%20%20%20%20%23%20improbable%20after%20the%20distraction.%0A%20%20%20%20greedy_valid%20%3D%20np.where(_valid_mask%2C%20_true%2C%200.0)%0A%20%20%20%20if%20greedy_valid.sum()%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20greedy_valid%20%3D%20np.where(_valid_mask%2C%201.0%2C%200.0)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20greedy_valid%20%3D%20greedy_valid%20%2F%20greedy_valid.sum()%0A%20%20%20%20%23%20Paper's%20point%3A%20greedy%20prune%20can%20make%20an%20improbable%20token%20dominate.%0A%20%20%20%20%23%20We%20model%20this%20as%20shifting%20mass%20from%20the%20desired%20candidate%20to%20a%20legal-but-irrelevant%20token.%0A%20%20%20%20_desired_idx%20%3D%201%0A%20%20%20%20_distractor_idx%20%3D%202%0A%20%20%20%20greedy_distorted%20%3D%20greedy_valid.copy()%0A%20%20%20%20greedy_distorted%5B_desired_idx%5D%20*%3D%20(1%20-%20_removed)%0A%20%20%20%20greedy_distorted%5B_distractor_idx%5D%20%2B%3D%20greedy_valid%5B_desired_idx%5D%20*%20_removed%0A%20%20%20%20greedy_distorted%20%2F%3D%20greedy_distorted.sum()%0A%0A%20%20%20%20%23%20AdapTrack%20intent-preserving%20adjustment%3A%20instead%20of%20allowing%20the%20distortion%2C%20it%20backtracks%0A%20%20%20%20%23%20and%20rescales%20so%20the%20conditioned%20distribution%20follows%20the%20original%20likelihood%20ordering%20among%0A%20%20%20%20%23%20valid%20tokens.%0A%20%20%20%20adaptrack_valid%20%3D%20np.where(_valid_mask%2C%20_true%2C%200.0)%0A%20%20%20%20if%20adaptrack_valid.sum()%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20adaptrack_valid%20%3D%20np.where(_valid_mask%2C%201.0%2C%200.0)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20adaptrack_valid%20%3D%20adaptrack_valid%20%2F%20adaptrack_valid.sum()%0A%20%20%20%20return%20adaptrack_valid%2C%20greedy_distorted%2C%20model_names%0A%0A%0A%40app.cell%0Adef%20_(adaptrack_valid%2C%20greedy_distorted%2C%20model_names%2C%20np)%3A%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(7%2C%203.2))%0A%20%20%20%20_x%20%3D%20np.arange(len(model_names))%0A%20%20%20%20_width%20%3D%200.22%0A%20%20%20%20_ax.bar(_x%20-%20_width%2C%20_true%2C%20_width%2C%20label%3D%22Model%20intent%20(no%20constraint)%22%2C%20color%3D%22%234e79a7%22)%0A%20%20%20%20_ax.bar(_x%2C%20greedy_distorted%2C%20_width%2C%20label%3D%22Constrained%20decoding%22%2C%20color%3D%22%23e15759%22)%0A%20%20%20%20_ax.bar(_x%20%2B%20_width%2C%20adaptrack_valid%2C%20_width%2C%20label%3D%22AdapTrack%22%2C%20color%3D%22%2359a14f%22)%0A%20%20%20%20_ax.set_xticks(_x)%0A%20%20%20%20_ax.set_xticklabels(model_names%2C%20rotation%3D18)%0A%20%20%20%20_ax.set_ylim(0%2C%201)%0A%20%20%20%20_ax.set_ylabel(%22Probability%22)%0A%20%20%20%20_ax.set_title(%22Greedy%20prune%20vs%20AdapTrack%20on%20API%20completion%20intent%22)%0A%20%20%20%20_ax.legend(frameon%3DFalse)%0A%20%20%20%20_ax.spines%5B%5B%22top%22%2C%20%22right%22%5D%5D.set_visible(False)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%20(plt%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20How%20large%20a%20backtrack%20helps%0A%20%20%20%20The%20paper%20argues%20limited%20backtrack%20distances%20(2%E2%80%938%20tokens)%20capture%20most%20of%20the%20benefit.%0A%20%20%20%20Below%3A%20pass%401%20lift%20over%20constrained%20decoding%20versus%20backtrack%20budget%20on%20two%20synthetic%20API%20settings.%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%20backtrack%20%3D%20mo.ui.slider(0%2C%2016%2C%201%2C%20value%3D8%2C%20label%3D%22Backtrack%20distance%20budget%22)%0A%20%20%20%20backtrack%0A%20%20%20%20return%20(backtrack%2C)%0A%0A%0A%40app.cell%0Adef%20_(backtrack%2C%20np)%3A%0A%20%20%20%20_steps%20%3D%20np.arange(0%2C%2017)%0A%20%20%20%20%23%20Shaped%20curves%20approximating%20curves%20from%20paper%20Figure%204%20style%20behavior%3A%0A%20%20%20%20%23%20For%20constrained%20decoding%2C%20benefit%20saturates%20quickly.%0A%20%20%20%20_tf1_lift%20%3D%201%20-%20np.exp(-_steps%20%2F%203.2)%0A%20%20%20%20_tf2_lift%20%3D%201%20-%20np.exp(-_steps%20%2F%204.8)%0A%20%20%20%20_selected%20%3D%20backtrack.value%0A%20%20%20%20if%20_selected%20%3C%200%3A%0A%20%20%20%20%20%20%20%20_selected%20%3D%200%0A%20%20%20%20_budget%20%3D%20np.clip(_steps%2C%200%2C%20_selected)%20%2F%20np.maximum(np.clip(_steps%2C%200%2C%20_selected)%2C%201)%0A%20%20%20%20_current_tf1%20%3D%20float(_tf1_lift%5Bmin(_selected%2C%20_tf1_lift.size%20-%201)%5D)%0A%20%20%20%20_current_tf2%20%3D%20float(_tf2_lift%5Bmin(_selected%2C%20_tf2_lift.size%20-%201)%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(plt)%3A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(7%2C%203.2))%0A%20%20%20%20_ax.plot(_steps%2C%20_tf1_lift%2C%20color%3D%22%234e79a7%22%2C%20linewidth%3D2%2C%20label%3D%22TFv1%20synthetic%20EM%20lift%22)%0A%20%20%20%20_ax.plot(_steps%2C%20_tf2_lift%2C%20color%3D%22%23f28e2b%22%2C%20linewidth%3D2%2C%20label%3D%22TFv2%20synthetic%20EM%20lift%22)%0A%20%20%20%20_ax.axvline(_selected%2C%20color%3D%22%23b07aa1%22%2C%20linestyle%3D%22--%22)%0A%20%20%20%20_ax.fill_between(_steps%5B%3A%20_selected%20%2B%201%5D%2C%200%2C%201%2C%20color%3D%22%23b07aa1%22%2C%20alpha%3D0.08)%0A%20%20%20%20_ax.scatter(%5B_selected%5D%2C%20%5B_current_tf1%5D%2C%20color%3D%22%234e79a7%22%2C%20zorder%3D5)%0A%20%20%20%20_ax.scatter(%5B_selected%5D%2C%20%5B_current_tf2%5D%2C%20color%3D%22%23f28e2b%22%2C%20zorder%3D5)%0A%20%20%20%20_ax.annotate(%0A%20%20%20%20%20%20%20%20f%22lift%E2%89%88%7B_current_tf1%3A.2f%7D%22%2C%0A%20%20%20%20%20%20%20%20(_selected%2C%20_current_tf1)%2C%0A%20%20%20%20%20%20%20%20textcoords%3D%22offset%20points%22%2C%0A%20%20%20%20%20%20%20%20xytext%3D(10%2C%206)%2C%0A%20%20%20%20%20%20%20%20color%3D%22%234e79a7%22%2C%0A%20%20%20%20)%0A%20%20%20%20_ax.annotate(%0A%20%20%20%20%20%20%20%20f%22lift%E2%89%88%7B_current_tf2%3A.2f%7D%22%2C%0A%20%20%20%20%20%20%20%20(_selected%2C%20_current_tf2)%2C%0A%20%20%20%20%20%20%20%20textcoords%3D%22offset%20points%22%2C%0A%20%20%20%20%20%20%20%20xytext%3D(10%2C%20-14)%2C%0A%20%20%20%20%20%20%20%20color%3D%22%23f28e2b%22%2C%0A%20%20%20%20)%0A%20%20%20%20_ax.set_xlabel(%22Backtrack%20distance%20(tokens)%22)%0A%20%20%20%20_ax.set_ylabel(%22Relative%20EM%20improvement%20over%20constrained%20decoding%22)%0A%20%20%20%20_ax.set_title(%22Benefit%20of%20backtrack%20budget%20saturates%20quickly%22)%0A%20%20%20%20_ax.set_ylim(0%2C%201.05)%0A%20%20%20%20_ax.spines%5B%5B%22top%22%2C%20%22right%22%5D%5D.set_visible(False)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%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%23%23%20Adaptive%20rejection-sample%20backtrack%0A%20%20%20%20AdapTrack%20maintains%3A%0A%20%20%20%20-%20%24Q%5Bx%5D%24%20%3D%20estimated%20probability%20of%20completing%20valid%20output%20after%20prefix%20%24x%24.%0A%20%20%20%20-%20%24N%5Bx%5D%24%20%3D%20sampled%20next%20token%20from%20the%20model%20distribution%20weighted%20by%20current%20%24Q%5Bx%5Cmathbin%7B%3A%3A%7Du%5D%24.%0A%20%20%20%20On%20failure%20to%20complete%2C%20%24Q%24%20is%20recomputed%20bottom-up%20and%20each%20sampled%20edge%20%24N%5Bx%5D%24%20is%20tested%3A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Ctext%7Baccept%20with%20%7D%20%5Cmin%5CBigl(1%2C%5C%3B%20%5Cfrac%7BP_%7BQ'%7D(N%5Bx%5D%5Cmid%20x)%7D%7BP_Q(N%5Bx%5D%5Cmid%20x)%7D%5CBigr)%0A%20%20%20%20%24%24%0A%20%20%20%20Rejection%20means%20backtrack%20and%20resample%20%24x%24%20without%20the%20rejected%20token.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20np)%3A%0A%20%20%20%20def%20_prefix_probability(_seq%2C%20_base_probs%2C%20_valid%2C%20start%3D0)%3A%0A%20%20%20%20%20%20%20%20%23%20Crude%20prefix%20probability%20estimate%20for%20toy%20trace%20visualization.%0A%20%20%20%20%20%20%20%20_probs%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20_i%20in%20range(len(_seq))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20_i%20%3C%20start%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_probs.append(1.0)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%20%20%20%20%20%20%20%20%20%20%20%20_p%20%3D%20float(_base_probs%5B_seq%5B_i%5D%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20_valid.get(_seq%5B_i%5D%2C%20True)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_p%20%3D%200.0%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20_p%20%3D%3D%200.0%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_probs.append(0.0)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_probs.append(_p)%0A%20%20%20%20%20%20%20%20return%20_probs%0A%0A%20%20%20%20%23%20Initial%20explorer%20path%3A%20follows%20the%20greedy%20intention%20token%20%22matrix_rank%22%0A%20%20%20%20_tokens%20%3D%20%5B%22matrix_rank%22%2C%20%22linalg.matrix_rank%22%2C%20%22matrix_power%22%2C%20%22linalg.svd%22%5D%0A%20%20%20%20_base%20%3D%20np.array(%5B0.5744%2C%200.1921%2C%200.2335%2C%200.0500%5D)%0A%20%20%20%20_base%20%2F%3D%20_base.sum()%0A%20%20%20%20_valid%20%3D%20%7B%22matrix_rank%22%3A%20False%2C%20%22linalg.matrix_rank%22%3A%20True%2C%20%22matrix_power%22%3A%20True%2C%20%22linalg.svd%22%3A%20True%7D%0A%0A%20%20%20%20selected%20%3D%20mo.ui.dropdown(%0A%20%20%20%20%20%20%20%20options%3D_tokens%2C%20value%3D%22matrix_rank%22%2C%20label%3D%22Chosen%20after%20backtrack%22%0A%20%20%20%20)%0A%20%20%20%20selected%0A%20%20%20%20return%20(selected%2C)%0A%0A%0A%40app.cell%0Adef%20_(np%2C%20plt%2C%20selected)%3A%0A%20%20%20%20_reject_idx%20%3D%20_tokens.index(%22matrix_rank%22)%0A%20%20%20%20_chosen_idx%20%3D%20_tokens.index(selected.value)%0A%20%20%20%20_chosen%20%3D%20selected.value%0A%20%20%20%20_reject%20%3D%20_tokens%5B_reject_idx%5D%0A%20%20%20%20_p_reject%20%3D%20float(_base%5B_reject_idx%5D)%20*%20(0.0%20if%20not%20_valid%5B_reject%5D%20else%201.0)%0A%20%20%20%20_p_chosen%20%3D%20float(_base%5B_chosen_idx%5D)%0A%0A%20%20%20%20if%20_p_reject%20%3E%200%20and%20_p_chosen%20%3C%20_p_reject%3A%0A%20%20%20%20%20%20%20%20_acceptance%20%3D%20_p_chosen%20%2F%20_p_reject%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_acceptance%20%3D%201.0%0A%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(6.6%2C%202.8))%0A%20%20%20%20_status%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%22Rejected%22%20if%20_i%20%3D%3D%20_reject_idx%20and%20not%20_valid%5B_reject%5D%20else%20(%22Chosen%22%20if%20_i%20%3D%3D%20_chosen_idx%20else%20%22Candidate%22)%0A%20%20%20%20%20%20%20%20for%20_i%20in%20range(len(_tokens))%0A%20%20%20%20%5D%0A%20%20%20%20_colors%20%3D%20%5B%22%23e15759%22%20if%20s%20%3D%3D%20%22Rejected%22%20else%20(%22%2359a14f%22%20if%20s%20%3D%3D%20%22Chosen%22%20else%20%22%23a0cbe8%22)%20for%20s%20in%20_status%5D%0A%20%20%20%20_ax.bar(np.arange(len(_tokens))%2C%20_base%2C%20color%3D_colors)%0A%20%20%20%20_ax.set_xticks(np.arange(len(_tokens)))%0A%20%20%20%20_ax.set_xticklabels(_tokens%2C%20rotation%3D24)%0A%20%20%20%20_ax.set_ylabel(%22Model%20probability%22)%0A%20%20%20%20_ax.set_title(f%22Acceptance%20probability%20for%20chosen%20successor%3A%20%7B_acceptance%3A.2f%7D%22)%0A%20%20%20%20_ax.spines%5B%5B%22top%22%2C%20%22right%22%5D%5D.set_visible(False)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%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%23%23%20Takeaways%20for%20CPU-runnable%20evaluation%0A%20%20%20%20-%20An%20offline%20simulator%20for%20AdapTrack%20can%20quantify%20backtrack%20budget%20and%20constraint%20hardness.%0A%20%20%20%20-%20Adding%20a%20verifier%20on%20top%20of%20a%20small%20model%20mirrors%20this%3A%20small%20models%20gain%20more%20from%20intent%20preservation.%0A%20%20%20%20-%20Sampler-side%20intervention%20does%20not%20require%20retraining%3B%20it%20only%20needs%20a%20verify-oracle.%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%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20References%0A%20%20%20%20-%20AdapTrack%3A%20Constrained%20Decoding%20without%20Distorting%20LLM's%20Output%20Intent.%20Li%2C%20Li%2C%20Li%2C%20Jin.%202510.17376.%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