Data attribution, which traces a model's prediction back to specific training data, is an important tool for interpreting sophisticated AI models. The widely used TRAK algorithm addresses this challenge by first approximating the underlying model with a kernel machine and then leveraging techniques developed for approximating the leave-one-out (ALO) risk. Despite its strong empirical performance, the theoretical conditions under which TRAK approximations are accurate, as well as the regimes in which they break down, remain largely unexplored. In this paper, we provide a theoretical analysis of the TRAK algorithm, characterizing its performance and quantifying the errors introduced by the approximations on which the method relies. We show that although the approximations can incur significant errors, TRAK preserves the separation between highly influential and weakly influential data points. We corroborate our theoretical results through extensive simulations and empirical studies.
We study fixed-policy evaluation for finite Markov chains that may be reducible and periodic. Classical evaluation methods with gain and bias decomposition are not always diagnostic: the gain records only invariant Ces\`aro averages, while persistent phase-dependent behavior is absorbed into the bias together with genuinely transient effects. We identify the real peripheral invariant subspace $\mathcal{K}(P)$ of the transition matrix $P$ as the source of this ambiguity. Quotienting by $\mathcal{K}(P)$ is the minimal exact quotient that removes all non-decaying modes and makes the remaining dynamics strictly stable. After choosing a gauge projection $\Pi$ with kernel $\mathcal{K}(P)$, the reward admits a unique decomposition $r = g_\Pi^\star + (I-P)v_\Pi^\star$, where $g_\Pi^\star$ is a persistent regime profile and $v_\Pi^\star$ is a gauge-fixed transient component. An exact comparison with classical normalized gain and bias shows that the new pair reallocates the same information into a persistent profile and a gauge-fixed representative of the transient quotient. This decomposition reconstructs finite-horizon returns, recovers statewise average reward, admits a transient-cost interpretation, and yields a stable estimator under a generative model.
Bayesian neural networks promise calibrated uncertainty but require $O(mn)$ parameters for standard mean-field Gaussian posteriors. We argue this cost is often unnecessary, particularly when weight matrices exhibit fast singular value decay. By parameterizing weights as $W = AB^{\top}$ with $A \in \mathbb{R}^{m \times r}$, $B \in \mathbb{R}^{n \times r}$, we induce a posterior that is \emph{singular} with respect to the Lebesgue measure, concentrating on the rank-$r$ manifold. This singularity captures structured weight correlations through shared latent factors, geometrically distinct from mean-field's independence assumption. We derive PAC-Bayes generalization bounds whose complexity term scales as $\sqrt{r(m+n)}$ instead of $\sqrt{m n}$, and prove loss bounds that decompose the error into optimization and rank-induced bias using the Eckart-Young-Mirsky theorem. We further adapt recent Gaussian complexity bounds for low-rank deterministic networks to Bayesian predictive means. Empirically, across MLPs, LSTMs, and Transformers on standard benchmarks, our method achieves competitive predictive performance while using up to $33\times$ fewer parameters than 5-member Deep Ensembles. It substantially improves OOD detection and often improves calibration relative to mean-field and perturbation baselines, while Deep Ensembles can still be stronger on in-distribution likelihood-based metrics.
Text-to-image (T2I) models struggle with compositional prompts requiring multiple objects, bound attributes, accurate counts, or spatial relations. Existing benchmarks face limits from fixed detector vocabularies, aggregated scores, and varying human-validation protocols. We introduce SANEval (Spatial, Attribute, and Numeracy Evaluation), an open-vocabulary compositional benchmark with two key contributions: (i) a public dataset of ~5,250 prompts paired with images from six state-of-the-art T2I models, split into algorithmically generated Simple and fully human-written Hard partitions across five compositional tasks, and (ii) a modular evaluation pipeline combining an LLM prompt parser, an open-vocabulary detector with LLM synonym expansion, and scorers for attribute binding, spatial relations, and numeracy. Each scorer emits numerical scores and diagnostic feedback for missing, extra, or mis-bound elements. Validated against human ratings (n = 500), SANEval achieves positive Pearson correlations across all categories (overall r = 0.222), outperforming CompBench++ (-0.016) and GenEval (-0.030). A manual audit of 24,012 detections shows 96.58% precision, with ablation confirming that synonym expansion is critical to pipeline performance. Dataset: [https://huggingface.co/datasets/saneval/saneval-release](https://huggingface.co/datasets/saneval/saneval-release). Pipeline: [https://github.com/nielseni6/saneval](https://github.com/nielseni6/saneval).
Probing studies what information is encoded in a frozen LLM's layer representations by training a lightweight predictor on top of them. Beyond analysis, probes are often used operationally in probe-then-steer pipelines: a learned concept vector is extracted from a probe and injected via additive activation steering by adding it to a layer representation during the forward pass. The effectiveness of this pipeline hinges on estimating concept vectors that are accurate, directionally stable under ablation, and inexpensive to obtain. Motivated by these desiderata, we propose RAPTOR (Ridge-Adaptive Logistic Probe), a simple L2-regularized logistic probe whose validation-tuned ridge strength yields concept vectors from normalized weights. Across extensive experiments on instruction-tuned LLMs and human-written concept datasets, RAPTOR matches or exceeds strong baselines in accuracy while achieving competitive directional stability and substantially lower training cost; these quantitative results are supported by qualitative downstream steering demonstrations. Finally, using the Convex Gaussian Min-max Theorem (CGMT), we provide a mechanistic characterization of ridge logistic regression in an idealized Gaussian teacher-student model in the high-dimensional few-shot regime, explaining how penalty strength mediates probe accuracy and concept-vector stability and yielding structural predictions that qualitatively align with trends observed on real LLM embeddings.
Vision-language models (VLMs) show promise for large-scale affect inference from images, but their alignment with human affective ratings remains unclear. We benchmarked 16 proprietary and open-source VLMs on three validated affective image datasets: the International Affective Picture System, the Nencki Affective Picture System, and the Library of AI-Generated Affective Images. In a zero-shot setting, models performed two tasks: (i) top-emotion classification and (ii) prediction of human ratings on 1--7/9 Likert scales for discrete emotions and affective dimensions. We also evaluated rater-conditioned prompting on the LAI-GAI dataset using de-identified participant metadata. Best models showed strong discrete emotion classification (60%--80% on six-emotion tasks; 60%--75% on a 12-category task), with anger and surprise predicted least accurately. For continuous ratings, models aligned moderately to strongly with human judgments ($r > 0.75$) but showed consistent biases, especially weaker performance on arousal and overestimation of response intensity. Rater-conditioned prompting had only small, inconsistent effects. Overall, VLMs captured broad affective patterns but sometimes lacked the nuance, although newer model generations performed better.
Recent works on language identification and generation have established tight statistical rates at which these tasks can be achieved. These works typically operate under a strong realizability assumption: that the input data is drawn from an unknown distribution necessarily supported on some language in a given collection. In this work, we relax this assumption of realizability entirely, and impose no restrictions on the distribution of the input data. We propose objectives to study both language identification and generation in this more general "agnostic" setup. Across both problems, we obtain novel interesting characterizations and nearly tight rates.
Inference efficiency in Large Language Models (LLMs) is fundamentally limited by their serial, autoregressive generation, especially as reasoning becomes a key capability and response sequences grow longer. Speculative decoding (SD) offers a powerful solution, providing significant speed-ups through its lightweight drafting and parallel verification mechanism. While existing work has nearly saturated improvements in draft effectiveness and efficiency, this paper advances SD from a new yet critical perspective: the verification cost. We propose TriSpec, a novel ternary SD framework that, at its core, introduces a lightweight proxy to significantly reduce computational cost by approving easily verifiable draft sequences and engaging the full target model only when encountering uncertain tokens. TriSpec can be integrated with state-of-the-art SD methods like EAGLE-3 to further reduce verification costs, achieving greater acceleration. Extensive experiments on the Qwen3 and DeepSeek-R1-Distill-Qwen/LLaMA families show that TriSpec achieves up to 35% speedup over standard SD, with up to 50% fewer target model invocations while maintaining comparable accuracy.
More than 80% of the 1.6B English speakers do not use Standard American English (SAE), yet LLMs often fail to correctly identify non-SAE dialects and generate stereotyped responses for their speakers. We introduce DialectLLM, the first large-scale framework for generating high-quality multi-dialectal conversational data encompassing the three pillars of written dialect---lexical (vocabulary), orthographic (spelling), and morphosyntactic (grammar) features. DialectLLM produces a dialect-parallel dialog dataset spanning nine English dialects. Partnering with native linguists, we design and validate SAE-to-dialect transformation rules, ensuring authenticity. Our approach challenges the prevailing practice of applying a single morphosyntactic feature set to both user utterances and model responses, showing that models should not reproduce up to 90% of the grammatical features of a dialect. Human evaluation confirms data quality, with annotators preferring DialectLLM over prior methods in 98.8% of pairwise comparisons for dialect naturalness. We then construct DialectLLM-Bench, a dialect-parallel benchmark with 50k+ dialogs, resulting in 97k+ QA pairs, and evaluate 17 LLMs on dialect identification and response generation tasks. Even frontier models achieve around 70% accuracy, fail to reach 40% for prominent dialects like Canadian English, and systematically misclassify non-SAE dialects as American or British. Beyond benchmarking, we show that DialectLLM data also serve as a scalable LLM post-training resource, suggesting a practical path toward dialect-aware conversational AI.
Best-of-N selection for code generation is only as good as the judge that picks among the candidates, and the judges accurate enough for this role are costly frontier LLMs. A small judge can be trained on abundant earlier data labeled by unit tests, but it is deployed on newer problems: the earlier problems are easier and teach it little, their tasks differ, and their candidates come from weaker policies, while data close to deployment are scarce. We present JURIST, a 1.7B reasoning judge that uses a small meta set close to deployment to learn how much each earlier training sample should count. In our experiments, we find that training on the meta set alone, or adding it to the training data, does not help, while using it to weight the earlier data through a bi-level objective adds 3.9 Hard points over adding it to the training data. On LiveCodeBench, JURIST reaches 84.7% Overall and matches Claude Sonnet 4.6 on Hard problems (73.8%) at 463x lower cost; on BigCodeBench, it surpasses all test-time scaling baselines (35.9%). Abundant earlier data, guided by a small deployment-aligned meta set, can thus train a compact judge that selects accurately at low cost.
Generating clinical reports that summarize abnormal patterns, diagnostic findings, and clinical interpretations from long-term EEG recordings remains labor-intensive. We present CELM, the first clinical EEG-to-Language foundation model capable of summarizing long-duration, variable-length EEG recordings and performing end-to-end clinical report generation at multiple scales. CELM integrates pretrained EEG foundation models with language models to enable scalable multimodal learning. We curate a large-scale clinical EEG dataset containing 9,922 reports paired with approximately 11,000 hours of EEG recordings from 9,048 patients to train CELM, and release the benchmark pipeline, with automated report-structuring, to facilitate future research. Experimental results show that CELM consistently outperforms existing methods across all evaluation settings. Importantly, we further conduct human evaluation with clinical experts, demonstrating that CELM generates reports that are more clinically coherent, diagnostically reliable, and better aligned with expert interpretation. We release our model code and benchmark construction pipeline at https://github.com/Jathurshan0330/CELM.
We develop diffusion-based samplers for target distributions known up to a normalising constant. To this end, we rely on the well-known diffusion path that smoothly interpolates between a simple base distribution and the target, popularised by diffusion models. We tackle the score estimation problem by developing an efficient sequential Monte Carlo sampler that evolves auxiliary variables from conditional distributions along the path, providing principled score and density estimates for time-varying distributions. To control the variance of score estimates, we further propose practical control variate schedules that incur minimal overhead. We adapt this general framework to paths induced by the Ornstein-Uhlenbeck (OU) time-reversal process, stochastic interpolants, and diffusion annealed Langevin dynamics, outlining their trade-offs. Finally, we provide theoretical guarantees and empirically demonstrate the effectiveness of our method on several synthetic and real-world datasets. Our code is available at: https://github.com/akyildiz-group/dpsmc.
We study privacy amplification for differentially private model training with matrix factorization under random allocation (also known as the balls-in-bins model). Recent work by Choquette-Choo et al. (2025) proposes a sampling-based Monte Carlo approach to compute amplification parameters in this setting. However, their guarantees either only hold with some high probability or require random abstention by the mechanism. Furthermore, the required number of samples for ensuring $(\epsilon,\delta)$-DP is inversely proportional to $\delta$. In contrast, we develop sampling-free bounds based on R\'enyi divergence and conditional composition. The former is facilitated by a dynamic programming formulation to efficiently compute the bounds. The latter complements it by offering stronger privacy guarantees for small $\epsilon$, where R\'enyi divergence bounds inherently lead to an over-approximation. Our framework applies to arbitrary banded and non-banded matrices. Through numerical comparisons, we demonstrate the efficacy of our approach across a broad range of matrix mechanisms used in research and practice.
We address the problem of accurate, training-free guidance for conditional generation in trained diffusion models. Existing methods typically rely on point-estimates to approximate the posterior score, often resulting in biased approximations that fail to capture multimodality inherent to the reverse process of diffusion models. We propose a sequential Monte Carlo (SMC) framework that integrates over the full denoising distribution to construct an unbiased Monte Carlo estimator of $p_\theta(y\mid x_t)$, avoiding the bias of existing training-free guidance methods. To ensure computational tractability, we incorporate variance-reduction schemes based on MultiLevel Monte Carlo (MLMC). We evaluate on three image generation tasks. On CIFAR-10 and ImageNet classifier guidance, our approach improves accuracy by 18 and 38 points respectively over the strongest training-free baseline. On CelebA-HQ joint-attribute guidance, three existing methods (DPS, LGD, TFG) used within our framework all outperform the closest SMC-based comparator (TDS) by 25-37 points, and improve over the strongest baseline by 2-14 points.
Optimisation algorithms typically assume oracle access to objectives and constraints, yet in many applications these functions are available only through data. To recover them, a common approach learns a surrogate for each function and optimises over it, but this can be unnecessarily expensive and may fail to preserve the structure required by the optimisation algorithm. We study this problem for monotone optimisation, a class of non-convex problems which can be solved globally using the Polyblock Outer Approximation (POA) algorithm. Our key observation is that POA does not require the constraint functions themselves, but only a scalar quantity we call the radial inverse. Learning this quantity directly replaces POA's inner bisection loop over a learned constraint with a single radial inverse evaluation. We characterise exactly which maps arise as radial inverses of monotone constraints, and relate their approximation error to perturbations of the feasible set and optimal objective value. Building on these results, we introduce Homogeneous-Monotone Radial Inverse (HM-RI) networks, structured min-max models that enforce the monotonicity and positive-homogeneity of radial inverses by construction. Across synthetic monotone optimisation problems and a cellular transmit-power application, HM-RI improves solution quality over direct constraint surrogates and local optimisation baselines by up to 58.6 percentage points relative to the Oracle reference objective, while achieving a 3.0-4.7x speed-up over bisection-based POA with learned constraint surrogates, with the largest solution-quality gains occurring in data-limited and piecewise-constant settings.
Differential privacy limits how much each training image can influence a medical imaging model but lowers its diagnostic accuracy. Pretrained weights reduce this loss. It is unknown whether the pretraining objective or domain improves accuracy more. Here we train chest radiograph classifiers under differential privacy from five starting points on 591,884 radiographs from five datasets. We show that supervised pretraining on chest radiographs ranks first in 24 of 25 combinations of dataset and privacy setting. Measured by the area under the receiver operating characteristic curve, its lead over supervised pretraining on natural images is 2.6% without privacy and 15.3% at the strictest budget. From that starting point, changing the domain improves this area 2.2 to 3.7 times more than changing the objective. Private pretraining on chest radiographs lowers this area by about 5% under privacy but outperforms the three starting points without patient data at the strictest budget. Under privacy, supervised chest radiograph pretraining leads in the worst-performing demographic subgroup.
The success of RL is fundamentally tied to having a reward function that accurately reflects the task objective. Yet, designing reward functions is time-consuming and prone to misspecification. To address this, our first goal is to understand how to support RL practitioners in specifying reward functions. We leverage the Trajectory Alignment Coefficient (TAC), a metric that evaluates how closely a reward function's induced preferences match those of a domain expert. To test whether TAC provides meaningful support, we evaluated its impact in a user study where practitioners designed reward functions for Lunar Lander. We found that iterative TAC feedback allowed participants to design more performant reward functions while reporting lower cognitive workload than unguided reward design. However, the study also underscored that manual reward design, even with TAC, remains labor-intensive. This motivated our second goal: to learn a reward model that maximizes TAC directly. Specifically, we propose Soft-TAC, a differentiable approximation of TAC that can be used as a loss function to train reward models from human preferences. Our theoretical analysis guarantees that under Soft-TAC optimization: (1) the optimal reward function belongs exclusively to the set of human-aligned reward functions, and (2) this optimal reward remains invariant whether trained on clean or noisy preference data. Empirically, we first evaluate reward models trained with Soft-TAC across three robotic tasks, demonstrating that they consistently match or outperform those trained with the Cross-Entropy loss. Next, we evaluate Soft-TAC in the GT7 driving simulator, and find that Soft-TAC-trained reward models more successfully capture preference-specific driving styles. Ultimately, this work demonstrates that TAC can serve as both a practical tool for guiding reward design and a reward learning objective in complex domains.
Interactive narrative (IN) authors craft spaces of divergent narrative possibilities for players to explore, with the player's input determining which narrative possibilities they actually experience. Generative AI can enable new forms of IN by improvisationally expanding on pre-authored content in response to open-ended player input. However, this extrapolation risks widening the gap between author-envisioned and player-experienced stories, potentially limiting the strength of plot progression and the communication of the author's narrative intent. To bridge the gap, we introduce Elsewise: an authoring tool for LLM-based INs that implements a novel Bundled Storyline concept to enhance author's perception and understanding of the narrative possibility space, allowing authors to explore similarities and differences between possible playthroughs of their IN in terms of open-ended, user-configurable narrative dimensions. A user study (n=12) shows that our approach improves author anticipation of player-experienced narrative, leading to more effective control and exploration of the narrative possibility spaces.
Robust machine learning for regulatory genomics is studied under biologically and technically induced distribution shifts. Deep convolutional and attention-based models achieve strong in-distribution performance on DNA regulatory sequence prediction tasks but are usually evaluated under i.i.d. assumptions, even though real applications involve cell type-specific programs, evolutionary turnover, assay protocol changes, and sequencing artifacts. We introduce a robustness framework that combines a mechanistic simulation benchmark with real data analysis on an MPRA dataset to quantify performance degradation, calibration failures, and uncertainty-based reliability. In simulation, motif-driven regulatory outputs are generated with cell type-specific programs, PWM perturbations, GC bias, depth variation, batch effects, and heteroscedastic noise, and CNN, BiLSTM, and transformer models are evaluated. Models remain accurate and reasonably calibrated under mild GC content shifts but show higher error, severe variance miscalibration, and coverage collapse under motif effect rewiring and noise-dominated regimes, revealing robustness gaps invisible to standard i.i.d. evaluation. Adding simple biological structural priors, motif-derived features in simulation, and global GC content in MPRA improves distribution error and yields consistent robustness gains under biologically meaningful genomic shifts, while providing only limited protection against strong assay noise. Uncertainty-aware selective prediction offers an additional safety layer, and risk coverage analyses on simulated and MPRA data show that filtering low-confidence inputs recovers low-risk subsets, including under GC-based OOD conditions, although reliability gains diminish when noise dominates.
In a multi-agent system (MAS), removing an expensive agent call can save tokens without making the answer arrive sooner: latency depends on the longest chain of dependent calls. We present Latency-Aware Multi-agent System (LAMaS), a modular platform for learning low-latency orchestration. A framework can send a query to LAMaS and receive an answer and execution trace, or keep its own graph builder and use the platform's latency-optimized orchestration layer. During training, the layer gives operators on the critical path the greatest latency penalty while targeting a floor on expected accuracy. During inference, a lightweight controller uses partial results to decide whether to skip work that has yet to start. Across four benchmarks, LAMaS reduces end-to-end latency by over 50% relative to the evaluated learning-based MAS baselines while maintaining competitive or better accuracy. GDesigner and AgentDropout retain their graph builders and train the controller on their own traces; both achieve lower latency across the same benchmarks.