Nikki Lijing Kuang

Author directory

2026

Large language models (LLMs) tool-use policies select actions from a finite candidate set at each reasoning step. We observe that this action-selection interface is already a linear contextual scorer: under a restricted output vocabulary, the LM head computes per-action scores from the last-layer hidden state, exactly the structure assumed by linear contextual bandits. Existing online adaptation methods for LLMs operate indirectly through prompts or retrieved memories, leaving this decision layer fixed and providing no explicit mechanism for action-level credit assignment or uncertainty-aware exploration. We propose OLIVIA, a lightweight online adaptation layer that replaces the frozen action-selection head with a per-action contextual linear bandit, selects actions via LinUCB exploration on the LLM’s hidden state and updates only its per-action sufficient statistics from step-level feedback via a rank-one Sherman–Morrison rule. Without backbone fine-tuning and additional LLM calls per step, OLIVIA adds negligible overhead relative to LLM inference. Across four tool-use benchmarks ToolBench, TaskBench, TaskBench-MM, BFCL and two backbones Qwen3-4B, Mistral-7B, OLIVIA outperforms prompt-based online adaptation and contextual-bandit variants.
Interactive text-to-SQL agents solve database tasks through multi-turn interactions involving schema exploration, query execution, feedback interpretation, and decision revision. Long-term memory helps agents reuse past experiences, but existing retrieval methods remain limited. Static methods use fixed similarity heuristics, whereas dynamic methods often learn from sparse terminal outcomes and retrieve at a single decision horizon. This is insufficient when memory usefulness changes across interaction stages, since memories useful for initial planning may differ from those needed for local, state-conditioned execution. We propose MERIT, a dynamic multi-horizon memory retrieval framework. MERIT maintains episode-level memory for global strategic guidance and turn-level memory for local decision support. Both levels use learned retrieval policies optimized with reinforcement learning. To train turn-level retrieval despite limited intermediate supervision, MERIT uses a lightweight Process Reward Model to provide dense proxy rewards for local memory selection. On BIRD-Interact, MERIT achieves higher success rates and fewer interaction turns than no-memory, static-retrieval, and dynamic-retrieval baselines. Transfer results on Spider2-Snow further show that the learned policies transfer without degradation, staying above all memory baselines without benchmark-specific tuning. These results suggest that multi-horizon retrieval improves experience reuse in interactive text-to-SQL agents. Our code is available at https://github.com/YiboWANG214/MERIT.