ICML 2026 论文列表

Mixture-of-Experts (MoE) models are typically pre-trained with explicit load-balancing constraints to ensure statistically balanced expert routing. Despite this, we observe that even well-trained MoE models exhibit significantly imbalanced routing. This behavior is arguably natural—and even desirable—as imbalanced routing allows models to concentrate domain-specific knowledge within a subset of experts. Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices, but with a less-discussed assumption of balanced routing. Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures on overloaded devices during post-training or inference, where explicit load balancing is often inapplicable. We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones. This ensures that all devices complete their workloads within the minimum collective latency while respecting memory constraints. Across different model scales, LLEP achieves up to 5x speedup and 4x reduction in peak memory usage compared to standard EP. This enables faster and higher-throughput post-training and inference, with ~1.9x faster for gpt-oss-120b. We support our method with extensive theoretical analysis and comprehensive empirical evaluations, including ablation studies. These results illuminate key trade-offs and enable a principled framework for hardware-specific hyper-parameter tuning to achieve optimal performance.

Responsible investing aims to generate positive impact across Environment (E), Society (S), and Governance (G), and rating companies along these dimensions is now widespread, making ESG scores highly popular. Allocating retail capital with sustainability in mind could be transformational, yet it remains unclear how individual investors can do so in practice. Current ESG solutions cannot model high-dimensional, multi-modal time series capturing the joint evolution of ESG risks, financial returns, news, and sentiment, even though this domain requires jointly reasoning over distinct numerical signals where both numerical proximity and semantic type must be preserved. To bridge this gap, we introduce a novel domain-aware $\textbf{representation learning framework}$ that learns geometry-preserving representations for heterogeneous time series using value-aware tokens with block-wise $\textbf{orthogonal embeddings}$. To capture trajectory-level structure, we introduce $\textbf{FACET}$ tokens and train the model using a geometry-preserving loss. The resulting model jointly learns to forecast future values and to organize entities in a representation space that reflects their temporal evolution. Trained on ESG, returns, news, and sentiment, the domain-aware LLM learns a representation space that enables accurate ESG forecasting, trajectory-based grouping, and latent-space search for superior asset selection and downstream application like portfolio rebalancing

Hypergraph contrastive learning has shown strong potential for modeling high-order relations, but most existing methods implicitly assume that nodes within the same hyperedge are semantically homogeneous. In real-world hypergraphs, this assumption often breaks down: heterogeneous node--hyperedge interactions may introduce noisy signals during augmentation and message passing, leading to feature contamination and false-positive contrastive alignment. To address this issue, we propose H$^2$CL, a heterogeneity-aware hypergraph contrastive learning framework for robust representation learning. H$^2$CL first estimates node--hyperedge heterogeneity from input features and uses it to guide a Heterogeneity-Aware View Generator, which selectively masks high-heterogeneity incidences and constructs cleaner contrastive views. It further introduces a Heterogeneity-Aware Hypergraph Encoder that dynamically reweights information propagation in both node-to-hyperedge and hyperedge-to-node aggregation, enabling hyperedges to aggregate more homogeneous signals while suppressing heterogeneous noise. We also provide theoretical analysis showing that the encoder corresponds to a coordinate descent step for minimizing a heterogeneity-weighted Dirichlet energy. Extensive experiments on standard benchmarks and larger-scale hypergraphs demonstrate that H$^2$CL achieves competitive or superior performance compared with recent baselines, remains robust under structural noise, and learns reweighting patterns that effectively reduce heterogeneity during message passing. Our code is available at: https://github.com/sxu-yaokx/HHCL

Pretraining optimizers are tuned to produce the strongest possible base model, on the assumption that a stronger starting point yields a stronger model after subsequent changes like post-training and quantization. This overlooks the geometry of the base model which controls how much of the base model's capabilities survive subsequent parameter updates. We study three pretraining optimization approaches that bias optimization toward flatter minima: Sharpness-Aware Minimization (SAM), large learning rates, and shortened learning rate annealing periods. Across model sizes ranging from 20M to 150M parameters, we find that these interventions consistently improve downstream performance after post-training on five common datasets with up to 80\% less forgetting. These principles hold at scale: a short SAM mid-training phase applied to an existing OLMo-2-1B checkpoint reduces forgetting by 31\% after MetaMath post-training and by 40\% after 4-bit quantization.

Efficient deployment of large language models (LLMs) requires extreme quantization, forcing a critical trade-off between low-bit efficiency and performance. Residual binarization promises hardware-friendly, matmul-free inference by stacking binary ($\pm$1) layers, but is plagued by pathological feature co-adaptation. We identify a key failure mode, which we term inter-path adaptation: during Quantization-Aware Training (QAT), parallel residual binary paths learn redundant features, degrading the error-compensation structure and crippling the model's expressive capacity. While prior work relies on heuristic workarounds (e.g., path freezing) that limit model capacity, we propose RaBiT, a novel quantization framework that resolves co-adaptation by algorithmically enforcing a residual hierarchy. Its core mechanism sequentially derives each binary path from a single shared full-precision weight, ensuring each path corrects its predecessor's error. This process is stabilized by a robust initialization that prioritizes functional preservation over mere weight approximation. RaBiT redefines the 2-bit accuracy-efficiency frontier: it achieves state-of-the-art performance, rivals even hardware-intensive Vector Quantization (VQ) methods, and delivers a 4.49$\times$ inference speed-up over full-precision models on an RTX 4090.

Test-time adaptation (TTA) adapts pretrained models to test data on-the-fly. Current TTA methods have focused on what to adapt: lightweight domain-aware components (prompts, normalization statistics) updated with consistency-aware self-supervised losses. This work investigates the more fundamental yet underexplored optimization process, providing insights and guidelines on how to appropriately update models for TTA. By analyzing the optimization error during TTA, we identify a pivotal stability-plasticity trade-off: the model should adapt to novel distributions while retaining learned knowledge, which motivates our design of a CONfidence-and-Gradient-Aware scheduler (CONGA) to constrain model learning rate (LR) within an adaptive exploration interval. For each iteration, the lower bound encourages model exploration on informative confident samples, while the upper bound prevents aggressive overfitting to noisy optimization gradients. Based on our theoretical findings, an adaptation-progress-conditioned cosine decay function decides the specific LR within the interval. As an LR scheduler, CONGA is naturally applicable on existing TTA methods as a plug-in module, introducing little computation overheads. Extensive experiments and analysis demonstrate the superiority and validness of CONGA.

Large language models (LLMs) have shown promising potential in persuasion, but existing works on training LLM persuaders are still preliminary. Notably, while humans are skilled in modeling their opponent's thoughts and opinions proactively and dynamically, current LLMs struggle with such Theory of Mind (ToM) reasoning, resulting in limited diversity and opponent awareness. To address this limitation, we introduce Theory of Mind Augmented Persuader (**ToMAP**), a novel approach for building more flexible persuader agents by incorporating two theory of mind modules that enhance the persuader's awareness and analysis of the opponent's mental state. Specifically, we instruct the persuader to consider possible objections to the target claim, and train a module to predict the opponent’s agreement level on these objections. Experiments show that the ToMAP persuader, while containing only 3B parameters, outperforms much larger baselines, like GPT-4o, with a relative gain of 39.4% across multiple persuadee models and diverse corpora. Notably, ToMAP exhibits complex reasoning chains and reduced repetition during training, which leads to more diverse and effective arguments. These results underscore ToMAP's potential for developing more persuasive language agents. Code is available at: https://github.com/ulab-uiuc/ToMAP.

We introduce Hybrid Space-aware Stochastic Convolution Attention Noise (HySCAN), a hybrid randomized defense that helps close the long-standing gap between provable robustness under ℓ2 certificates and empirical robustness against strong ℓ∞ attacks, while maintaining strong generalization across diverse imaging benchmarks. HySCAN jointly explores complementary sources of stochasticity at both training and inference: (i) implicit weight-space randomness via stochastic-aware Random Weights, and (ii) explicit feature-space randomness via Stochastic Attention Noise Injection modules. By incorporating randomness at both the parameter and representation levels, HySCAN enables meaningful certified guarantees while improving empirical robustness in practice. Comprehensive experiments on diverse imaging datasets, e.g., CelebA, CIFAR-10, and CIFAR-100, ImageNet-1k, HAM10000, and NIH Chest X-ray, demonstrate that HySCAN outperforms existing certified and empirical defenses, improving certified robustness by up to ≈ 9.6% and empirical robustness by up to ≈ 5% without reducing clean accuracy.

This paper studies fairness-aware graph adaptation, aiming to transfer knowledge from a labeled source graph to an unlabeled target graph while addressing fairness. Most prior methods require target-domain attributes to learn invariant graph representations of sensitive attributes, which are often unavailable in practice. To address this limitation, we introduce \underline{C}ausality-attended Repres\underline{e}ntation Dientang\underline{l}ement with Structural A\underline{l}ignment (CELL) for fairness-aware graph adaptation without requiring target sensitive labels. CELL constructs a causal graph to model the graph-generation mechanism and guide fair representation disentanglement. Specifically, CELL uses sensitive and causal encoders to extract sensitive and causal factors, respectively, and promotes disentanglement by minimizing their conditional mutual information. To leverage unlabeled target data, we further generate pseudo-labels for both target task labels and sensitive attributes, and use similarity relations to derive unbiased node representations. Finally, to further mitigate domain shift, we build a fairness-aware bipartite graph that provides additional structural supervision for cross-domain alignment. Experiments on benchmarks show that CELL consistently outperforms strong baselines in both predictive performance and fairness.

Dataset distillation aims to synthesize a compact subset of the original data, enabling models trained on it to achieve performance comparable to those trained on the original large dataset. Existing distribution-matching methods are confined to Euclidean spaces, making them only capture linear structures and overlook the intrinsic geometry of real data, e.g., curvature. However, high-dimensional data often lie on low-dimensional manifolds, suggesting that dataset distillation should have the distilled data manifold aligned with the original data manifold. In this work, we propose a geometry-aware distribution-matching framework, called GeoDM, which operates in the Cartesian product of Euclidean, hyperbolic, and spherical manifolds, with flat, hierarchical, and cyclical structures all captured by a unified representation. To adapt to the underlying data geometry, we introduce learnable curvature and weight parameters for three kinds of geometries. At the same time, we design an optimal transport loss to enhance the distribution fidelity. Our theoretical analysis shows that the geometry-aware distribution matching in a product space yields a smaller generalization error bound than the Euclidean counterparts. Extensive experiments conducted on standard benchmarks demonstrate that our algorithm outperforms state-of-the-art data distillation methods and remains effective across various distribution-matching strategies for the single geometries.

Inference-time steering offers a promising way to control language models (LMs) without retraining. However, standard approaches typically rely on activation addition, which inevitably alters the hidden-state magnitudes raising concerns about representation collapse and degraded open-ended generation. In this work, we explore *Spherical Steering*, a training-free primitive that resolves this trade-off through activation rotation. Rather than shifting activations with a fixed vector, our method rotates them along a geodesic toward a target direction, preserving signal integrity while steering toward the target concept. To further enhance adaptivity, we incorporate a confidence gate that dynamically modulates steering strength based on input uncertainty. Extensive experiments across multiple-choice benchmarks demonstrate that Spherical Steering significantly outperforms addition-based baselines (notably by +10\% on TruthfulQA, COPA, and Storycloze), while simultaneously maintaining the model’s general open-ended generation quality. This work highlights the value of geometric consistency, suggesting that norm-preserving rotation is a robust and effective primitive for precise inference-time control. The code is available at: https://github.com/chili-lab/Spherical-Steering.

Asynchronous reinforcement learning (RL) accelerates LLM training through parallel data collection, but introduces *policy lag*: experiences are collected under stale policy weights, which destabilizes learning—especially on hard tasks. We identify why: **gradient variance scales exponentially with task difficulty under staleness**, because hard tasks have narrow solution spaces corresponding to sharp loss-landscape curvature (high Hessian eigenvalues). We formalize this as a *staleness budget optimization problem* and prove that the optimal allocation follows an exponential decay: $\eta^*(d) = \eta_{\text{base}} \cdot e^{-\lambda d}$ where $\lambda = \alpha/2$ is half the Hessian growth rate. Building on this principle, we propose **GAS (Gradient-Aware Scheduling)**, a drop-in recipe for fast and stable asynchronous RL with three components: (i) a bandit-based curriculum that selects task difficulty by learning signal (ACB); (ii) execution-aware staleness budgets that keep slow tasks from bottlenecking rollout (EAAS); and (iii) curriculum–staleness coupling that tightens the staleness budget as difficulty rises (CSC). Our mechanistic analysis validates the theoretical predictions: the "safe zone" of gradient coherence follows the derived exponential boundary. On code generation benchmarks, GAS improves Pass@1 from 39.7% to 60.1% while training at $2.3\times$ the throughput of synchronous GRPO, showing that matching each task's staleness budget to its difficulty (and thus to its loss-landscape curvature) makes asynchronous RL both fast and stable.

While prior studies have successfully compressed vision Transformers (ViTs) through various pruning techniques, most have concentrated on width pruning to achieve significant reductions in model size. Depth pruning, which removes entire layers from a ViT, is notoriously difficult for accuracy recovery despite its potential to deliver higher speedups, limiting the acceleration achieved by existing joint width-and-depth pruning methods. In this work, we reveal that the failure of existing depth pruning methods lies in their neglect of heterogeneity between different layers, and we introduce HetDPT, a heterogeneity-aware depth pruning method that avoids dimension mismatch. Comprehensive experiments on ImageNet-1K, CIFAR-100, COCO, and ADE20K validate our method: HetDPT achieves a 1.58$\times$ speedup for DeiT-B while maintaining accuracy and a 1.39$\times$ speedup for DeiT-S with nearly no accuracy degradation. Furthermore, when combined with width pruning, HetDPT+ sets a new state-of-the-art record in extreme ViT pruning, enhancing the acceleration ratio from 4.24$\times$ to 5.19$\times$ for the Isomorphic-Pruning-2.6G configuration while maintaining near-lossless accuracy; our code is available at \url{https://github.com/Efficient-AI-for-All/HetDPT}.

Sharpness-Aware Minimization (SAM) is widely used to seek flatter minima—often linked to better generalization. In its standard implementation, SAM updates the current iterate using the loss gradient evaluated at a point perturbed by distance $\rho$ along the normalized gradient direction. We show that, for some choices of $\rho$, SAM can stall at points where this shifted (perturbed-point) gradient vanishes despite a nonzero original gradient, and therefore, they are not stationary points of the original loss. We call these points hallucinated minimizers, prove their existence under simple nonconvex landscape conditions (e.g., the presence of a local minimizer and a local maximizer), and establish sufficient conditions for local convergence of the SAM iterates to them. We corroborate this failure mode in neural network training and observe that it aligns with SAM's performance degradation often seen at large $\rho$. Finally, as a practical safeguard, we find that a short initial SGD warm-start before enabling SAM mitigates this failure mode and reduces sensitivity to the choice of $\rho$.

The rapid growth of deep neural networks (DNNs) has intensified the demand for efficient hardware acceleration under quantization. While prior research has successfully reduced weight and activation precision, partial sums generated during accumulation often retain high precision, resulting in significant energy overhead. In this work, we analyze psum distributions in tiled architectures and reveal that within-tile outliers are input-dependent. We propose PsumQuant, a post-training, input-aware quantization that predicts psum scales on-the-fly. By leveraging the crest factor of input activations, our learnable scale predictor effectively bounds the psum bit-width while handling the extreme outliers in DNNs. Experimental results on a $128 \times 128$ systolic array demonstrate that PsumQuant compresses psum precision down to 8-bit within only a 1\% accuracy drop on ResNet-18 and a marginal 0.04 perplexity increase on Llama-3.1. Furthermore, bit-width reduction with PsumQuant results in a 45\% reduction in total energy with minimal accuracy loss, demonstrating that PsumQuant provides a highly efficient solution for actual NPU architectures.

Quantization-aware training (QAT) is an effective method to drastically reduce the memory footprint of LLMs while keeping performance degradation at an acceptable level. However, the optimal choice of quantization format and bit-width presents a challenge in practice. The full design space of quantization is not fully explored in the context of QAT, and the precise trade-off between quantization and downstream performance is poorly understood, as comparisons often rely solely on perplexity-based evaluations. In this work, we address these shortcomings with an empirical study of QAT in the low-bit regime. We show that k-means based weight quantization outperforms integer formats and can be implemented efficiently on standard hardware. Furthermore, we find that, under a fixed inference memory budget, the best performance on generative downstream tasks is achieved with $1$-bit quantized weights.

Identifying high-utility candidates from massive discrete spaces under expensive evaluations is a recurring challenge across the sciences, with structure-based drug discovery as a prominent example. While surrogate-based optimization can increase sample efficiency by reducing the number of expensive evaluations, modern molecular libraries have reached billions to trillions of compounds, making full-library surrogate inference itself a major computational bottleneck. We introduce BOBA, a bandit-guided surrogate optimization framework that eliminates full-library inference by adaptively allocating computation across partitions of the action space. By treating partitions as arms in a multi-armed bandit, BOBA concentrates inference and evaluations on empirically promising partitions while maintaining principled exploration. Experiments on real-world synthesis-on-demand libraries demonstrate that optimism-under-uncertainty bandits, combined with meaningful action space partitioning, are essential for effective allocation of inference and evaluations. Our findings reveal a tunable tradeoff between screening performance and surrogate inference cost, which supports practical optimization over current libraries, and establishes a viable route to ultra-large library virtual screening.

Compute-In-Memory (CIM) accelerators are promising for energy-efficient edge inference, yet they faces fundamental challenges when deploying Deep Neural Networks (DNNs), as hardware-induced weight perturbations from intrinsic noise and device drift degrade accuracy and impede reliable inference. To tackle this challenge, we propose Diversity-aware Weight Perturbation (DWP), an immune-system-inspired training method that emulates affinity-based selection by exploiting sample-level prediction disagreement under diverse noise realizations to guide adaptive sample weighting, building robustness to weight perturbation. Experiments show that DWP-trained models consistently yield superior robustness, achieving over 15\% accuracy improvements compared to standard-trained models under severe weight perturbations (mismatch level up to 70\%) and maintaining inference accuracy at 90\% over a simulated one-year CIM operation with only 2\%–4\% variation in accuracy. Moreover, under matched model and inference configurations, deployment on low-precision CIM hardware reduces inference energy by 38\% compared to a GPU baseline. These results demonstrate that DWP enables robust and energy-efficient neural network deployment on resource-constrained edge devices with inherent hardware uncertainties.

Text-to-image diffusion models remain computationally intensive: generating a single image typically requires dozens of passes through large transformer backbones (e.g., SDXL uses ~50 evaluations of a 2.6B-parameter model). Few-step variants reduce the step count to 2–8 but still rely on large, full-precision backbones, making inference impractical on resource-constrained platforms. Existing post-training quantization (PTQ) methods are further hampered by their dependence on full-precision calibration. We introduce **Q-Sched**, a scheduler-level PTQ approach that adapts the diffusion sampler while keeping the quantized weights fixed. By adjusting the few-step sampling trajectory with quantization-aware preconditioning coefficients, Q-Sched matches or surpasses full-precision quality while delivering a **4× reduction** in model size and preserving a single reusable checkpoint across bit-widths. To learn these coefficients, we propose a reference-free **Joint Alignment–Quality (JAQ)** loss, which combines text–image compatibility with an image-quality objective for fine-grained control. JAQ requires only a handful of calibration prompts and avoids any full-precision inference during calibration. Empirically, Q-Sched yields substantial gains: a **15.5%** FID improvement over the FP16 4-step Latent Consistency Model and a **16.6%** improvement over the FP16 8-step Phased Consistency Model, demonstrating that quantization and few-step distillation are complementary for high-fidelity generation. A large-scale user study with **80,000+** annotations further validates these results on both FLUX.1[schnell] and SDXL-Turbo. Code: [https://github.com/enyac-group/q-sched](https://github.com/enyac-group/q-sched)