NeurIPS 2026 论文列表

Conference on Neural Information Processing Systems

录用类型

6,437篇论文

按官方录用论文列表统计,可能与会议公布的数字略有出入。

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

浏览全部论文

KVFocus: A Perturbation-Theoretic Token-Risk Score for Selective KV Cache Reuse in RAG

Shizhuo Zhang, Nuowen Kan, Chenglin Li, Rui-Xiao Zhang, Yifan Zhao, Yuanpeng He, Huixin Zhang, FAN He, Hang Xu, Hao Zhang, Wenrui Dai, Junni Zou, Hongkai Xiong

展开摘要

In retrieval-augmented generation (RAG) scenarios, the prefill delay of the large language model (LLM) long inputs is the dominant component of time-to-first-token (TTFT). Reusing KV caches across different LLM inputs reduces this cost, but the resulting cross-context mismatch on reused segments has a significant degradation on answer quality via propagating through prefill attention into decode-time logits. Existing selective-recomputation methods mitigate this issue by recomputing tokens chosen by empirical importance heuristics, which overlook the propagation of cache mismatches into answer-side errors, thereby failing to effectively strike the quality-TTFT tradeoff. To address this issue, we propose a perturbation-theoretic KV cache reuse framework, KVFocus, which selects reused KV caches through a perturbation-theoretic token-risk score during the prefill process. Specifically, we first develop a first-order analysis that traces reuse error in three stages: source mismatch on the reused segment, suffix contamination during prefill, and decode-time logit stability. This analysis reveals that the leading-order answer-side error decomposes into a product of source-side V-drift and downstream suffix-to-segment attention concentration. Guided by this multiplicative structure, we define a token-risk score that operationalizes the bound during the inference. A plug-and-play selector for KV cache reuse in existing RAG-based LLM serving is then designed by recomputing the top-$r\%$ tokens directly suppresses the predicted leading-order answer-side error, covering both sides of the pathway that prior one-sided heuristics miss. Empirical results across three 7-8B instruction-tuned LLMs (Mistral, Qwen2.5, Qwen3) and four multi-hop QA benchmarks demonstrate that the proposed KVFocus achieves a superior trade-off between the ansewer quality and the TTFT in comparison to state-of-the-art baselines, with up to $3\times$ TTFT speedup over full prefill at matched answer quality.

Poster
KV-PRM: Efficient Process Reward Modeling via KV-Cache Transfer for Multi-Agent Test-Time Scaling

Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

展开摘要

rocess Reward Models (PRMs) have been proven to be highly effective in guiding test-time scaling (TTS) methods, which significantly boost the capabilities of LLM-based multi-agent systems. However, existing PRMs are *text-based*: they re-encode the entire trajectory text from scratch. In long multi-agent rollouts, the scoring cost, growing quadratically with respect to sequence length $L$, creates a severe computational bottleneck, severely limiting PRMs' application in long-context scenarios. To resolve this, we introduce KV-PRM, a highly efficient process reward model that eliminates the heavy text re-encoding by directly reading the KV cache produced naturally during the LLM's generation phase. By processing a single "verify token" against the pre-existing KV cache, KV-PRM reduces the scoring cost from $O(L^2)$ to $O(L)$. We formally prove that the KV cache contains strictly greater information capacity than text, and is more efficient for downstream reward modeling. Empirically, across the MATH, GSM8K, and AIME benchmarks, KV-PRM matches or strictly outperforms text-PRMs under various TTS methods such as Beam Search, MCTS, and Weighted Voting, with up to a *5,000*$\times$ reduction in scoring FLOPs, a *37*$\times$ reduction in latency, and a *34*$\times$ reduction in per-sequence memory footprint compared to text-based PRMs.

Poster
KV Packet: Recomputation-Free Context-Independent KV Caching for LLMs

Chuangtao Chen, Grace Li Zhang, Xunzhao Yin, Cheng Zhuo, Bing Li, Ulf Schlichtmann

展开摘要

Large Language Models (LLMs) rely heavily on Key-Value (KV) caching to minimize inference latency. However, standard KV caches are context-dependent: reusing a cached document in a new context requires recomputing KV states to account for shifts in attention distribution. Existing solutions such as CacheBlend, EPIC, and SAM-KV mitigate this issue by selectively recomputing a subset of tokens; however, they still incur non-negligible computational overhead (FLOPs) and increased Time-to-First-Token (TTFT) latency. In this paper, we propose KV Packet, a recomputation-free cache reuse framework that treats cached documents as immutable ``packets'' wrapped in light-weight trainable soft-token adapters, which are trained via self-supervised distillation to bridge context discontinuities. Experiments on Llama-3.1 and Qwen2.5 demonstrate that the proposed KV Packet method achieves near-zero FLOPs and lower TTFT than recomputation-based baselines, while retaining F1 scores comparable to those of the full recomputation baseline. Code and experimental results are available at https://anonymous.4open.science/r/kv_packet_anonymous-46FC/.

Poster
KNN Implementation Details Can Dramatically Change Performance: An Example from Cover Trees

Amol Khanna, Edward Raff

展开摘要

Nearest neighbor search algorithms are widely used and have a long history, but we find that key implementation details are seldom reviewed in articles about them. These details can materially impact performance, and thus change the conclusions one makes about a given algorithm's superiority. In this article, we highlight these impacts by reviewing the Cover Tree as an example, producing RustKNN, a unified Rust implementation of three construction variants of the Cover Tree with three query modes. Our analysis reveals: (1) the optimal query strategy depends on $k$ and on the dataset — single-tree search is fastest at $k{=}1$ on three of our four largest datasets, while dual-tree search is fastest at $k{\ge}5$ on the three largest — and choosing the wrong query mode can cost up to a $7.5\times$ factor in runtime; (2) the cover tree base parameter, set to 2 in theory and 1.3 in practice, is not a safe default for dual-tree search: 1.3 is optimal on only one of seven datasets, and on covtype it is $1.6\times$ slower than the optimum; (3) among five micro-optimizations, only early-exit distance computation has a consistent effect, and that effect grows with dimension (none at $d{=}32$, $1.31\times$ at $d{=}54$, and $1.49\times$ at $d{=}784$); and (4) all-nearest-neighbor and held-out benchmarks yield different implementation rankings. All main timing tables report medians over 20 independent allocations with bootstrap 95\% confidence intervals. RustKNN is open source at https://github.com/CrowdStrike/RustKNN.

Poster
KL for a KL: On-Policy Distillation with Control Variate Baseline

Minjae Oh, Sangjun Song, Gyubin Choi, Yunho Choi, Yohan Jo

展开摘要

On-Policy Distillation (OPD) has emerged as a dominant post-training paradigm for large language models, especially for reasoning domains. However, OPD remains unstable in practice due to the high gradient variance of its single-sample Monte Carlo estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy Distillation with a control variate baseline), which casts OPD as policy-gradient RL and stabilizes it by introducing a control variate baseline—canonically a value function—from the RL literature. We show that the OPD value function admits a closed form as the per-token negative reverse KL divergence between the student and the teacher, available directly from the already-computed forward pass with no additional critic or inference. Existing stabilization methods either compute the full token-level reverse KL over the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample estimator, subtracting the value function as a detached baseline to keep the gradient unbiased while reducing variance. Furthermore, we show that a top-k approximation of the baseline further lowers cost without compromising performance. Across mathematical and scientific reasoning benchmarks, vOPD consistently outperforms vanilla OPD and matches the most expensive full-vocabulary baseline, offering an efficient stabilization of On-Policy Distillation through principled RL variance reduction.

Poster
KINDER: Kernel-based Independence for Fair Representation Learning via Prototype-space Erasure

Abtin Mogharabin, Jiaee Cheong, Alp Toykan Kaplan, Sinan Kalkan

展开摘要

Concept erasure is a prominent approach to achieving fairness in machine learning through the removal of sensitive attributes from learned representations. Prior concept erasure methods typically define debiasing indirectly through the failure of a chosen adversary, probe, or fairness penalty, making the target of erasure dependent on a particular decoder family or optimization setup. Rather than depending on a particular architecture, task type, or training paradigm, we introduce KINDER, which aims to directly remove sensitive-attribute information from a model's intermediate representations and is compatible with any deep model that contains an intermediate feature space. KINDER operates in a random Fourier feature space, where it estimates a sensitive subspace from protected group prototypes and projects representations onto its orthogonal complement, thereby forming an explicit representation cleaning mechanism. We conduct extensive experiments across diverse settings and show the effectiveness of KINDER on (i) unimodal and multimodal datasets, (ii) supervised and self-supervised settings, (iii) classification, regression, and image segmentation tasks, and (iv) diverse data modalities, including visual, textual, and tabular data.

Poster
KG-Guard: Graph-Based Hallucination Detection for Knowledge Base Question Answering

Albert Sawczyn, Piotr Bielak, Tomasz Kajdanowicz

展开摘要

Large language models (LLMs) are increasingly used for knowledge base question answering (KBQA), where answering requires selecting entities from a question-specific knowledge-graph subgraph. Yet LLMs are known to hallucinate across tasks, and KBQA is no exception: even when we provide a graph as the knowledge source, the model may rely on parametric knowledge instead of graph evidence or perform invalid reasoning over the given relations. Such hallucinated answer nodes can limit the practical deployment of KBQA systems, especially in high-stakes domains such as healthcare. We formulate hallucination detection in KBQA as an answer-node classification problem and propose a lightweight graph-based framework that treats the answering LLM as a black box. KG-Guard represents each KBQA instance as an augmented graph. It initializes node features with semantic representations of KG entities, marks topic entities and LLM-proposed answer nodes with learned vectors, and connect a virtual question node to the topic entities. A graph encoder then produces verification-oriented node representations, and a small MLP classifies each proposed answer node using its graph representation together with the question embedding. Experiments on WebQSP, ComplexWebQuestions, and PUGG show that our detector achieves the highest F1 on all three benchmarks ($82.0$, $87.4$, and $84.3$), outperforming LLM-judge and sampling-based baselines, while having $\sim305\times$ fewer parameters than the reference approaches. Beyond detection, the node-level feedback is actionable: when flagged answers are fed back to the KBQA system for iterative refinement, downstream KBQA F1 improves by $13.0$-$14.5$ points and Exact Match by $16.9$-$17.6$ points.

Poster
KDASO: Knowledge-Data Dual-Driven Automated Skill Optimization for Liability Adjudication Task

Zi-Jian Cheng, Weiming Wu, Jie Meng, Zhen Peng, Shan Huang, Tracy Li, Guobin Wu, Lan-Zhe Guo

展开摘要

Liability adjudication tasks (e.g., legal charge determination), require analyzing event descriptions against domain-specific skills to determine whether a party should bear liability. While Large Language Models (LLMs) have shown promise in this area by leveraging executable skills for interpretable reasoning, existing approaches face three critical limitations: **skill libraries become stale as real-world scenarios dynamically evolve**, **skill descriptions often suffer from a semantic execution gap where LLMs deviate from intended expert logic**, and **growing skill collections exceed context budgets causing critical skill omission**. To address these challenges, we propose **K**nowledge-**D**ata Dual-Driven **A**utomated **S**kill **O**ptimization (KDASO), which continuously evolves a skill library through neighborhood-conditioned induction from recent adjudication data, execution-aligned refinement based on LLM trajectory analysis, and context-budgeted routing for scalable skill selection. Experiments across ride-hailing adjudication, legal liability determination, and logical reasoning judgment demonstrate that KDASO significantly outperforms current general-purpose and domain-specific baselines. KDASO further achieves an 8\% improvement when transferring skills to concept-similar scenarios and maintains consistent performance gains through continuous optimization, underscoring the high transferability and stability of generated skills.

Poster
K12-KGraph: A Curriculum-Aligned Knowledge Graph for Benchmarking and Training Educational LLMs

Hao Liang, Qihan Lin, Mingrui Chen, Hengyi Feng, Zhaoyang Han, Xiaochen Ma, Zhen H Wong, Meiyi Qiang, Linzhuang Sun, Wentao Zhang

展开摘要

Large language models (LLMs) are increasingly deployed in K–12 education, yet existing benchmarks such as C-Eval, CMMLU, GaokaoBench, and EduEval measure only whether a model can answer an exam question, i.e., factual recall. Effective educational AI further requires curriculum cognition: the structured understanding of how knowledge is organized, including prerequisite chains, concept taxonomies, experiment–concept links, and pedagogical sequencing. Curriculum cognition is neither probed by current benchmarks nor explicitly taught by current instruction-tuning data. To close this gap, we introduce K12-KGraph, a curriculum-aligned knowledge graph extracted from the official People's Education Press textbooks, covering mathematics, physics, chemistry, and biology across primary, middle, and high school, with seven node types (Concept, Skill, Experiment, Exercise, Section, Chapter, Book) and nine relation types spanning taxonomy, prerequisite, association, verification, assessment, location, and order. Building on this single graph, we derive two complementary resources: K12-Bench, a 23,640-question multi-select benchmark across five graph-derived task families (Ground, Prereq, Neighbor, Evidence, and Locate) that jointly probe curriculum cognition; and K12-Train, a KG-guided supervised fine-tuning corpus of ~2,300 QA pairs synthesized from node attributes and edge semantics. Experiments expose a clear gap and a clear remedy: on K12-Bench, even the strongest proprietary model (Gemini-3-Flash) reaches only 57% exact match and the best open-source model (Gemma-4-31B-IT) only 46%, with Prereq and Neighbor being the hardest; yet under a strictly matched 2,300-sample SFT budget on Qwen3-4B-Base and Llama3.1-8B-Base, K12-Train consistently outperforms equally sized subsets of eight mainstream instruction-tuning corpora (OpenHermes, Infinity, UltraChat, WizardLM, DataFlow, LMSYS, SmolTalk, Tulu-3) on both GaokaoBench and EduEval, showing that structural curriculum grounding is remarkably sample-efficient for educational SFT. We release the graph, benchmark, training data, and the full construction pipeline.

Poster
K-PWM: Control-Oriented Structured World Models under Partial Observation

Santosh M Rajkumar, Sriram Narayanan, Samuel E. Otto, Debdipta Goswami

展开摘要

World models enable prediction, planning, and decision-making by learning internal simulators of environment dynamics. Although substantial progress has been made in world models from pixel-based observations, many physical systems are instead observed through vector-based measurements that are noisy and only partially informative of the underlying state. We introduce K-PWM, a Koopman-structured probabilistic world model for prediction and control under partial observation. K-PWM combines a parameterized Koopman-structured state-space model with a nonlinear observation decoder, and performs probabilistic latent-state inference directly through the learned model parameters rather than through a separately trained inference network. We train K-PWM using a generalized expectation conditional maximization (GECM) procedure with principled initialization. Across Gymnasium MuJoCo tasks, K-PWM improves sample efficiency, yields reliable long-horizon predictions, and supports downstream control through model predictive path integral (MPPI) control. These results suggest that structured probabilistic latent dynamics provide a useful route toward data-efficient world modeling and control in partially observable vector-valued settings.

Poster
Justitia: Fair and Efficient Scheduling of Task-parallel LLM Agents with Selective Pampering

Mingyan Yang, Guanjie Wang, Manqi Luo, Yifei Liu, Chen Chen, Han Zhao, Yu Feng, Quan Chen, Minyi Guo

展开摘要

LLM agents, which often comprise parallel inference tasks, are commonly adopted to solve real-world problems. When serving such task-parallel LLM agents in shared GPU servers, the scheduler is expected to attain fast agent completion with guaranteed worst-case performance. For that objective, our insight is to selectively pampering agents based on their completion order under idealized fair-sharing. We design Justitia, a fair and also efficient scheduler for task-parallel LLM agents. Noticing that memory is prevalently a bottleneck in LLM serving, Justitia quantifies the true agent cost in a memory-centric manner. It also adopts a light-weight yet accurate method to predict agent costs. Finally, Justitia adopts a virtual-time based fair queuing algorithm to reduce the overall performance with guaranteed worst-case delay. We have implemented Justitia atop vLLM, and experimental results involving diverse agents show that it can substantially enhance the scheduling efficiency with fairness preserved.

Poster
Just Ramp-Up: Debiasing Regression-based Estimator for A/B Tests under Network Interference

Qianyi Chen, Bo Li

展开摘要

Network interference complicates A/B testing on online platforms, such as social networks and marketplaces, where causal methods based on single experiments often suffer from significant bias due to complex interference patterns. This paper demonstrates the statistical benefits of merging data from multiple experiments with varying treatment proportions. Sequential experimentation with increasing traffic, or ramp-up, is widely used in tech companies for risk management and cost control. Beyond operational benefits, we show that regression-based estimators trained on merged data achieve substantial bias reduction, even under simple randomization schemes and regression models. We focus on the global average treatment effect (GATE), a key estimand in the tech industry, and consider a general interference pattern that extends beyond the 1-hop setting. We present an exact bias–variance analysis of the linear regression estimator and show that, in practical settings, the bias term typically dominates. In addition, we characterize the substantial bias reduction achieved by our proposed approach, which merges experimental data collected at different ramp-up stages to improve the training of the regression function. We also offer an intuitive explanation for this reduction and highlight the synergy between cluster-level randomization and our approach. Furthermore, we consider a refined estimator based on graph neural networks (GNN). Extensive simulations across novel and challenging scenarios confirm that our methodology significantly improves the accuracy of regression-based estimators.

Poster
Jump Start Your Policy Learning with Lessons from 145,000 Training Runs

Nabil Omi, Chung Yik Edward Yeung, Eric Bae, Siddhartha Sen, Ali Farhadi

展开摘要

Reliable progress in offline policy learning depends on a small set of methodological practices, including careful reporting, well-tuned baselines, and evaluation across diverse conditions. Prior work has noted that results can be sensitive to reporting choices, hyperparameter tuning, and dataset properties independently, but these sources of variability have not been systematically investigated at the scale needed to understand how they shape conclusions. To address this gap, we present a large-scale empirical study of offline reinforcement and imitation learning, training over 145,000 policies across 114 datasets. At this scale, no algorithm dominates: aggregate performance across top methods is often close, but the leaders differ substantially across environments. We find that proper hyperparameter tuning frequently reshuffles perceived algorithm rankings; and simple baselines, including behavior cloning, are often stronger than is commonly assumed after extensive tuning. We also study hyperparameter transfer and sensitivity across environments, identifying a simple strategy that generalizes well. From these analyses we distill practical recommendations, and release JumpStart: a resource suite of trained models, per-model hyperparameter and reward data, strong baselines across all environments, and a website to make retrieval and analysis trivial. Together, these resources aim to make offline policy learning research more reliable and to open new directions for work beyond the scope of this study.

Poster
JudgmentBench: Comparing Rubric and Preference Evaluation for Quality Assessment

Russell Yang, Ruishi Chen, Pierce Kelaita, Riya Ranjan, Sibo Ma, Charles Dickens, Matthew Guillod, Megan Ma, Julian Nyarko

展开摘要

Two methodologies dominate current practices of benchmarking: rubric-based scoring evaluates items against predefined criteria, whereas comparative judgment elicits pairwise preferences between outputs. Although both methodologies are widely used, the choice between them is rarely justified. We release JudgmentBench, a benchmark of 30 real-world legal tasks, paired with 1,539 rubric scores and 1,530 pairwise preference judgments collected from practicing attorneys—including at major U.S. law firms—with substantial experience. The annotations constitute the first publicly available dataset in a high-expertise domain in which both supervision signals are elicited from the same experts on the same items. Using LLM-generated outputs at three constructed quality levels, we provide an initial empirical comparison: comparative judgments recover the intended quality ordering substantially better than rubrics (mean Spearman's rank correlation of 0.908 vs. 0.150, $\widehat{D}=0.758$ [0.494, 1.021]) while requiring less than half the annotation time. The patterns hold for human annotators and LLM autograders. Beyond this initial comparison, the paired structure of the dataset supports a broader research agenda on how expert judgment should be elicited, aggregated, and used as supervision in domains without verifiable ground truth.

Poster
JuICE: A Benchmark for Evaluating LLM-Judge in Identifying Cultural Errors

Jiho Jin, Junho Myung, Juhyun Oh, Junyeong Park, Rifki A Putri, Sunipa Dev, Vinodkumar Prabhakaran, Alice Oh

展开摘要

As large language models (LLMs) are increasingly deployed to users around the world, they are integrated into everyday tasks across diverse cultural contexts, from drafting personal communications to brainstorming creative ideas. These tasks are inherently cultural: they require contextual appropriateness, symbolic resonance, and tacit cultural expectations that native speakers draw on instinctively, meaning that a response can be factually plausible yet unmistakably wrong to a local reader. Existing cultural benchmarks have treated culture as a flat set of facts via fact verification or norm entailment methods, and have adopted LLM-as-a-Judge without examining whether they can capture such thick cultural errors. To address this gap, we present JuICE (Benchmark for LLM-Judge in Identifying Cultural Errors), a multilingual dataset of 7,470 span-level annotations of cultural and linguistic errors, collected from native speakers in long-form LLM responses. It covers 1,050 query-response pairs from four countries (United States, South Korea, Indonesia, and Bangladesh), in both English and their countries' main languages. Using JuICE, we find that even the strongest LLM-judge achieves only an F1 of 0.52 against human-annotated errors, and that this gap is systematic. LLM-judges reliably detect surface-level linguistic and factual mistakes but consistently miss thick cultural errors that local residents readily identify. Our findings suggest that robust cultural evaluation must move beyond surface-level detection toward frameworks that account for the depth and situatedness of cultural meaning.

Poster
Jointly Robust Fairness: Overcoming Simultaneous Label and Attribute Noise

Gaurav Jain

展开摘要

Fairness in machine learning typically relies on the assumption of clean training data. However, in high-stakes applications like healthcare and criminal justice, privacy mechanisms and proxy variables frequently introduce simultaneous noise into both sensitive attributes and target labels. While existing methods address label or attribute noise in isolation, they fail under this "double noise" regime, often inadvertently amplifying bias. In this paper, we propose $\textbf{Jointly Robust Fairness (JRF)}$, a unified framework that guarantees fairness under simultaneous data corruption. JRF integrates a Forward Loss correction to handle label noise with a novel Matrix-Inverse Robust Loss (MIRL) that algebraically recovers the true demographic parity gap from noisy attribute observations. We provide rigorous theoretical guarantees for our estimator, including a finite-sample concentration bound demonstrating that the sample complexity scales quadratically with the inverse of the noise transition determinant. Furthermore, our framework naturally generalizes to both symmetric and highly asymmetric noise distributions. Extensive experiments on the Adult, Bank, and COMPAS datasets demonstrate that JRF reduces fairness violations by over 91\% in high-noise regimes ($\rho \ge 0.3$) compared to standard robust baselines, maintaining strict fairness without significant degradation in predictive utility.

Poster
Jointly Reinforcing Diversity and Quality in Language Model Generations

Tianjian Li, Yiming Zhang, Ping Yu, Swarnadeep Saha, Daniel Khashabi, Jason Weston, Jack Lanchantin, Tianlu Wang

展开摘要

Post-training of Large Language Models (LMs) often prioritizes accuracy and helpfulness at the expense of diversity. This creates a tension: while post-training improves response quality, it also sharpens output distributions and reduces the range of ideas, limiting the usefulness of LMs in creative and exploratory tasks such as brainstorming, storytelling, or problem solving. We address this challenge with Diversity-Aware Reinforcement Learning (DARLING), a framework that jointly optimizes for response quality and semantic diversity. At its core, DARLING introduces a learned partition function to measure diversity beyond surface-level lexical variations. This diversity signal is then combined with a quality reward during online reinforcement learning, encouraging models to generate outputs that are both high-quality and distinct. Experiments across multiple model families and sizes show that DARLING generalizes to two regimes: non-verifiable tasks (instruction following and creative writing) and verifiable tasks (competition math). On five benchmarks in the first setting, DARLING consistently outperforms quality-only RL baselines, producing outputs that are simultaneously of higher quality and novelty. In the second setting, DARLING achieves higher pass@1 (solution quality) and pass@k (solution variety). Most strikingly, explicitly optimizing for diversity catalyzes exploration in online RL, which manifests itself as higher-quality responses.

Poster
Joint Treatment Effect Estimation from Incomplete Healthcare Data: Temporal Causal Normalizing Flows with LLM-driven Evolutionary MNAR Imputation

Olivia Jullian Parra, Sara Zoccheddu, David C Cerezo, Tom Forzy, Franziska S Ulrich, William Sutcliffe, Jakob M Burgstaller, Oliver Senn, Patrick Owen, Nicola Serra

展开摘要

Target trial emulation (TTE) provides a framework for answering causal questions using observational data when randomized controlled trials (RCTs) are infeasible. However, standard methods for treatment effect estimation have been developed in isolation, failing to jointly address the compounding challenges inherent to analyses of observational data such as electronic health records (EHRs). In particular, these challenges include time-varying confounding and missing-not-at-random (MNAR) missingness reaching 50\%-80\% for critical biomarkers. To address this gap, we propose a two-stage pipeline that jointly handles MNAR missingness and causal structure across time. CausalFlow-T, a Directed Acyclic Graph (DAG)-constrained normalizing flow with Long Short-Term Memory (LSTM)-encoded patient history, performs exact invertible counterfactual inference, eliminating the approximation errors and confounding biases where existing variational and adversarial methods fail silently. Ablations on four synthetic and one semi-synthetic dataset with known counterfactuals confirm that its two core design choices address strictly non-overlapping failure modes (DAG constraints for confounding separation, exact inference for structural propagation) with neither compensating for the absence of the other. To handle the incomplete data CausalFlow-T receives as input, we propose an LLM-driven evolutionary imputer and evaluate it with three LLM backends, including two open-source models. Across 30\%-80\% MNAR missingness, the imputer achieves the best pooled rank across biomarker and causal metrics, leading on point-wise accuracy and temporal extrapolation while maintaining average treatment effect (ATE) recovery where statistical methods progressively degrade. Applied to a cohort of adults with type 2 diabetes in Swiss primary care initiating a GLP-1 receptor agonist or SGLT-2 inhibitor, the pipeline recovers a per-protocol weight-loss difference of $-0.98$ kg [$95\%$ CI $-1.01$, $-0.96$] favoring GLP-1 receptor agonists, consistent with RCT evidence and estimated directly from realistically incomplete real-world EHR data.

Poster
Joint Sequence--Vocabulary Selection for Efficient LLM Distillation

Xueli Geng, Weicheng Zhao, Xutong Mu, Tianrui Wei, Yanbiao Ma, Yulong Shen

展开摘要

Selective knowledge distillation has become increasingly important for compressing large language models, where dense logit-based knowledge transfer over long sequences and large vocabularies incurs substantial computational overhead. Distillation supervision is naturally organized over a sequence–vocabulary matrix, with its highest-information components often concentrated on a small subset of specific sequence–vocabulary pairs. Existing methods typically select sequence positions or vocabulary candidates independently, relying on coarse-grained one-dimensional heuristics that weaken the precision and effectiveness of knowledge transfer. To address this, we propose $\textbf{J}$oint $\textbf{S}$equence–$\textbf{V}$ocabulary $\textbf{D}$istillation ($\textbf{JSVD}$), a plug-and-play selective distillation framework that identifies and focuses supervision on informative sequence–vocabulary pairs. For each training example, JSVD constructs bidirectional teacher--student candidate supports, scores sequence--vocabulary pairs using prediction discrepancies, and progressively focuses distillation on unresolved informative pairs during training. Extensive experiments across diverse distillation objectives, model pairs, and benchmarks show that JSVD consistently improves downstream performance while reducing distillation overhead and accelerating convergence.

Poster
Joint Optimization of Tool Creation and Use for Large Language Model Agents

Zhi Rui Tam, Chieh-Yen Lin, Yun-Nung (Vivian) Chen, Shao-Hua Sun, Hung-yi Lee

展开摘要

Tool-augmented language models are bounded by the APIs humans bothered to write; existing tool-creation systems patch this by prompting a frozen LLM at inference time, leaving the model that writes a tool decoupled from the one that uses it, with no signal that the schemas it produces are schemas it can invoke. We propose **SMITH** (Schema-grounded Multi-task Iterative Tool Honing), a reinforcement learning framework that jointly trains tool creation and tool use inside a single policy. Each rollout is either a **build** task (write a tool from a few examples) or a **use** task (invoke a pooled tool on a held-out question). Three separate reward axes catch schema, code, and outcome failures independently, so each failure mode contributes its own gradient. A 4B Qwen3 trained with SMITH on 13 procedural reasoning tasks with exact verifiers reaches $79.8$ macro-average accuracy on held-out tasks, the best across all evaluated methods and ahead of an untrained 30B-A3B tool-writer. It also reaches $40.4$ on TabMWP-Hard and $42.6$ on out-of-domain GQA ($+7.6$ over the best same-backbone inference-time baseline), without any visual or tabular training data. When invoked by a frozen 350M student, tools written by our 4B match those produced by a writer an order of magnitude larger. The same recipe also lifts Qwen3-8B and Granite-3.3-8B without modification.

Poster