ICML 2026 Papers

International Conference on Machine Learning

Decision breakdown

6,341papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

Browse all papers

FAIR-Calib: Frontier-Aware Instability-Reweighted Calibration for Post-Training Quantization of Diffusion Large Language Models

Haoyu Huang, Linlin Yang, Sheng Xu, Boyu Liu, Guodong Guo, Zhongqian Fu, Hang Zhou, Baochang Zhang

Expand abstract

Diffusion Large Language Models (dLLMs) refine tokens iteratively but commit them irreversibly, leading to a "stability lag" where early decisions remain fragile even after being written. We reveal that Post-Training Quantization (PTQ) error easily flips these borderline decisions at the write frontier, which are then permanently locked in and amplified. To address this, we propose Frontier-Aware Instability-Reweighted Calibration (FAIR-Calib), a two-stage PTQ framework for dLLMs. Stage I probes a full-precision teacher to estimate a position prior that combines frontier hits and masked-stage reliability. Stage II performs off-policy, layer-wise calibration by minimizing a reweighted hidden-state MSE, effectively prioritizing the protection of fragile frontier states without requiring expensive end-to-end diffusion rollouts. We further theoretically justify our weighted objective as a surrogate for output KL divergence. Empirically, FAIR-Calib consistently outperforms state-of-the-art baselines on LLaDA and Dream (W4A4), significantly reducing frontier decision flips and suppressing post-commit mismatches across diverse benchmarks.

Poster
HTAC: Hierarchical Task-Aware Composition for Continual Offline Reinforcement Learning

Qiyang Zhou, Xu Ruihang, Peng Wang, WenJie Lu, Xiaochun Cao, Naiqiang Tan, Li Shen

Expand abstract

Continual Offline Reinforcement Learning (CORL) enables building long-term autonomous agents from static datasets. However, it faces heterogeneity in environment dynamics, reward functions, and behavior policies across tasks. Combined with the inherent distribution shift in offline learning, this requires agents to selectively reuse shared knowledge during transfer while isolating task-specific features. The flat knowledge sharing mechanisms employed by existing methods struggle to capture such distinctions, limiting cross-task generalization. To address this, we propose Hierarchical Task-Aware Composition (HTAC), which balances plasticity and stability through dual-level task encoding and soft composition mechanisms. HTAC comprises four modules: (1) a Hierarchical Semantic Task Representation that decomposes tasks into domain-level and task-level embeddings; (2) a Dual-level Expert Network that creates domain and task experts on demand for parameter-efficient knowledge isolation; (3) an Adaptive Knowledge Composition module that integrates historical expert outputs via attention mechanisms for knowledge reuse; (4) Task Adapters that preserve historical routing weights to prevent forgetting. Experiments on Offline Continual World show that HTAC outperforms existing baselines, demonstrating better knowledge reuse and transfer capabilities.

Poster
Merge to Remember: Sharpness-Aware Isotropic Merging for Continual Learning

Qun Yang, Enneng Yang, Wei Chen, Li Shen, Long Lan

Expand abstract

Continual learning with large pre-trained models offers significant potential for cross-task knowledge accumulation, but faces critical challenges such as catastrophic forgetting and parameter interference, especially when historical data is unavailable. Existing approaches typically rely on sequential fine-tuning or model merging strategies, yet often overlook the impact of loss landscape sharpness and dominant singular value directions, which leads to subspace misalignment and severe knowledge forgetting. In this paper, we propose the Sharpness-Aware Isotropic Merging (SAIM) framework, which introduces targeted optimizations in both the fine-tuning and merging stages to address these issues. Specifically, SAIM consists of two synergistic modules: (1) a Sharpness-Aware Block Coordinate Descent (SA-BCD) optimizer that guides the model toward flatter minima and selectively updates the most task-sensitive parameters, thereby mitigating parameter interference and enhancing robustness; (2) an adaptive isotropic merging algorithm that dynamically balances the singular value spectrum across tasks, effectively preventing the model from overemphasizing any single task direction, maintaining balanced knowledge representation, and improving subspace alignment. Extensive experiments on vision and language benchmarks demonstrate that SAIM achieves 5-10\% higher accuracy than existing methods and maintains robust performance as the number of tasks increases. Our code are available at \url{https://github.com/Yangqun123456/SAIM}.

Poster
QiMeng-PerceptOS: Semantic-Aware Kernel Optimization for OS-Intensive Workloads via Hardware-Software Alignment

Huilai Chen, Yuanbo Wen, Liangfeng Li, Shaohui Peng, Jingzhe Zhu, Jun Bi, Xuzhi Zhang, Qi Guo, Ling Li, Yunji Chen

Expand abstract

Optimizing OS kernels for specific applications is vital for peak performance, yet existing LLM-based methods struggle with a semantic mismatch between generalized reasoning and low-level system behaviors. As a result, these static, open-loop approaches suffer from runtime blindness, configuration fragmentation, and search drift, ultimately failing to unlock the system’s performance potential. To address this, we propose QiMeng-PerceptOS, an autonomous framework that shifts the paradigm to perception-driven tuning. QiMeng-PerceptOS integrates: (1) a Perception Module that aligns raw telemetry into high-fidelity semantic fingerprints; (2) a Global Search Module utilizing a Bi-level Hierarchical Induction Tree (BHIT) for global navigation and efficient pruning; and (3) a Posterior Enhancement Module to suppress hallucinations via trajectory synthesis. Experiments across diverse workloads show that it achieves significant performance breakthroughs by optimizing kernel configurations, reaching 296.6% of default Redis throughput and surpassing SOTA baselines by 32.6% within only 15 iterations. By establishing a perception-driven closed-loop, QiMeng-PerceptOS provides new insights for fully automated, large-scale system optimization.

Poster
Unlocking Cross-Modal Biosignal Synthesis: A Temporally-Aware VAE-Diffusion Model

Chenyang Xu, Dezhen Wang, Hao Wang

Expand abstract

Synthesizing authentic phonocardiograms (PCG) from ubiquitous electrocardiograms (ECG) is a critical task for accessible cardiac monitoring. Existing generative models, however, struggle to capture the heart's complex electromechanical coupling, failing to meet the dual requirements of temporal precision and physiological fidelity needed for clinically relevant waveform analysis. We introduce the Temporally-Aware VAE-Diffusion model, a synergistic hybrid architecture that resolves this trade-off. Our architecture enforces tight physiological coupling through an Enhanced Condition Fusion mechanism and explicitly models long-range cardiac dynamics via Temporal Attention Blocks. On the EPHNOGRAM benchmark, our model sets a new state-of-the-art, achieving a Pearson correlation of 0.910$\pm$0.008, 95.95\% S1 detection accuracy, and a precise 12.0 ms timing error, significantly outperforming leading diffusion and Transformer baselines. Crucially, our work provides a reproducible zero-shot transfer evaluation for ECG-to-PCG synthesis. Evaluated on the synchronized PhysioNet/CinC 2016 training-a/MITHSDB subset without target-domain training, our model preserves high waveform fidelity and clinically relevant timing structure under domain shift, including on pathological recordings. These results support cross-dataset robustness of the proposed synthesis framework, while downstream diagnostic validation remains an important direction for future work.

Poster
Causal-aware Anomaly Detection for Tabular Data

Dang Nguyen, Tu Anh Hoang Nguyen, Thuc Duy Le, Svetha Venkatesh, Trung Le, Sunil Gupta

Expand abstract

Anomaly detection (AD) methods often ignore causal dependencies and treat anomalies as outliers, which is brittle when anomalies are primarily mechanism violations rather than extreme values. We propose CausalAno, a causal-aware detector that trains a causal GAN on normal data and leverages its discriminator to learn mechanism-consistent representations. We score test samples by fitting a Gaussian model in this feature space and computing the Mahalanobis distance, measuring deviation from the normal causal manifold. We demonstrate the effectiveness of CausalAno with extensive experiments on 28 tabular datasets (18 continuous-only and 10 mixed-type), comparing against 16 SOTA baselines. Our results show consistent improvements across both mixed-type and numerical-only settings. Our ablation studies further confirm that the gains come from the causal factorization in the causal GAN rather than a generic GAN backbone. Our CausalAno offers a practical and effective solution for real-world applications where anomalies often arise from mechanism-violating behaviors.

Poster
RAST-MoE-RL: A Regime-Aware Spatio-Temporal MoE Framework for Deep Reinforcement Learning in Ride-Hailing

Yuhan Tang, Kangxin Cui, Jung Ho Park, Yibo Zhao, Xuan Jiang, Haoze He, Jiangbo Yu, Haris Koutsopoulos, Jinhua Zhao

Expand abstract

Ride-hailing platforms must balance passenger waiting times with overall system efficiency under highly uncertain supply--demand conditions. Adaptive delayed matching, which controls the holding intervals for batched sets of requests and vehicles, exposes an inherent trade-off between matching and pickup delays. The non-stationary nature of request arrivals and dynamic congestion exposes a key limitation of existing methods, which rely on shallow encoders that cannot capture these spatio-temporal regime shifts. We introduce the Regime-Aware Spatio-Temporal Mixture-of-Experts (RAST-MoE) framework, which formalizes adaptive delayed matching as a regime-aware Markov Decision Process and equips RL agents with a self-attention Mixture-of-Experts encoder, letting different experts specialize automatically across operational regimes while keeping per-sample computation modest. On real-world TNC trajectory data from San Francisco, our 12M-parameter model reduces average matching delay by 10% and pickup delay by 15%, trains stably without reward hacking, generalizes zero-shot to two additional cities, and exhibits clear expert specialization across demand regimes. These results demonstrate the value of MoE-enhanced RL for large-scale decision-making tasks with complex spatio-temporal dynamics.

Poster
Expected Returns and Policy Inconsistency-Aware Offline Federated Deep Reinforcement Learning

Meng XU, Zhongying Chen, Weiwei Fu, Yan Li, Shuguang Wang, Jianping Wang

Expand abstract

Offline Federated Deep Reinforcement Learning (FDRL) methods aggregate multiple client-side offline Deep Reinforcement Learning (DRL) models, each trained locally, to facilitate knowledge sharing while preserving privacy. Existing offline FDRL methods assign client weights during global aggregation using either simple averaging or Q-values, but they neglect the combined consideration of Q-values and policy inconsistency, the latter of which reflects the distributional discrepancy between the learned policy and the policy from offline data. This causes clients with no significant advantages in one aspect but obvious disadvantages in the other to disproportionately affect the global model, thereby degrading its capabilities in that aspect. During local training, clients in existing methods are compelled to fully adopt the global model, which negatively impacts clients when the global model is weak. To this end, we propose a novel Federated Learning (FL) framework that can be seamlessly integrated into current offline FDRL approaches to improve their performance. Our method considers both policy inconsistency and Q-values to determine the weights of client models, with the latter adjusted by a scaling factor to avoid significant numerical discrepancies with the former. The aggregated global model is then distributed to clients to facilitate their learning from the global model. The impact of the global model on the local models is reduced when a client's model performance exceeds that of the global model, thereby mitigating the influence of a weaker global model. Experiments on the Datasets for Deep Data-Driven Reinforcement Learning (D4RL) demonstrate that our method improves seven state-of-the-art (SOTA) offline FDRL methods across several metrics.

Poster
Structure-aware Granular-Ball based Information Bottleneck for Multi-modal Clustering

Zhengzheng Lou, Yuhan Zhan, Mingyang Lv, Yingxuan Li, Yuyang Du, Shizhe Hu

Expand abstract

Multi-modal clustering, which integrates information from diverse sources and feature modalities, has shown great potential in data mining and computer vision. However, existing methods relying on single-granularity relationships often struggle with complex data distributions, leading to limited performance, as fine-grained features are prone to local heterogeneity and redundant perturbations while coarse-grained representations tend to lose local structural information. To address these limitations, we introduce granular-balls (GBs), adaptive multi-granularity hyperspheres that enclose similar samples, and propose the Structure-aware Granular-Ball based Information Bottleneck (SGB-IB) algorithm. This method initializes the dataset as a single GB and recursively splits GBs based on a purity metric, which quantifies the average mutual information between sample features and K-means-derived pseudo-labels across all modalities. It also balances local structure preservation and global redundancy suppression through a structure-aware objective function. Extensive experiments on benchmark datasets demonstrate that our method outperforms state-of-the-art approaches, validating the effectiveness of fusing GB structures with information-theoretic principles.

Poster
SI-IGCL: Subject Invariance-aware Inverse Graph Contrastive Learning for Psychiatric Disorder Identification

Jiayu Lu, Yujin Wang, Xiaofeng Liu, Dandan Li, Bin Wang

Expand abstract

Functional brain network analysis plays an important role in understanding and diagnosing psychiatric disorders. However, current methods struggle with subject variations, impairing the model’s generalization ability to the test set. To address this issue, we propose the Subject Invariance-aware Inverse Graph Contrastive Learning (SI-IGCL) model, which adopts a two-stage paradigm with self-supervised subject-invariant pre-training followed by supervised fine-tuning for identification. During the pre-training phase, we construct an inverse contrastive objective that reshapes the embedding space by repelling intra-subject and attracting inter-subject embeddings to learn subject-invariant representations, with an auxiliary correction term to avoid early optimization plateaus. Meanwhile, we incorporate a structure-preserving reconstruction constraint to preserve discriminative information. Moreover, a Hierarchical Topology Enhanced Transformer (HTET) module is designed to enable multi-level modeling of subject-invariant functional patterns. During the fine-tuning phase, a supervised classifier is integrated to perform psychiatric disorder classification. Extensive experiments demonstrate that our method outperforms all state-of-the-art methods. The code is available at https://anonymous.4open.science/r/SI-IGCL.

Poster
Unsat Core Prediction through Polarity-Aware Representation Learning over Clause-Literal Hypergraphs

Zhenchao Sun, Shuai Ma, Ping Lu, Chongyang Tao

Expand abstract

Graph neural networks have been widely used in Boolean satisfiability (SAT) tasks to learn structural information from SAT formulas. The goal of these studies is to solve SAT instances or to enhance SAT solvers, including tasks such as unsat-core prediction. However, most existing approaches model a SAT formula as a bipartite graph or a directed acyclic graph, which are less direct in capturing clause-level and higher-order interactions among literals and clauses. Moreover, these approaches are limited in modeling intrinsic polarity-related properties of SAT, such as the complementary relationship between the positive and negative literals of a variable. To address these limitations, we propose a polarity-aware representation learning framework over clause-literal hypergraphs. We model SAT formulas as clause-literal hypergraphs augmented with a clause incidence graph to capture higher-order structural interactions. We then introduce a polarity-aware decomposition mechanism that separates variable representations into polarity invariant and equivariant components, explicitly modeling the relationship between positive and negative literals, with the resulting literal representations propagated along the hypergraph structure. We further incorporate a polarity-inversion consistency regularization to reinforce polarity-consistent representations during training. Experimental results on multiple SAT datasets demonstrate the effectiveness of the proposed approach.

Poster
TUR-DPO: Topology- and Uncertainty-Aware Direct Preference Optimization

Abdulhady abas, Fatemeh Daneshfar, Seyedali Mirjalili, Mourad Oussalah

Expand abstract

Aligning large language models (LLMs) with human preferences is commonly done via reinforcement learning from human feedback (RLHF) with Proximal Policy Optimization (PPO) or, more simply, via Direct Preference Optimization (DPO). While DPO is stable and RL-free, it treats preferences as flat winner vs. loser signals and is sensitive to noisy or brittle preferences arising from fragile chains of thought. We propose TUR-DPO, a topology- and uncertainty-aware variant of DPO that rewards how answers are derived, not only what they say, by eliciting lightweight reasoning topologies and combining semantic faithfulness, utility, and topology quality into a calibrated uncertainty signal. A small learnable reward is factorized over these signals and incorporated into an uncertainty-weighted DPO objective that remains RL-free and relies only on a fixed or moving reference policy. Empirically, across open 7–8B models and benchmarks spanning mathematical reasoning, factual question answering, summarization, and helpful/harmless dialogue, TUR-DPO improves judge win-rates, faithfulness, and calibration relative to DPO while preserving training simplicity and avoiding online rollouts. We further observe consistent gains in multimodal and long-context settings, and show that TUR-DPO matches or exceeds PPO on reasoning-centric tasks while maintaining operational simplicity.

Poster
Deep Multi-view Graph Clustering via Attribute-aware Bidirectional Structural Refinement and Pseudo-label Guided Multi-level Fusion

Youqing Wang, Tianxiang Zhao, Mengyuan Xin, Ye Su, Jiapu Wang, Tengfei Liu, Junbin Gao, Jipeng Guo

Expand abstract

Deep multi-view graph clustering (DMGC) typically leverages graph neural networks for representation learning, but most existing methods excessively depend on local and static graph structures and only utilize simplistic cross-view fusion strategies. To this end, this paper proposes **A**ttribute-aware Bidirectional Structural Refinement (ABSR) and **P**seudo-label Guided Multi-level Fusion (PGMF) for DM**GC**, termed **APGC**. Specifically, ABSR selectively strengthens high-quality connections and suppresses semantically conflicting relationships, achieving bidirectional refinement of the graph structure based on attribute similarity. It incorporates global attribute semantics into the graph structure, thereby promoting the homophilic connections for discriminative graph representation learning. Guided by reliable pseudo-labels, PGMF achieves adaptive weighted fusion at both the node-level and the view-level, effectively balancing the differentiated contributions of multi-view information. Experiments on six homophilic and heterophilic datasets demonstrate the superior clustering performance of the proposed APGC method. The code of APGC could be available at https://github.com/TianxiangZhao0474/APGC.git.

Poster
HelioX: A GPU-Native Framework for Simulation and Training of Biophysically Detailed Networks

JUNFENG LU, Zijie Yu, Shaoyang Cui, Gan He, Ruiqin Xiong, Kai Du, Tiejun Huang

Expand abstract

Biophysically detailed neural networks represent a promising frontier for brain-inspired AI, offering intrinsic spatio-temporal dynamics to enhance the expressivity and computational density of deep learning systems. However, general-purpose deep learning frameworks suffer from a fundamental mismatch between their dense parallel optimizations and the irregular, tree-structured complexity of biological mechanisms. In this work, we propose **HelioX**, a \textbf{GPU-native} framework designed to unify high-performance simulation with scalable training. Unlike approaches that adapt biology to existing deep learning tools, **HelioX** adopts a “GPU-to-Biophysics'' paradigm. We tailor the underlying GPU parallelism to biological structures by implementing custom-fused CUDA kernels for both the Dendritic Hierarchical Scheduling (DHS) algorithm and its gradient propagation. This design eliminates the runtime overhead of generic automatic differentiation and enables multi-stream concurrency for spike generation and equation assembly. Experimental results demonstrate that **HelioX** outperforms standard simulators (NEURON) by orders of magnitude and surpasses prior GPU-based solvers in both speed and scalability. We successfully train deep biophysical MLPs and whole-brain-scale biophysical circuits (e.g., the BAAIWorm \textit{C. elegans} model) on a single consumer-grade GPU. **HelioX** establishes a new standard for computational efficiency, enabling the training of biophysically detailed models at scales previously unattainable.

Poster
Topology-Aware Contrastive Learning: Regulating Representation Connectivity via Persistent Homology

Jiaxin Sun, Yuhua Qian, Yang Wang

Expand abstract

Standard contrastive learning minimizes geometric distance between positive pairs, implicitly assuming that strict compactness optimizes discrimination. However, this topology-agnostic confusion neglects intrinsic data structures and topological complexity, leading to class confusion—particularly when aggressive augmentations induce semantic drift. To address this, we propose Topology-Aware Contrastive Learning, a framework that shifts the objective from geometric singularity to topological connectivity. Leveraging Persistent Homology, we explicitly regulate the connectivity of the latent space, ensuring positive pairs maintain an $\alpha\text{--}\beta$ connectivity that balances intra-class cohesion with separability. Theoretically, we formally define the topology-agnostic confusion phenomenon, prove that excessive compactness strictly lower-bounds the probability of confusion and derive a generalization bound demonstrating that richer topological connectivity tightens downstream risk. Furthermore, we establish a measure-theoretic framework to mitigating the sensitivity of our method against varying augmentation strengths. Empirical results on benchmarks confirm that our approach enhances representation quality and reduces reliance on specific augmentation strategies compared to standard baselines.

Poster
GeoEvo: Identity-Aware Potential Game with Geometric Evolution for Personalized Multimodal Federated Learning

Chen Wang, Yongli Hu, Huajie Jiang, Kan Guo, Tengfei Liu, Junbin Gao, Yanfeng Sun, Baocai Yin

Expand abstract

We reconceptualize Personalized Multimodal Federated Learning (PMFL) by treating missing modalities as intrinsic structural identities that constrain each client to a distinct Riemannian submanifold, rather than as deficiencies to be compensated. To reconcile the tension between identity preservation and cross-client collaboration, we cast PMFL as an identity-aware potential game and seek a geometry-consistent equilibrium instead of a monolithic full-modality optimum. We propose GeoEvo, a federated approximate solver that combines curvature-adaptive Fisher descent with manifold-lifted evolutionary search: Natural Evolution Strategies for basin escape, and Particle Swarm updates anchored to a server-broadcast Fr\'echet prototype for cross-client transfer. A monotone acceptance rule drives per-step potential dissipation, yielding an $\mathcal{O}(1/\sqrt{T})$ stationarity rate and convergence toward first-order Nash equilibria in non-convex regimes. Empirically, GeoEvo improves personalization and robustness across diverse modality-missing patterns.

Poster
MetaMoE: Diversity-Aware Proxy Selection for Privacy-Preserving Mixture-of-Experts Unification

Weisen Jiang, Shuhao Chen, Sinno Jialin Pan

Expand abstract

Mixture-of-Experts (MoE) models scale capacity by combining specialized experts, but most existing approaches assume centralized access to training data. In practice, data are distributed across clients and cannot be shared due to privacy constraints, making unified MoE training challenging. We propose **MetaMoE**, a privacy-preserving framework that unifies independently trained, domain-specialized experts into a single MoE using public proxy data as surrogates for inaccessible private data. Central to MetaMoE is diversity-aware proxy selection, which selects client-domain–relevant and diverse samples from public data to effectively approximate private data distributions and supervise router learning. These proxies are further used to align expert training, improving expert coordination at unification time, while a context-aware router enhances expert selection across heterogeneous inputs. Experiments on computer vision and natural language processing benchmarks demonstrate that MetaMoE consistently outperforms recent privacy-preserving MoE unification methods. Code is available at https://github.com/ws-jiang/MetaMoE.

Poster
Trajectory-Aware Heuristic Learning for Combinatorial Search

Mustafa Seddiqi, Marta Kersten-Oertel, Tiberiu Popa

Expand abstract

Learning effective value heuristics for combinatorial search is difficult, as prior methods rely on surrogate supervision or costly downstream search to assess progress. We introduce a trajectory-aware probabilistic framework that models uncertainty in cost-to-go labels instead of treating them as fixed targets. Heuristic learning is cast as inference over state trajectories using an HMM-style model, where estimated depth-change dynamics define transitions and forward-backward inference yields soft supervision. To evaluate heuristic quality without search, we propose a large-scale local ranking metric that measures a model's ability to order neighboring states. On the Rubik's Cube, our approach consistently improves local ranking accuracy and downstream search performance under matched computational budgets.

Poster
ArborKV: Structure-Aware KV Cache Management for Scaling Tree-based LLM Reasoning

Yeqiu Chen, Ziyan Liu, Zhenxin Huang, Runquan Gui, Hong Wang, Lei Liu

Expand abstract

Recent progress in LLM reasoning has increasingly shifted from single-pass generation to explicit search over intermediate reasoning states. Tree-of-Thoughts (ToT) organizes inference to tree-structured search with branching and backtracking, but it substantially amplifies the key--value (KV) cache: retaining KV states for a frontier of partial trajectories quickly becomes a memory bottleneck that limits throughput and constrains search depth and width under fixed hardware budgets. We address this challenge by observing that KV reuse in ToT-style inference is governed by search dynamics: near-term decoding depends primarily on the active branch and its ancestors, whereas inactive subtrees have low short-term reuse probability yet must remain recoverable for backtracking. Motivated by this, we propose **ArborKV**, a structure-aware eviction framework that couples a lightweight value estimator with a tree-aware allocation policy, and performs purely token-extractive eviction with lazy rehydration to support revisits. Experiments on ToT-style reasoning benchmarks show that ArborKV achieves up to $\sim4\times$ peak KV-memory reduction while preserving near-full-retention accuracy, enabling larger search configurations under fixed device budgets that would otherwise run out of memory.

Poster
FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

Felix X.-F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, LINSONG CHU, Davis Wertheimer

Expand abstract

Entropic optimal transport (EOT) via Sinkhorn iterations is widely used in modern machine learning, yet GPU solvers remain inefficient at scale. Tensorized implementations suffer quadratic HBM traffic from dense $n\times m$ interactions, while existing online backends avoid storing dense matrices but still rely on generic tiled map-reduce reduction kernels with limited fusion. We present **FlashSinkhorn**, an IO-aware EOT solver for squared Euclidean cost that rewrites stabilized log-domain Sinkhorn updates as row-wise LogSumExp reductions of biased dot-product scores, the same normalization as transformer attention. This enables FlashAttention-style fusion and tiling: fused Triton kernels stream tiles through on-chip SRAM and update dual potentials in a single pass, substantially reducing HBM IO per iteration while retaining linear-memory operations. We further provide streaming kernels for transport application, enabling scalable first- and second-order optimization. On A100 GPUs, FlashSinkhorn achieves up to $32\times$ forward-pass and $161\times$ end-to-end speedups over state-of-the-art online baselines on point-cloud OT, improves scalability on OT-based downstream tasks. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn.

Spotlight