ICML 2026 Papers

Simulating trajectories of dynamical systems is a fundamental problem in a wide range of fields such as molecular dynamics, biochemistry, and pedestrian dynamics. Machine learning has become an invaluable tool for scaling physics-based simulators and developing models directly from experimental data. In particular, recent advances in deep generative modeling and geometric deep learning enable probabilistic simulation by learning complex trajectory distributions while respecting intrinsic permutation and time-shift symmetries. However, trajectories of N-body systems are commonly characterized by high sensitivity to perturbations leading to bifurcations, as well as multi-scale temporal and spatial correlations. To address these challenges, we introduce STFlow (Spatio-Temporal Flow), a generative model based on graph neural networks and hierarchical convolutions. By incorporating data-dependent couplings within the Flow Matching framework, STFlow denoises starting from conditioned random-walks instead of Gaussian noise. This novel informed prior simplifies the learning task by reducing transport cost, increasing training and inference efficiency. We validate our approach on N-body systems, molecular dynamics, and human trajectory forecasting. Across these benchmarks, STFlow achieves the lowest prediction errors with fewer simulation steps and improved scalability.

Accurate survival modeling in biomedical studies is often hindered by rare events, limited effective sample sizes, and settings with limited or partially observed information (e.g., covariates of interest that are difficult or expensive to collect, highly structured sampling designs, or nuisance parameters omitted by conditioning). Knowledge distillation can leverage external predictive information without sharing individual-level data, but existing approaches are largely built for fully specified likelihoods or probability-based survival models and do not extend to settings where outcome distributions are only partially specified. To address this challenge, we propose a knowledge distillation framework based on a composite-likelihood Kullback--Leibler divergence that aligns teacher and student models within components. Our key insight is that, although composite likelihoods do not define a global outcome distribution, each likelihood component induces a well-defined probability model on its restricted outcome space, enabling a principled KL divergence. Simulation studies and biomedical case studies show improved discrimination and predictive accuracy in rare-event, heterogeneous settings without requiring access to external individual-level data.

Standard mixed-precision training of neural networks requires many bytes of accelerator memory for each model parameter. These bytes reflect not just the parameter itself, but also its gradient and one or more optimizer state variables. With each of these values typically requiring 4 bytes, training even a 7 billion parameter model can be impractical for researchers with less than 100\,GiB of accelerator memory. We introduce FlashOptim, a suite of optimizations that reduces per-parameter memory by over 50% while preserving model quality and API compatibility. Our approach introduces two key techniques. First, we improve master weight splitting by finding and exploiting a tight bound on its quantization error. Second, we design companding functions that greatly reduce the error in 8-bit optimizer state quantization. Together with 16-bit gradients, these techniques reduce AdamW memory from 16 bytes to 7 bytes per parameter, or 5 bytes with gradient release. They also cut model checkpoint sizes by more than half. Experiments with FlashOptim applied to SGD, AdamW, and Lion show no measurable quality degradation across a collection of standard vision and language benchmarks, including Llama-3.1-8B finetuning.

End-to-end online map generation is a core component of autonomous driving perception systems. However, in complex traffic scenes, Bird’s-Eye-View (BEV) with vectorized mapping suffers from limitations such as target misclassification, spatial localization drift, and ambiguous semantic segmentation. Introducing uncertainty quantification can alleviate these problems, so we propose MapUQ, a robust BEV vectorized mapping method guided by uncertainty-aware optimization. Specifically, we quantify uncertainty at the feature level to enhance semantic perception, apply an error-driven dynamic receptive field adaptation mechanism at the decoding stage to enforce geometric consistency, and leverage negative sample information at the output head to improve lane classification accuracy. Experimental results on the nuScenes and Argoverse 2 datasets show that our method outperforms prior approaches in AP across three road types, achieving an average improvement of 1.5% over the baseline with marginal computational overhead. In addition, our method surpasses the baseline on uncertainty metrics such as ECE and NLL, significantly improving robustness and mapping accuracy in complex scenarios. Our code has been released at github: https://anonymous.4open.science/r/MapUQ-D287.

Policy mirror descent (PMD) provides a principled framework for reinforcement learning (RL) by iteratively solving KL-regularized policy improvement subproblems. While this approach has been adopted in training advanced LLMs such as Kimi K1.5/K2, the ideal closed-form PMD updates require reliable partition function estimation, a significant challenge when working with limited rollouts in the vast action spaces of LLMs. We investigate a practical algorithm, termed PMD-mean, that approximates the log-partition term with the mean reward under the sampling policy and performs regression in log-policy space. Specifically, we characterize the population solution of PMD-mean and demonstrate that it implicitly optimizes mirror descent subproblems with an adaptive mixed KL--$\chi^2$ regularizer. This additional $\chi^2$ regularization constrains large probability changes, producing more conservative updates when expected rewards are low and enhancing robustness against finite-sample estimation errors. Experiments on math reasoning tasks show that PMD-mean improves stability and remains competitive in the stale-rollout settings while enabling higher throughput through larger rollout batches. These findings deepen our understanding of PMD-mean and illuminate pathways toward principled improvements in RL algorithms for LLMs. Code is available at https://github.com/horizon-llm/OpenKimi.

The performance of sparse direct solvers is fundamentally governed by fill-in, i.e. new nonzero entries arising from the LU factorization of a sparse matrix, as they dictate memory footprint and subsequent computation time. For decades, a variety of graph-theoretic algorithms have aimed to minimize fill-in, a problem known to be both NP-hard and critically important. While recent deep learning methods, optimizing surrogate fill-in objectives, show empirical promise and can outperform classical algorithms on certain matrices, they offer limited interpretability into the underlying mechanism of fill-in generation. To address this, we propose a novel reordering approach, Causal Triplet Structure Learning (CTS), which is grounded in the Fill-Path Theorem and reduces arbitrary-length fill-paths to length-two candidate triplets, identifies the causal structures that trigger fill-in, and intervenes to block their formation. Empirically, we design a multigrid-style GAT with KAN activations to learn vertex embeddings and introduce a causal triplet loss that discourages such structures during training. Experiments on the SuiteSparse Matrix Collection demonstrate that our method reduces fill-in by 6$\times$, leading to 12$\times$ speedup in factorization time compared to state-of-the-art methods on Chemical Process Simulation and Computational Fluid Dynamics matrices.

Continual learning enables neural networks to learn tasks sequentially without forgetting previously acquired knowledge. However, neural networks suffer from catastrophic forgetting, where learning new tasks degrades performance on earlier ones. We address this problem with **Shapley Neuron Valuation (SNV)**, a principled framework that quantifies Neuron importance in continual learning, grounded in cooperative game theory. SNV selectively freezes important Neurons while keeping others plastic, enabling **buffer-free** continual learning without expanding architecture. Experiments on ImageNet-1k show that SNV consistently outperforms existing buffer-free methods. In particular, SNV improves accuracy by **+2.88\%** in the class incremental learning and **+6.46\%** in the task incremental learning scenarios compared to the second baseline.

Latent Factorization-of-Tensors (LFT) models are an effective approach for relation analysis. Existing LFT models assume each mode of the target tensor corresponds to an entity set and the relationships between entity sets are bipartite graphs to explore the relationships among entities within a mode. However, when the topological structure of entities in a mode is known, for example, entities are ordered physical quantities, such as time or coordinates, the relations between such modes forms a more complicated structure, i.e., aligned bipartite networks, and existing LFT models cannot accurately capture this structure. This work is the first to recognize and analyze this issue, and proposes an Asymmetric Latent Factorization-of-Tensors (ALFT) model to address it. ALFT can model aligned bipartite networks in mode pairs of a tensor by imposing constraints between particular mode pairs in the tensor network. Experimental results on real-world datasets demonstrate the existence of this issue and confirm that the proposed ALFT model can effectively resolve it.

Policies learned via continuous actor-critic methods often exhibit erratic, high-frequency oscillations, making them unsuitable for physical deployment. Current approaches attempt to enforce smoothness by directly regularizing the policy's output. We argue that this approach treats the symptom rather than the cause. In this work, we theoretically establish that policy non-smoothness is fundamentally governed by the differential geometry of the critic. By applying implicit differentiation to the actor-critic objective, we prove that the sensitivity of the optimal policy is bounded by the ratio of the Q-function's mixed-partial derivative (noise sensitivity) to its action-space curvature (signal distinctness). To empirically validate this theoretical insight, we introduce PAVE (Policy-Aware Value-field Equalization), a critic-centric regularization framework that treats the critic as a scalar field and stabilizes its induced action-gradient field. PAVE rectifies the learning signal by minimizing the Q-gradient volatility while preserving local curvature. Experimental results demonstrate that PAVE achieves smoothness comparable to policy-side smoothness regularization methods, while maintaining competitive task performance, without modifying the actor.

A natural recipe for intelligent robotic decision-making is initializing from pretrained generative control policies, which have summarized offline experience, and adapting them to self-collected online experience. We present DF-ExpEnse, an exploration technique that improves the quality of online experience collection, thus increasing finetuning sample-efficiency. DF-ExpEnse leverages the multimodal modeling capabilities of the generative control policy to create an expressive and tractably evaluatable candidate set. It then utilizes an ensemble of critics to identify the action that best balances quality with high exploration interest. In fleet settings, DF-ExpEnse further enables cross-agent communication to facilitate collaborative exploration as a group. DF-ExpEnse can be seamlessly integrated with existing strategies that finetune pretrained generative control policies via reinforcement learning. We experimentally validate consistent sample-efficiency benefits through DF-ExpEnse across a variety of manipulation and locomotion tasks, compared to default finetuning and alternative action selection schemes. Project can be found at [df-expense.github.io](https://df-expense.github.io).

Multimodal alignment can fail under deployment shift because standard objectives entangle cross-modal dependence with marginal-sensitive geometry. Stable Spectral Copula Alignment (SSCA) provides a deployment protocol for copula-stable dependence under approximately coordinate-wise monotone marginal distortions, together with auditable, label-free diagnostics for monitoring and mitigation. SSCA combines (i) clipped soft-rank Gaussianization that suppresses marginal effects while tracking tie and approximation errors, (ii) dependence-weighted sliced Wasserstein hub coupling for globally coherent multiway alignment with cycle auditing, and (iii) diagonal-stabilized block-spectral learning with eigengap-normalized Davis-Kahan diagnostics, yielding an actionable subspace-risk inequality. A calibrated gate maps diagnostic proxies to a reliability signal with a measurable false-alarm/miss trade-off, enabling stability-mode updates, budgeted remediation, and conservative no-update fallback for out-of-scope drift. Evaluations on MOSEI/MELD, MSCOCO, and CC3M-500K show improved performance under perturbation and substantially reduced degradation under controlled monotone distortions, raw-pipeline drifts, and frozen-feature retrieval stress tests.

Spurious correlations that lead models to correct predictions for the wrong reasons pose a critical challenge for robust real-world generalization. Existing research attributes this issue to group imbalance and addresses it by maximizing group-balanced or worst-group accuracy, which heavily relies on expensive bias annotations. A compromise approach involves predicting bias information using extensively pretrained foundation models, which requires large-scale data and is limited to physically interpretable biases. To address these challenges, we offer a novel perspective by reframing the spurious correlations as imbalances/mismatches in class-conditional distributions caused by general biases, whether interpretable or not, and propose a simple yet effective robust learning method that mitigates the need for bias annotations or predictions. With the goal of maximizing the conditional entropy (uncertainty) of the label given spurious factors, our method leverages a sample reweighting strategy to achieve class-conditional distribution balancing, which automatically highlights minority groups and classes, effectively mitigating spurious correlations and producing a debiased data distribution for classification. Extensive experiments and analysis demonstrate that our approach consistently delivers strong overall performance, and is competitive with methods that rely on bias supervision.

High-fidelity 3D generation depends on 3D VAEs that can compress and reconstruct complex geometry at high resolution. Recent methods often convert raw meshes into signed distance fields (SDFs), but this preprocessing can be lossy, especially for open or non-watertight assets. Render-supervised VAEs such as TripoSF avoid this conversion by matching rendered depth and normal maps, yet render losses supervise only the visible geometry for each view, leaving many latent voxels weakly constrained while still incurring unnecessary decoding and attention costs. We introduce Focusing, a view-consistent sparse-voxel training scheme for efficient 3D VAEs. Given a training view, Focusing performs depth-driven voxel carving directly in the structured latent space: voxels inconsistent with the rendered depth are removed before decoding, allowing the decoder and attention layers to operate only on locally relevant geometry. This view-dependent sparsity reduces memory and computation while concentrating learning on surface regions that contribute to the render. To stabilize training across shapes and viewpoints, we further propose adaptive zooming, which adjusts camera intrinsics to keep the number of active voxels within a target range and strengths supervision for fine details. The VAE is trained with render-based depth, normal, mask, and perceptual losses, together with sparse-voxel total variation and a brief TSDF warm-up to improve convergence and suppress holes. Across standard reconstruction benchmarks, Focusing improves Chamfer Distance and F-score over strong baselines while substantially reducing video random access memory (VRAM) consumption, enabling $1024^3$-resolution VAE training with as little as 50 GB of VRAM. These results demonstrate that local, view-consistent sparsity is an effective path toward higher-resolution and more efficient 3D VAE training.

Reward-maximizing RL methods have shown to be capable of enhancing the reasoning performance of LLMs, but often lead to reduced generation diversity. Recent works address this issue by adopting GFlowNets, training LLMs to match a target distribution while jointly learning its partition function. In contrast to prior works that treat this partition function solely as a normalizer, we reinterpret it as a per-prompt expected-reward (i.e., online accuracy) signal, leveraging this unused information to improve sample efficiency. Specifically, we first establish a theoretical relationship between the partition function and per-prompt accuracy estimates. Building on this key insight, we propose \textbf{Pa}rtition Fun\textbf{c}tion-Guid\textbf{ed} \textbf{RL} (PACED-RL), a post-training framework that leverages accuracy estimates to prioritize informative question prompts during training, and further improves sample efficiency through an accuracy estimate error–prioritized replay. Crucially, both components reuse information already produced during GFlowNet training, effectively amortizing the compute overhead into the existing optimization process. Extensive experiments across diverse benchmarks demonstrate strong performance improvements over GRPO and prior GFlowNet approaches, highlighting PACED-RL as a promising direction for a more sample efficient distribution-matching training for LLMs.

We propose a scalable method for training prediction (machine learning) models in the predict-then-optimize paradigm, where model outputs serve as coefficients for a subsequent linear optimization task. Directly minimizing the empirical decision regret is intractable for linear programming and combinatorial optimization since the decision mapping is piecewise constant, and the gradients are zero almost everywhere. While existing methods address this by smoothing the differentiation process, they suffer from scalability issues, since a computationally expensive solver call is required for every gradient evaluation. To address this, we propose a decision-focused learning pipeline based on a measure transformation principle, which yields a new surrogate loss that is completely optimization-solver-free during training. We establish theoretical guarantees, including Fisher consistency and excess risk bounds. Empirically, our method achieves decision quality competitive with state-of-the-art methods while reducing training time by orders of magnitude.

Graph Convolutional Networks (GCNs) are defined based on aggregating the information of adjacent nodes, that are usually treated as equally important and may limit the representational power of existing GCNs. To address this shortcoming, we propose a novel Global Interacted Graph Convolutional Network (GI-GCN), that leverages the solution vectors maintained during the iterative updates of the Dominant Set to adaptively characterize the global importance distribution over all nodes. Specifically, at each convolution layer, this distribution is adopted to adaptively modulate the importance weights of node features before performing the local message passing. We show that this convolution strategy can effectively capture the highly correlated information between nonadjacent nodes through the Dominant Set algorithm, not only emphasizing the critical graph-level information but also enhancing the discriminative power of graph representations. Furthermore, we optimize the memory complexity of the framework, significantly reducing the memory overhead associated with the global interaction modeling. Experiments demonstrate the effectiveness of the proposed GI-GCN model.

Online learning for the Cox model is challenging because its partial likelihood is non-decomposable, with each risk set requiring a summation over all samples. We propose Collaborative Operation of Linked Survival Analysis (COLSA), an online learning framework that replaces the partial likelihood with the full likelihood using sieve approximation of the baseline hazard. The proposed likelihood function is decomposable and eliminates the need to store historical data in memory, enabling efficient online updates. Moreover, COLSA maintains sufficient statistics for a higher-order basis and employs data-driven basis projection to adaptively scale model complexity to the effective sample size. Unlike existing online Cox methods, COLSA achieves asymptotic normality and attains the same statistical efficiency as the pooled-data partial likelihood estimator, without accessing full data and only requiring constant memory. Simulation studies and application to kidney transplant data demonstrate that COLSA outperforms existing online methods and matches the performance of full-data estimation.

We introduce Flower, a neural architecture for learning PDE solution operators built entirely from multihead warps. Aside from pointwise channel mixing and a multiscale scaffold, Flowers use no Fourier multipliers, no dot-product attention, and no convolutional mixing. Each head predicts a displacement field and warps the mixed input features. Motivated by physics and computational efficiency, displacements are predicted pointwise, without any spatial aggregation, and nonlocality enters *only* through sparse sampling at source coordinates, *one* per head. Stacking warps in multiscale residual blocks yields Flowers, which implement adaptive, global interactions at linear cost. We theoretically motivate this design through three complementary lenses: flow maps for conservation laws, waves in inhomogeneous media, and a kinetic-theoretic continuum limit. Flowers achieve excellent performance on a broad suite of 2D and 3D time-dependent PDE benchmarks, particularly flows and waves. A compact 17M-parameter model consistently outperforms Fourier, convolution, and attention-based baselines of similar size, while a 150M-parameter variant improves over recent transformer-based foundation models with much more parameters, data, and training compute.

Spiking Neural Networks (SNNs) enable low-power, event-driven computation, but Transformer-based SNNs remain difficult to deploy on neuromorphic hardware due to dense operations and communication overhead. We propose two simplified attention mechanisms, \textbf{Simplified Spiking Attention (SSA)} and \textbf{Ultra-Simplified Spiking Attention (USSA)}, which replace matrix multiplication with Hadamard products and eliminate hardware-unfriendly components such as multi-head attention and scaling. We show that consecutive masking is redundant and analyze a spiking-order effect in which early spiking contributes more temporal information to attention modulation. On CIFAR-10, CIFAR-100, and DVS-Gesture, SSA achieves accuracies of 96.38\%, 79.45\%, and 97.56\%, respectively, while reducing computational complexity from $\mathcal{O}(N^2D)$ to $\mathcal{O}(ND)$ and communication complexity from $\mathcal{O}(N^2)$ to $\mathcal{O}(ND)$. USSA further reduces communication complexity to $\mathcal{O}(N)$ with only marginal accuracy degradation. On ImageNet-1K, SSA and USSA achieve 76.91\% and 77.27\% accuracy, respectively, demonstrating scalability to large-scale classification.