论文检索

Deep Learning · Generative Models and Autoencoders

Langzhang Liang, Ming Yang, Yi Feng, Junfan Li, Shirui Pan, Xu Yinghui, Tianlei Ying, YIZHEN ZHENG, Zenglin Xu

Protein sequence generation for engineering requires samples that are biophysically plausible and, when targeting a family/domain, remain recognizable members while exploring within-family diversity. Current discrete generative models typically start from uniform or masked-token noise, which discards strong position-specific constraints induced by evolution and forces the model to reconstruct conserved residues from scratch, leading to weak family control and low foldability. We propose \emph{LineageFlow}, a simplex-valued flow-matching model that initializes generation from lineage priors derived from ancestral sequence reconstruction, turning generation into structured mutation from an evolved scaffold. Across diverse protein families, LineageFlow achieves family validity close to held-out natural sequences and improves predicted structural confidence over uniform-/mask-initialized baselines while maintaining substantial within-family novelty and diversity, even surpassing a large pretrained baseline trained on substantially more data. Finally, we introduce \emph{rerouting}, a single intermediate-time mutate--select--amplify intervention that enables objective-guided sampling without per-step predictor guidance and yields further gains in plausibility, including a zero-shot enzyme generation case study.

Deep Learning · Large Language Models

Zhen Zhang, Wanjing Zhou, Juncheng Li, Hao Fei, Jun Wen, Wei Ji

Compared with individual agents, large language model based multi-agent systems have demonstrated great capabilities across a wide range of tasks, including code generation, mathematical reasoning, and planning, etc. Despite their impressive performance, the effectiveness and robustness of these systems heavily rely on their communication topology, which is often fixed or generated in a single step. This restricts fine-grained structural exploration and flexible composition, leading to excessive token consumption for simple tasks or performance bottlenecks for complicated ones. To address this challenge, we introduce RADAR, a redundancy-aware and query-adaptive generative framework that actively reduce communication overhead. Inspired by conditional discrete graph diffusion models, we formulate communication topology synthesis as a step-by-step generation process, guided by the effective size of the graph. Comprehensive experiments on six benchmarks demonstrate that RADAR consistently outperforms recent baselines, achieving higher accuracy, lower token consumption, and greater robustness across diverse scenarios. Our source code and data are available at https://anonymous.4open.science/r/RADAR-8430.

Applications · Health / Medicine

Keunho Byeon, Jin Tae Kwak

Spatial transcriptomics offers spatially resolved gene expression profiling within tissue sections, but its cost and limited throughput hinder large-scale deployment. To extend this capability to routine practice, recent computational methods aim to infer spatial gene expression directly from ubiquitous hematoxylin and eosin-stained histology slides. However, most existing models assume Cartesian or geometry-agnostic locality, despite the hexagonal sampling of widely used spot-array platforms, and point-wise regression objectives often yield over-smoothed gene expression profiles, obscuring gene-specific spatial heterogeneity. To address these, we propose HEXST, a geometry-aligned Transformer for spatial gene expression prediction from histology. HEXST operates directly on hexagonal spot coordinates to enable efficient local-to-global contextual modeling via tailored shifted-window attention mechanism and hexagonal rotary positional encoding. To enhance gene-wise spatial contrast, HEXST complements point-wise regression with a contrast-sensitive differential objective and transcriptomic priors from a pretrained single-cell foundation model during training. Across seven spatial transcriptomics datasets, HEXST consistently outperforms state-of-the-art models, providing accurate and robust spatial gene expression predictions while preserving gene-wise contrast and spatial heterogeneity.

General Machine Learning · Hardware and Software

Size Zheng, Xuegui Zheng, Hanshi Sun, Qi Hou, Wenlei Bao, Shiyu Li, Haojie Duanmu, Jin Fang, Chenli Xue, Chenhui Huang 等

The scaling of large language models (LLMs) is currently bottlenecked by the rigidity of distributed programming. While high-performance libraries like CuBLAS and NCCL provide optimized primitives, they lack the flexibility required for rapidly evolving model architectures. Conversely, existing tensor compilers fail to address the complex memory hierarchy of distributed clusters effectively. To bridge this gap, we propose DITRON, a scalable tile-level compiler that democratizes high-performance distributed kernel development. DITRON introduces a novel hierarchical programming abstraction spanning Core, Device, and Task levels to map tensor programs efficiently onto heterogeneous distributed hardware. This abstraction allows DITRON to support diverse parallelism strategies while abstracting away the complexity of inter-node and intra-node communication. Evaluated across large-scale clusters, DITRON achieves performance parity with or exceeding expert-tuned CUDA libraries, delivering speedups of 6%–30% on isolated kernels and 5%–30% on end-to-end inference in vLLM. Furthermore, DITRON demonstrates strong portability, achieving significant speedups on both NVIDIA and AMD platforms.

Applications · Neuroscience, Cognitive Science

Luca M. Schulze Buschoff, Konstantinos Voudouris, Can Demircan, Eric Schulz

Pre-trained vision language models do not have good intuitions about the physical world. Recent work has shown that supervised fine-tuning can improve model performance on simple physical tasks. However, fine-tuned models do not appear to learn robust physical rules that can generalize to new contexts. Based on research in cognitive science, we hypothesize that models need to interact with an environment to properly learn its physical dynamics. We train models that learn through interaction with the environment using reinforcement learning. While learning from interaction allows models to improve their within-task performance, it fails to produce models with generalizable physical intuitions. We find that models trained on one task do not reliably generalize to related tasks, even if the tasks share visual statistics and physical principles, and regardless of whether the models are trained through interaction.

Theory · Reinforcement Learning and Planning

Ziyue Chu, Leonardo Stella

Multi-agent reinforcement learning (MARL) has received increasing attention for solving complex decision-making tasks. Networked MARL approaches offer a decentralized solution for parameter sharing to accelerate training via value aggregation. However, existing federated aggregations rely on convex averaging that may fail to converge to global optima and cause learning rollback in the online learning setting. In this paper, we formally characterize the learning rollback phenomenon arising from aggregating value estimates with unequal uncertainty under heterogeneous online update depths. We propose a novel adaptive global consensus (AGC) mechanism for Q-value aggregation in decentralized MARL policy evaluation, which dynamically adjusts aggregation weights based on agents’ uncertainty. We establish theoretical guarantees on accelerated convergence and bounded learning variance with empirical validations, advancing the state-of-art MARL theory.

Social Aspects · Safety

Tanqiu Jiang, Yuhui Wang, Jiacheng Liang, Ting Wang

LLM agents are increasingly deployed in long-horizon, complex environments to solve challenging problems, but this expansion exposes them to long-horizon attacks that exploit multi-turn user–agent–environment interactions to achieve objectives infeasible in single-turn settings. To measure agent vulnerabilities to such risks, we present AgentLAB, the first benchmark dedicated to evaluating LLM agent susceptibility to adaptive, long-horizon attacks. Currently, AgentLAB supports five novel attack types including intent hijacking, tool chaining, task injection, objective drifting, and memory poisoning, spanning 28 realistic agentic environments, and 644 security test cases. Leveraging AgentLAB, we evaluate representative LLM agents and find that they remain highly susceptible to long-horizon attacks; moreover, defenses designed for single-turn interactions fail to reliably mitigate long-horizon threats. We anticipate that AgentLAB will serve as a valuable benchmark for tracking progress on securing LLM agents in practical settings.

Thiziri Nait Saada, Louis Béthune, Michal Klein, David Grangier, Marco Cuturi, Pierre Ablin

Large-scale models are pretrained on massive web-crawled datasets containing documents of mixed quality, making data filtering essential. A popular method is Classifier-based Quality Filtering (CQF), which trains a binary classifier to distinguish between pretraining data and a small, high-quality set. It assigns each pretraining document a quality score defined as the classifier's score and retains only the top-scoring ones. We provide an in-depth analysis of CQF. We show that while CQF improves downstream task performance, it does not necessarily enhance language modeling on the high-quality set. Importantly, we find that training on CQF-selected data can outperform training directly on the high-quality set, even when the latter is sufficiently large. This finding alone is particularly striking, given the substantial effort and cost recently devoted to augmenting high-quality data. We explain this paradox by the fact that CQF implicitly filters the high-quality dataset as well as the low-quality one. Finally, we introduce an optimization-driven notion of data quality and demonstrate that it can be reliably estimated using small-scale proxy experiments. Altogether, our results both elucidate the mechanisms behind CQF and deepen our understanding of data selection methods widely used in practice.

Deep Learning · Large Language Models

Jun Yang, Yuechun Sun, Yi Wu, Rodrigo Caridad, Yongwei Yuan, Jianan Yao, Shan Lu, Kexin Pei

Large Language Models (LLMs) have shown promising results in automating formal verification. However, existing approaches treat proof generation as a static, end-to-end prediction over source code, relying on limited verifier feedback and lacking access to concrete program behaviors. We present EXVERUS, a counterexample-guided framework that enables LLMs to reason about proofs using behavioral feedback via counterexamples. When a proof fails, EXVERUS automatically generates and validates counterexamples, and then guides the LLM to generalize them into inductive invariants to block these failures. Our evaluation shows that EXVERUS significantly improves proof accuracy, robustness, and token efficiency over the state-of-the-art prompting-based Verus proof generator.

Erdun Gao, Jake Fawkes, Dino Sejdinovic

Estimating the Conditional Average Treatment Effect (CATE) is constrained by the high cost of obtaining outcome measurements, making active learning essential. However, conventional strategies suffer from a fundamental objective mismatch: they reduce uncertainty in model parameters or observable outcomes rather than the unobservable causal quantities of interest. We address this via the principle of causal objective alignment, positing that acquisition functions must target unobservables like potential outcomes or CATE directly. We operationalize this through Causal-EPIG, a framework adapting Expected Predictive Information Gain to quantify uncertainty reduction in causal quantities. We derive two distinct strategies: a comprehensive approach that robustly models the full causal mechanisms via the joint potential outcomes, and a focused approach that directly targets the CATE estimand for maximum sample efficiency. We provide theoretical justification for our framework, establishing a formal link between our information-theoretic objective and the minimization of CATE estimation error. Extensive experiments demonstrate that our strategies consistently outperform standard baselines, and crucially, reveal that the optimal strategy is context-dependent, contingent on the base estimator and data complexity. Our framework thus provides a principled guide for sample-efficient CATE estimation in practice.

Deep Learning · Generative Models and Autoencoders

Jannis Chemseddine, Gregor Kornhardt, Richard Duong, Gabriele Steidl

The default Gaussian latent in flow-based generative models poses challenges when learning certain distributions such as heavy-tailed ones. We introduce a general framework for learning data-adaptive latent distributions using one-dimensional quantile functions, optimized via the Wasserstein distance between noise and data. The quantile-based parameterization naturally adapts to both heavy-tailed and compactly supported distributions and shortens transport paths. Numerical results confirm the method’s flexibility and effectiveness achieved with negligible computational overhead.

General Machine Learning · Everything Else

Prayas Agrawal, Prateek Chanda, Ishita Khatri, Ganesh Ramakrishnan, Bamdev Mishra, Pratik Kumar Jawanpuria

Training Large Language Models (LLMs) on heterogeneous datasets requires optimizing domain representations to balance convergence speed and domain coverage. While recent methods reduce computational overhead by selecting high-quality data subsets, they typically perform selection independently per domain or rely on computationally expensive proxy models to determine continuous domain weights. In this paper, we propose a joint sample selection framework that learns better domain representations in the batch directly and formulate the selection objective as joint domain utility maximization subject to partition matroid constraints. This approach enforces domain-specific budgets while maximizing a validation-guided gradient matching utility across all domains simultaneously. Theoretically, we establish that this objective is weakly submodular, allowing us to employ a computationally efficient orthogonal matching pursuit algorithm with provable guarantees. Empirically, we demonstrate that our method significantly outperforms baselines that favor independent domain selection on mathematical reasoning and molecular generation benchmarks. Furthermore, our analysis show that our approach reduces the number of conflicting training gradient pairs significantly as compared to independent selection across domains or domain agnostic selection. Applied to Qwen2.5 and Llama-3 and trained on MethaMathQA and Mol-Instructions, our approach yields robust gains under multiple subset fractions, with improvements on nine math-reasoning and four molecule-generation benchmarks, highlighting cross-domain benefits of joint subset selection.

Social Aspects · Safety

Yudong Yang, Xuezhen Zhang, Zhifeng Han, Siyin Wang, Jimin Zhuang, Zengrui Jin, Jing Shao, Guangzhi Sun, Chao Zhang

Recent progress in large language models (LLMs) has enabled understanding of both speech and non-speech audio, but has also exposed new safety risks arising from complex audio inputs that are inadequately handled by current safeguards. We introduce SACRED-Bench (Speech–Audio Composition for RED-teaming) to evaluate the robustness of LLMs under complex audio-based attacks. Unlike existing perturbation-based methods that rely on noise optimization or white-box access, SACRED-Bench exploits speech–audio composition to enable effective black-box attacks. SACRED-Bench adopts three composition mechanisms: (a) speech overlap, (b) multi-speaker dialogue, and (c) mixtures of speech and non-speech audio. These mechanisms focus on evaluating safety in settings where benign and harmful intents co-occur within a single auditory scene. Moreover, questions in SACRED-Bench are designed to implicitly refer to content in the audio, such that no explicit harmful information appears in the text prompt alone. Experiments demonstrate that even Gemini 2.5 Pro, a state-of-the-art proprietary LLM with safety guardrails fully enabled, still exhibits a 66% attack success rate. To bridge this gap, we propose SALMONN-Guard, the first guard model that jointly inspects speech, audio, and text for safety judgments, reducing the attack success rate to 20\%. Our results highlight the need for audio-aware defenses to ensure the safety of multimodal LLMs.

Deep Learning · Large Language Models

Sayantan Dasgupta, Trevor Cohn, Tim Baldwin

The core learning signal used in language model distillation is the standard Kullback-Leibler (KL) divergence between the distribution of the student and the teacher. Traditional KL divergence tends to be dominated by the teacher’s highest-probability modes, thus diminishing the influence of less probable yet potentially informative components of the output distribution. We propose a new tail-aware divergence that decouples the contribution of the teacher model's top-$K$ predicted probabilities from that of lower-probability predictions, while maintaining the same computational profile as the KL Divergence. Our decoupled approach reduces the impact of the teacher modes and, consequently, increases the contribution of the tail of the distribution. Experimental results demonstrate that our modified distillation method yields competitive performance in both pre-training and supervised distillation of decoder models across various datasets. Furthermore, the distillation process is efficient and can be performed with a modest academic budget for large datasets, eliminating the need for industry-scale computing.

De Li, Zhou Tan, Qiyu Li, Zeming Gan, Tiange Xia, Chunpei Li, Xianxian Li

Federated Graph Learning (FGL) is a distributed machine learning paradigm based on graph neural networks, enabling secure and collaborative modeling of local graph data among clients. However, label noise in graph data can degrade the generalization performance of the global model. Existing federated label noise learning methods, primarily focused on computer vision tasks, often yield suboptimal results when directly applied to FGL. To address this issue, we propose a robust federated graph learning method with label noise, termed **FedRGL**. Specifically, FedRGL leverages the globally aggregated model and local subgraph structural information to implement a dual-perspective consistency noise-node filtering mechanism under class-aware dynamic thresholds. The resulting *class-aware dual-consistency filtering (CADF)* can also serve as a plug-and-play module, enhancing noise robustness across various subgraph federated learning frameworks. To better exploit the supervisory information from filtered noisy nodes, we employ the natural augmentation techniques from graph contrastive learning to assign high-confidence pseudo-labels to the noise nodes. Additionally, we measure model quality via the average predictive entropy of unlabeled nodes, enabling adaptive robust aggregation on the server side. Extensive experiments on real-world graph datasets show that FedRGL consistently outperforms existing methods under different noise rates, noise types, and client scales, achieving on average **5--8\%** higher accuracy and up to **30\%** improvement over the weakest baselines under noisy conditions. The anonymous source code is available at https://anonymous.4open.science/r/FedRGL_ICML26-376F.

Jake Fawkes, Jason Hartford

In GFlowNets and variational inference, it has been shown that the mean square error between target and model log probabilities is an effective, low variance, surrogate loss for training generative models. This loss has the property that when evaluated \emph{on-policy} its gradients correspond to those of the KL divergence, while \emph{off-policy} it remains a valid loss with the same global minimiser. In this work, we demonstrate that this construction can be extended to the whole family of $f$-divergences, leading to a family of losses whose on-policy gradients are that of the corresponding $f$-divergence, but retain the same global minimiser off-policy. Specifically, we show that the on-policy gradients lead to a one to one correspondence between translation invariant loss functions on the target and model log probabilities, and $f$-divergences. This equivalence allows us to design new surrogate loss functions for tuning a wide class of generative models that inherit the properties of the corresponding $f$-divergence, such as being more mode covering, whilst being applicable to off-policy data. We apply our losses on a range of tasks, including classic synthetic examples, SynFlowNets for molecule discovery, and asynchronous large language model (LLM) tuning, demonstrating that our models retain their predicted properties on- and off-policy and can be applied to a wide class of generative models.

François Bertholom, Khalid Oublal

Diffusion models achieve state-of-the-art performance in generative modeling but are limited by their reliance on Gaussian noise and the high computational cost of iterative sampling. Star-shaped diffusion addresses the former by introducing a non-Markovian forward process, yet this comes at the expense of temporal coherence in the reverse process. We propose a novel framework that resolves this trade-off by learning a Markovian projection of a star-shaped forward process, and its reversal. This design enables learning over a broad class of exponential models and recovers DDPM as a special case. It is particularly well-suited for knowledge distillation, allowing few-step or even single-step generation. Experiments demonstrate the effectiveness and flexibility of our approach across multiple generative tasks. Code is available at \url{https://anonymous.4open.science/w/MStar-Diffusion-B3EE/}.

Deep Learning · Large Language Models

Yunhao Fang, Weihao Yu, Shu Zhong, Qinghao Ye, Xuehan Xiong, Lai Wei

Long-sequence modeling faces a fundamental trade-off between the efficiency of compressive fixed-size memory in RNN-like models and the fidelity of lossless growing memory in attention-based Transformers. Inspired by the Multi-Store Model in cognitive science, we introduce a memory framework of artificial neural networks. Our method maintains a sliding window of the Transformer's KV cache as lossless short-term memory, while a learnable module termed Artificial Hippocampus Network (AHN) recurrently compresses out-of-window information into a fixed-size compact long-term memory. To validate this framework, we instantiate AHNs using modern RNN-like architectures, including Mamba2, DeltaNet, and GatedDeltaNet to augment open-weight base LLMs. We also propose an efficient self-distillation method where the base model' all parameters are frozen and only the parameters from AHNs are optimized. For inference, our method sets a default large sliding window size of 32k for attention, and AHNs activate only when the sequence length exceeds the 32k window, addressing the quadratic-complexity issue of attention that emerges at that scale. Extensive experiments on long-context benchmarks LV-Eval and InfiniteBench demonstrate that AHN-augmented models consistently outperform sliding window baselines and achieve performance comparable or even superior to full-attention models, while substantially reducing computational and memory requirements. For instance, augmenting the Qwen2.5-3B-Instruct with AHNs reduces inference FLOPs by 40.5% and memory cache by 74.0%, while improving its average score on LV-Eval (128k sequence length) from 4.41 to 5.88.

General Machine Learning · Causality

Ruicong Yao, Tim Verdonck, Mihaela van der Schaar, Jakob Raymaekers

Differentiable approaches to causal discovery have shown promise in learning DAG structures via continuous optimization, but their theoretical guarantees are largely restricted to models with homoscedastic noise or known noise distribution. In particular, existing methods based on mean squared error fail to identify the true DAG when noise distributions are non-Gaussian and vary in scale. In this paper, we address this gap in nonlinear additive noise models (ANMs) with arbitrary noise. Our approach extends NOTIME (Berrevoets et al. 2025) which minimizes an independence criterion among the residuals. We show that the global minimizer of the independence criterion corresponds to the true underlying DAG up to additional constant edges in general ANMs. To recover the exact structure, we introduce an adaptive group lasso penalty that regularizes entire columns of the first-layer weight matrix of an MLP, enabling the selective pruning of constant edges in a functionally meaningful way. Empirically, our method exhibits effective and stable performance across diverse noise types and variances, outperforming prior methods that lack identifiability guarantees in this setting.

Christoph Kern, Michael Kim, Angela Zhou

Estimating heterogeneous treatment effects is important to tailor treatments to those individuals who would most likely benefit. However, conditional average treatment effect predictors may often be trained on one population but possibly deployed on different, possibly unknown populations. We use methodology for learning multi-accurate predictors to post-process CATE T-learners (differenced regressions) to become robust to unknown covariate shifts at the time of deployment. The method works in general for pseudo-outcome regression, such as the DR-learner. We show how this approach can combine (large) confounded observational and (smaller) randomized datasets by learning a confounded predictor from the observational dataset, and auditing for multi-accuracy on the randomized controlled trial. We show improvements in bias and mean squared error in simulations with increasingly larger covariate shift, and on a semi-synthetic case study of a parallel large observational study and smaller randomized controlled experiment. Overall, we establish a connection between methods developed for multi-distribution learning and achieve appealing desiderata (e.g. external validity) in causal inference and machine learning.