论文检索

Probabilistic Methods · Bayesian Models and Methods

Sergio Calvo Ordoñez, Jonathan Plenk, Richard Bergna, Alvaro Cartea, Yarin Gal, Jose Miguel Hernandez-Lobato, Kamil Ciosek

Bayesian last layers (BLLs) provide a convenient and computationally efficient way to estimate uncertainty in neural networks. However, they underestimate epistemic uncertainty because they apply a Bayesian treatment only to the final layer, ignoring uncertainty induced by earlier layers. We propose a method that improves BLLs by leveraging a projection of Neural Tangent Kernel (NTK) features onto the space spanned by the last-layer features. This enables posterior inference that accounts for variability of the full network while retaining the low computational cost of a standard BLL. We show that our method yields posterior variances that are provably greater or equal to those of a standard BLL, correcting its tendency to underestimate epistemic uncertainty. To further reduce computational cost, we introduce a uniform subsampling scheme for estimating the projection matrix and for posterior inference. We derive approximation bounds for both types of sub-sampling. Empirical evaluations on UCI regression, contextual bandits, and out-of-distribution detection tasks in image and tabular datasets, demonstrate improved calibration and uncertainty estimates compared to standard BLLs and competitive baselines, while reducing computational cost.

General Machine Learning · Kernel methods

Rikhav Shah, Sandeep Silwal, Haike Xu

This paper studies the use of *kernel density estimation* (KDE) for linear algebraic tasks involving the *kernel matrix* of a collection of $n$ data points in $\mathbb{R}^d$. In particular, we improve upon the best existing algorithms for computing the following up to $(1+\varepsilon)$ relative error for a Gaussian kernel matrix and other kernels: matrix-vector products, matrix-matrix products, the spectral norm, and sum of all entries. The runtimes of our algorithms depend linearly on the dimension $d$, sub-quadratically in the number of points $n$, and polynomially on the target error $\varepsilon$. Importantly, the dependence on $n$ in each case is far lower when accessing the kernel matrix through KDE queries as opposed to reading individual entries. Our improvements over existing best algorithms (particularly those of [Backurs et al. ICML `21]) for these tasks reduce the polynomial dependence on $\varepsilon$, and additionally decrease the dependence on $n$ in the case of computing the sum of all entries of the kernel matrix. For example, we reduce the power of $1/\epsilon$ from $\approx 7.7$ to $\approx 3.2$ for a $1-\varepsilon$ relative error estimation of the spectral norm of a Gaussian kernel matrix. We complement our upper bounds with several lower bounds for related problems, which provide (conditional) quadratic time hardness results and additionally hint at the limits of KDE based approaches for the problems we study.

General Machine Learning · Causality

Alek Fröhlich, Vladimir Kostic, Karim Lounici, Daniel Rodrigues Perazzo, Daniel Tiezzi, Massimiliano Pontil

Conditional independence (CI) is central to causal inference, feature selection, and graphical modeling, yet it is untestable in many settings without additional assumptions. Existing CI tests often rely on restrictive structural conditions, limiting their validity. Kernel methods using partial covariance operators offer a more principled approach but suffer from limited adaptivity, slow convergence, and poor scalability. In this work, we explore whether representation learning can help address these limitations. Specifically, we focus on representations derived from the singular value decomposition of partial covariance operators and use them to construct a simple test statistic. We also introduce a bi-level contrastive algorithm to learn these representations. Our theory links representation learning error to test performance and establishes asymptotic validity and power guarantees. Experiments on real and synthetic data suggest that this approach offers a principled and statistically grounded path toward scalable CI testing, bridging kernel-based theory with modern representation learning.

Deep Learning · Large Language Models

Chi-Chih Chang, Wei-Cheng Lin, Chien-Yu Lin, Hung-Yueh Chiang, Yash Akhauri, Xilai Dai, Huiqiang Jiang, Yucheng Li, Kai-Chiang Wu, Luis Ceze 等

Long-context Large Language Models (LLMs) enable powerful applications but incur high memory costs due to the key–value states (KV-Cache). Recent studies attempt to share KV-Cache across layers, but these approaches either require expensive pretraining or rely on per-token cross-layer cosine similarity that is often limited in practice. We show, via Centered Kernel Alignment (CKA), that the dominant singular vectors of KV-Cache are well aligned across layers. Motivated by this observation, we propose xKV, a post-training compression method that jointly factorizes grouped-layer KV-Cache into a shared low-rank subspace, substantially reducing KV-Cache memory. Across widely used LLMs, xKV achieves up to 8× KV-Cache compression while preserving accuracy on long-context tasks and in multi-turn settings. To further improve efficiency, we introduce Selective Reconstruction (SR) at decode time. Combined with SR, xKV achieves up to 4.23× end-to-end speedup, surpassing notable baselines with 30% higher throughput under a similar accuracy level. Overall, xKV provides a plug-and-play approach to reduce both memory and latency for long-context LLM inference. Our code will be open-sourced.

Deep Learning · Foundation Models

Yanlong Chen, Amir Habibian, Luca Benini, Yawei Li

Vision-Language Models (VLMs) achieve strong multimodal performance but are costly to deploy, and post-training quantization often causes significant accuracy loss. Despite its potential, quantization-aware training for VLMs remains underexplored. We propose GRACE, a framework unifying knowledge distillation and QAT under the Information Bottleneck principle: quantization constrains information capacity while distillation guides what to preserve within this budget. Treating the teacher as a proxy for task-relevant information, we introduce confidence-gated decoupled distillation to filter unreliable supervision, relational centered kernel alignment to transfer visual token structures, and an adaptive controller via Lagrangian relaxation to balance fidelity against capacity constraints. Across extensive benchmarks on LLaVA and Qwen families, our INT4 models consistently outperform FP16 baselines (e.g., LLaVA-1.5-7B: 70.1 vs. 66.8 on SQA; Qwen2-VL-2B: 76.9 vs. 72.6 on MMBench), nearly matching teacher performance. Using real INT4 kernel, we achieve 3× throughput with 54\% memory reduction. This principled framework significantly outperforms existing quantization methods, making GRACE a compelling solution for resource-constrained deployment.

Theory · Learning Theory

Kedar Karhadkar, Alexander Sietsema, Deanna Needell, Guido Montufar

Motivated by recent work on benign overfitting in overparameterized machine learning, we study the generalization behavior of functions in Sobolev spaces $W^{k, p}(\mathbb{R}^d)$ that perfectly fit a noisy training data set. Under assumptions of label noise and sufficient regularity in the data distribution, we show that approximately norm-minimizing interpolators, which are canonical solutions selected by smoothness bias, exhibit harmful overfitting: even as the training sample size $n \to \infty$, the generalization error remains bounded below by a positive constant with high probability. Our results hold for arbitrary values of $p \in [1, \infty)$, in contrast to prior results studying the Hilbert space case ($p = 2$) using kernel methods. Our proof uses a geometric argument which identifies harmful neighborhoods of the training data using Sobolev inequalities.

General Machine Learning · Transfer, Multitask and Meta-learning

Jean-Baptiste Fermanian, Batiste Le Bars, Aurélien Bellet

Personalized Federated Learning enables a collection of agents to collaboratively learn individual models without sharing raw data. We propose a new approach in which each agent optimizes a weighted combination of all agents' empirical risks, with the weights learned from data rather than specified a priori. The novelty of our method lies in formulating the estimation of these collaborative weights as a kernel mean embedding estimation problem with multiple data sources, leveraging tools from multi-task averaging to capture statistical relationships between agents. This perspective yields a fully adaptive procedure that requires no prior knowledge of data heterogeneity and automatically transitions between global and local learning regimes. By recasting the objective as a high-dimensional mean estimation problem, we derive finite-sample guarantees on local excess risks for a broad class of distributions, explicitly quantifying the statistical gains of collaboration. To address communication constraints inherent to federated settings, we also propose a practical implementation based on random Fourier features, which allows one to trade communication cost for statistical efficiency. Numerical experiments validate our theoretical results.

Deep Learning · Large Language Models

Haodong WANG, Junjie Liu, Zicong Hong, Qianli Liu, Jian Lin, Song Guo, Xu Chen

4-bit quantization reduces the memory footprint and latency of large language model inference, but its aggressive precision reduction can severely degrade accuracy. Prior methods address this by decomposing each weight matrix into two components (e.g., via singular value decomposition) and quantizing them separately, assigning the bulk of values to a low-precision residual component while handling outliers with a high-precision low-rank component. However, such decompositions are designed to minimize the real-valued energy of the residual, rather than the post-quantization error of the residual and low-rank components. We propose TwinQuant, a 4-bit quantization framework that learns quantization-friendly decomposed subspaces and jointly reshapes both the low-rank and residual components. TwinQuant learns component-specific transformations via a joint optimization over the Stiefel and general linear manifolds, flattening their distributions and reducing dynamic-range imbalance. To enable efficient end-to-end execution, we further design a fused dual-component kernel that pipelines the two-stage low-rank computation on-chip and merges both components with a single epilogue, avoiding intermediate global-memory traffic. Across LLaMA3 and Qwen3 models, TwinQuant preserves near-FP16 accuracy and delivers up to $2.11\times$ end-to-end speedup over an FP16 baseline.

Probabilistic Methods · Everything Else

Ádám Jung, Domokos Kelen, Andras Benczur

A key challenge in probabilistic regression is ensuring that predictive distributions accurately reflect true empirical uncertainty. Minimizing overall prediction error often encourages models to prioritize informativeness over calibration, producing narrow but overconfident predictions. However, in safety-critical settings, trustworthy uncertainty estimates are often more valuable than narrow intervals. Realizing the problem, several recent works have focused on post-hoc corrections; however, existing methods either rely on weak notions of calibration (such as PIT uniformity) or impose restrictive parametric assumptions on the nature of the error. To address these limitations, we propose a novel nonparametric re-calibration algorithm based on conditional kernel mean embeddings, capable of correcting calibration error without restrictive modeling assumptions. For efficient inference with real-valued targets, we introduce a novel characteristic kernel over distributions that can be evaluated in $\mathcal{O}(n \log n)$ time for empirical distributions of size $n$. We demonstrate that our method consistently outperforms prior re-calibration approaches across a diverse set of regression benchmarks and model classes.

Haoyu Wang, Haiyan Zhao, Xingyu Yu, Zhangyang Yao, Xu Han, Zhiyuan Liu, Maosong Sun

Post-training quantization at the 2-bit level enables low-cost deployment and inference acceleration for large language models (LLMs). Scalar quantization (SQ) and vector quantization (VQ) are two primary quantization methods, however, the former suffers from significant performance degradation, and the latter incurs computational and storage overhead. We propose UniSVQ, a unified 2-bit quantization framework that bridges scalar and vector quantization by parameterizing codewords as an affine transform of integer lattices. This structure preserves compatibility with optimized integer kernels while retaining much of VQ's flexibility. We further introduce a data-driven block-wise fine-tuning strategy to directly minimize quantization reconstruction error. Extensive experiments across multiple LLM families and zero-shot benchmarks demonstrate that UniSVQ consistently outperforms state-of-the-art SQ methods and achieves performance comparable to advanced VQ methods, while providing higher inference throughput.

Deep Learning · Large Language Models

Senyu Han, Yilu Cao, Kai Yu, Lu Chen

Large language model (LLM) exists a subset of attention heads that are highly responsible for long-context processing. Existing work has identified different long-context heads in models, but their detection methods mainly rely on model inference on actual long texts and do not analyze the inherent properties of the head parameters. In this paper, we use kernel methods to analyze static *frequency kernels* formed by different rotation frequency components of attention heads, and we design a Long-context Potential Score (LPS) to measure the potential of attention heads in processing long contexts. Kernels of heads with high LPS exhibit concentrated low-frequency energy and low effective rank, which allow them to effectively capture highly specialized information from distant contexts. Experiments and analysis on long-context tasks and model behaviors show that the LPS metrics can well reflect the actual capability of heads on long contexts. Furthermore, by simply amplifying low-frequency kernels of heads with high retrieval potential, we can further improve model's performance on long-context tasks. Our metrics and head enhancement methods are fully static and offline, and they can be quickly conducted under low-resource constraints.

Deep Learning · Sequential Models, Time series

Egor Serov, Ilya Kuleshov, Alexey Zaytsev

Neural Controlled Differential Equations (Neural CDEs) provide a powerful continuous-time framework for sequence modeling, yet the roughness of the driving control path often restricts their efficiency. Standard splines introduce high-frequency variations that force adaptive solvers to take excessively small steps, driving up the Number of Function Evaluations (NFE). We propose a novel approach to Neural CDE path construction that replaces exact interpolation with Kernel and Gaussian Process (GP) smoothing, enabling explicit control over trajectory regularity. To recover details lost during smoothing, we propose an attention-based Multi-View CDE (MV-CDE) and its convolutional extension (MVC-CDE), which employ learnable queries to inform path reconstruction. This framework allows the model to distribute representational capacity across multiple trajectories, each capturing distinct temporal patterns. Empirical results demonstrate that our method, MVC-CDE with GP, achieves state-of-the-art accuracy while significantly reducing NFEs and total inference time compared to spline-based baselines.

Theory · Learning Theory

Jiaolong Wang, Fode Zhang, Lingrui Wang

Reliable estimation of Amari $\alpha$-divergences underpins variational inference, yet unconstrained neural critics are notoriously prone to instability. We propose a scalable estimator by constraining the critic to a Reproducing Kernel Hilbert Space (RKHS) ball and approximating the kernel via band-limited Random Fourier Features (RFF). This formulation yields a linear-time objective amenable to mini-batch stochastic optimization while avoiding the cubic complexity of Gram-matrix methods. We present a unified analysis based on a four-term error decomposition—comprising RKHS approximation, feature discretization, statistical deviation, and optimization residual. Under a spectral source condition, we derive non-asymptotic bounds establishing that the RKHS approximation bias scales as $\mathcal{O}(R^{-\gamma})$, the RFF discretization error as $\mathcal{O}(R D^{-1/2})$, and the statistical error as $\mathcal{O}(R n^{-1/2})$. We further show that statistical non-degeneracy induces intrinsic local curvature, enabling our proposed Armijo-SGD to achieve local linear convergence. Empirical evaluations demonstrate that the RFF-RKHS estimator outperforms varying-representation baselines in stability, and applying this spectral regularization to GAN critics significantly enhances the capture of high-frequency data components.

Applications · Neuroscience, Cognitive Science

Bozhou Li, Chubo Liu, Yan Ding, Yufeng Zhang, Zhuo Tang, Kenli Li

ANN-to-SNN conversion offers energy-efficient inference but faces a fidelity-latency trade-off due to open-loop error accumulation. While conversion-aware training mitigates this, it sacrifices the generality of using off-the-shelf ANNs. We propose a closed-loop fine-tuning framework that calibrates these errors without altering the source model. Our approach employs a Dual Alignment Mechanism, utilizing global Kullback-Leibler divergence for output distillation and introducing an adaptive local Centered Kernel Alignment constraint, weighted by initial conversion loss, for feature alignment. We uncover a critical time-dependent dynamic: local constraints are essential for stabilizing representations in low-latency regimes (e.g., $T=8$) where global gradients are unstable, whereas global alignment drives fidelity at higher time steps. Experiments on CIFAR-10 demonstrate that our method achieves over 99\% of source ANN accuracy at $T=32$ (e.g., ResNet-18: 96.38\% vs.\ 96.39\%). Furthermore, this fine-tuning acts as a regularizer, yielding SNNs with input noise robustness that matches or exceeds the source ANN.

Reinforcement Learning · Multi-agent

Shiyang Li, Zijian Zhang, Winson Chen, Yuebo Luo, Mingyi Hong, Caiwen Ding

Modern machine learning (ML) workloads increasingly rely on GPUs, yet achieving high end-to-end performance remains challenging due to dependencies on both GPU kernel efficiency and host-side settings. Although LLM-based methods show promise on automated GPU kernel generation, prior works mainly focus on single-kernel optimization and do not extend to end-to-end programs, hindering practical deployment. To address the challenge, in this work, we propose \textsc{StitchCUDA}, a multi-agent framework for end-to-end GPU program generation, with three specialized agents: a \textit{Planner} to orchestrate whole system design, a \textit{Coder} dedicated to implementing it step-by-step, and a \textit{Verifier} for correctness check and performance profiling using Nsys/NCU. To fundamentally improve the \textit{Coder}'s ability in end-to-end GPU programming, \textsc{StitchCUDA} integrates rubric-based agentic reinforcement learning over two atomic skills, task-to-code generation and feedback-driven code optimization, with combined rubric reward and rule-based reward from real executions. Therefore, the \textit{Coder} learns how to implement advanced CUDA programming techniques (e.g., custom kernel fusion, cublas epilogue), and we also effectively prevent \textit{Coder}'s reward hacking (e.g., just copy PyTorch code or hardcoding output) during benchmarking. Experiments on KernelBench show that \textsc{StitchCUDA} achieves nearly 100\% success rate on end-to-end GPU programming tasks, with 1.72$\times$ better speedup over the multi-agent baseline and 2.73$\times$ than the RL model baselines.

Applications · Chemistry, Physics, and Earth Sciences

Martin Jankowiak, Yerdos Ordabayev, Rudraksh Tuwani, Henry Ward, Hunter Nisonoff, James McFarland, Gevorg Grigoryan

Despite its importance to applications in protein design, predicting protein properties like binding affinity and thermostability from sparse experimental data remains a significant challenge. Accordingly, we introduce a class of sequence kernels that exploit evolutionary substitution matrices as well as local linearity and demonstrate that the resulting Gaussian processes provide data-efficient models of protein property landscapes, frequently outperforming alternatives that rely on foundation model embeddings. Furthermore--by learning what are in effect structure-aware substitution matrices--we show that our kernels can readily incorporate structural information from foundation models. We demonstrate that these structure-conditioned kernels are well suited to multi-task learning across multiple protein property landscapes and can decisively outperform local supervised learning methods.

Deep Learning · Large Language Models

Hoang Anh Duy Le, Sahil Joshi, Zeyu Yang, Zhaozhuo Xu, Anshumali Shrivastava

Self-attention dominates the computational and memory cost of long-context LLM inference across both prefill and decode phases. To address this challenge, we introduce **Sketch\&Walk** Attention, a training-free sparse attention method that determines sparsity with lightweight sketches and deterministic walk. Sketch\&Walk applies Hadamard sketching to get inexpensive approximations of attention scores, then aggregates these estimates across layers via a walk mechanism that captures attention influence beyond direct interactions between tokens. The accumulated walk scores are used to select top-$k$ attention blocks, enabling dynamic sparsity with a single training-free algorithm that applies uniformly to both the prefill and decode phases, together with custom sparse attention kernels. Across a wide range of models and tasks, Sketch\&Walk maintains near-lossless accuracy at 20\% attention density and can slightly outperform dense attention in some settings, while achieving up to $6\times$ inference speedup.

Xinyou Wang, Liang Hong, Jiasheng Ye, Zaixiang Zheng, Shujian Huang, Quanquan Gu

Proteins are shaped by gradual evolution under biophysical and functional constraints. Protein language models learn rich evolutionary constraints from large-scale sequence data, and discrete diffusion–based protein language models (e.g., DPLMs) have emerged as a promising framework for both understanding and generation. However, existing DPLMs typically rely on masking-based absorbing diffusion, which conflicts with a basic biological intuition: proteins evolve through accumulated edits rather than emerging from masked tokens. As a result, these frameworks lack explicit pretraining objectives for substitution and insertion/deletion (indel) operations, which in turn limits both optimization-style post-editing and flexible guided generation. To address these limitations, we present DPLM-Evo, an evolutionary discrete diffusion framework that explicitly predicts substitution, insertion, and deletion operations during denoising. \method decouples a fixed-length latent alignment space from the variable-length observed sequence space, making indel-aware generation tractable and enabling adaptive scaffold growth throughout the process with negligible computational overhead. To further align substitutions with real evolutionary dynamics, we introduce a contextual evolutionary noising kernel that induces biologically informed, context-dependent mutation patterns. Across tasks, \method improves sequence understanding and achieves state-of-the-art performance on ProteinGym in the single-sequence setting, while also enabling variable-length simulated evolution, guided generation, and post-editing or optimization of existing proteins via explicit edit trajectories.

General Machine Learning · Kernel methods

Anirban Chatterjee, Sayantan Choudhury, Rohan Hore

How can we generate samples from a conditional distribution that we never fully observe? This question arises across a broad range of applications in both modern machine learning and classical statistics, including image post-processing in computer vision, approximate posterior sampling in simulation-based inference, and conditional distribution modeling in complex data settings. In such settings, compared with unconditional sampling, additional feature information can be leveraged to enable more adaptive and efficient sampling. Building on this, we introduce Conditional Generator using MMD (CGMMD), a novel framework for conditional sampling. Unlike many contemporary approaches, our method frames the training objective as a simple, adversary-free direct minimization problem. A key feature of CGMMD is its ability to produce conditional samples in a single forward pass of the generator, enabling practical one-shot sampling with low test-time complexity. We establish rigorous theoretical bounds on the loss incurred when sampling from the CGMMD sampler, and prove convergence of the estimated distribution to the true conditional distribution. In the process, we also develop a uniform concentration result for nearest-neighbor based functionals, which may be of independent interest. Finally, we show that CGMMD performs competitively on synthetic tasks involving complex conditional densities, as well as on practical applications such as image denoising and image super-resolution.

Na Li, Hangguan Shan, Wei Ni, Wenjie Zhang, Xinyu Li

Actor-critic (AC) methods are a cornerstone of reinforcement learning (RL) but offer limited interpretability. Current explainable RL methods seldom use *state attributions* to assist training. Rather, they treat all state features equally, thereby neglecting the heterogeneous impacts of individual state dimensions on the reward. We propose *RKHS-SHAP-based Advanced Actor-Critic (RSA2C)*, an attribution-aware, kernelized, two-timescale AC algorithm, including Actor, Value Critic, and Advantage Critic. The Actor is instantiated in a vector-valued reproducing kernel Hilbert space (RKHS) with a Mahalanobis-weighted operator-valued kernel, while the Value Critic and Advantage Critic reside in scalar RKHSs. These RKHS-enhanced components use sparsified dictionaries: the Value Critic maintains its own dictionary, while the Actor and Advantage Critic share one. State attributions, computed from the Value Critic via RKHS-SHAP (kernel mean embedding for on-manifold and conditional mean embedding for off-manifold expectations), are converted into Mahalanobis-gated weights that modulate Actor gradients and Advantage Critic targets. We derive a global, non-asymptotic convergence bound under *state perturbations*, showing stability through the perturbation-error term and efficiency through the convergence-error term. Empirical results on three continuous-control environments show that RSA2C achieves efficiency, stability, and interpretability.