NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

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

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

MotionGrounder: Grounded Multi-Object Motion Transfer via Diffusion Transformer

Samuel Teodoro, Yun Chen, Agus Gunawan, Soo Ye Kim, Jihyong Oh, Munchurl Kim

Expand abstract

Motion transfer enables controllable video generation by transferring temporal dynamics from a reference video to synthesize a new video conditioned on a target caption. However, existing Diffusion Transformer (DiT)–based methods are limited to single-object videos, restricting fine-grained control in real-world scenes with multiple objects. In this work, we introduce MotionGrounder, a DiT-based framework that $\textit{firstly}$ handles motion transfer with $\textit{multi-object controllability}$. Our Flow-based Motion Signal (FMS) in MotionGrounder provides a stable motion prior for target video generation, while our Object-Caption Alignment Loss (OCAL) grounds object captions to their corresponding spatial regions. We further propose a $\textit{new}$ Object Grounding Score (OGS), which jointly evaluates (i) spatial alignment between source video objects and their generated counterparts and (ii) semantic consistency between each generated object and its target caption. Our experiments show that MotionGrounder consistently outperforms recent baselines across quantitative, qualitative, and human evaluations.

Poster
MotionCFG: Boosting Motion Dynamics via Semantic Motion Sharpening

Byungjun Kim, Soobin Um, Jong Chul Ye

Expand abstract

Despite recent advancements in Text-to-Video (T2V) synthesis, generating high-fidelity and dynamic motion remains a significant challenge. Existing methods primarily rely on Classifier-Free Guidance (CFG), but standard CFG treats all semantic dimensions uniformly and provides no mechanism to selectively resolve ambiguity along the motion axis. To address this, we propose MotionCFG, a training-free approach that introduces selective guidance along the motion subspace of the condition embedding. Specifically, we construct a motion-perturbed negative condition by injecting Gaussian noise into motion-related embeddings, and steer generation away from it to selectively amplify motion intent. We show theoretically that this procedure implicitly approximates *Semantic Laplacian Sharpening*, a second-order correction that suppresses motion-ambiguous regions of the score landscape and amplifies well-resolved dynamic peaks. Combined with a piecewise guidance schedule that confines intervention to the early denoising steps, MotionCFG consistently improves motion dynamics across state-of-the-art T2V frameworks with negligible overhead. We further demonstrate that this Laplacian sharpening principle generalizes beyond motion, effectively steering complex, non-linear concepts such as precise object numerosity that are typically difficult to modulate via standard text-based guidance.

Poster
Motion Forcing: Decoupling Ego and Object Motion via Sparse Inputs for Structured Video Generation

Tianshuo Xu, ZhiFei Chen, Leyi Wu, Hao Lu, Yingcong Chen

Expand abstract

Structured video generation (e.g., autonomous driving) faces a fundamental domain gap between simple human control and dense physical reality. Text-to-Video models accept accessible inputs but struggle with precise spatial guidance, while dedicated autonomous driving world models require expensive, dense 3D annotations that are difficult for users to provide. To bridge this gap, we introduce \textbf{Motion Forcing}, a structured video generation framework that achieves precise multi-agent control using extremely sparse 2D inputs. Our key insight is to explicitly bridge human intent and visual synthesis via a hierarchical \textbf{``Point-Shape-Appearance''} paradigm. This approach decomposes generation into verifiable stages: modeling user intent as sparse geometric points, expanding them into dynamic depth maps to explicitly resolve 3D geometry and decouple ego-motion from object dynamics, and finally rendering high-fidelity textures. Furthermore, to elevate the model from passive instruction-following to active physical reasoning, we employ a \textbf{Masked Point Recovery} strategy. By forcing the reconstruction of dynamic depth from occluded trajectories, the model internalizes latent physical laws, enabling causal inference. Extensive experiments demonstrate that Motion Forcing significantly outperforms state-of-the-art baselines on large-scale autonomous driving benchmarks. Additional validation in rigid-body physics and robotic manipulation confirms the structural integrity and generalizability of our framework. \textbf{To facilitate future research, we have released our model weights and code.}

Poster
Motion Cues from Image-based Point Tracking for LiDAR Scene Flow Estimation

Youngdong Jang, Gyeongrok Oh, Jong Wook Kim, Hyunju Ryu, Hyung-gun Chi, Seunghyeon Kim, Seungryong Kim, Jonghyun Choi, Sangpil Kim

Expand abstract

LiDAR scene flow estimation is essential for autonomous driving, as it provides 3D motion for each point. Self-supervised approaches use static-dynamic classification to mitigate the imbalance between static and dynamic points, deriving targeted supervision. However, existing methods rely on sparse geometric observations for this classification, making them vulnerable to data sparsity and occlusions. The resulting noisy labels provide incorrect motion guidance and degrade scene flow learning. To address this, we introduce TrackCue, a tracking-guided framework for improving dynamic object representation in LiDAR scene flow estimation. In particular, TrackCue repurposes point tracking to obtain dense image-space trajectories anchored to LiDAR points, providing motion cues beyond sparse geometric observations. Furthermore, we present a visually consistent motion compensation strategy that compares the tracked trajectories with ego-induced rigid trajectories in the image plane, effectively isolating true object motion from ego-induced apparent motion. To transfer these isolated motion cues back to the LiDAR domain, we perform visual motion cue lifting, which associates ego-compensated image trajectories with LiDAR points for static-dynamic label refinement. As a result, TrackCue produces more accurate static-dynamic classification and provides more reliable supervision for scene flow learning. Experimental results show that TrackCue significantly improves the precision and F1 score of dynamic labels, leading to performance gains in self-supervised scene flow estimation.

Poster
Motif-Mamba: network motif improved mamba for long-range sequence modeling

Chonghe Hao, Yue Sun, Jian Zhang, Yansong Wang, Wangzi Yao, Yao Yunjie, Tielin Zhang

Expand abstract

Efficient long-sequence modeling remains a central challenge for large language models, as self-attention scales quadratically with sequence length. Mamba offers a linear-time alternative through selective state space recurrence, but its predominantly diagonal state transitions restrict explicit interactions among state dimensions. We propose Motif-Mamba, a structured state space model that augments Mamba with a motif-constrained low-rank recurrent pathway. Inspired by the dynamics of three-node network motifs, the proposed pathway projects hidden states into a compact dynamical subspace, imposes motif-guided interactions, and maps the resulting dynamics back to the original state space. This design enhances cross-dimensional communication while preserving the linear-time recurrent structure of Mamba. Experiments on long-sequence extrapolation, language modeling benchmarks, and brain--computer interface decoding show consistent improvements over Mamba backbones, suggesting that motif-guided low-rank dynamics provide an effective structural prior for long-range sequence modeling.

Poster
Mosaic: A Benchmark Suite for Differentiable Physics Solvers

Andrin Rehmann, Heiko Zimmermann, Dion Häfner

Expand abstract

Differentiable physics solvers underpin solver-in-the-loop ML training, gradient-based optimal control, and inverse problems, yet the practical cost of obtaining correct, usable gradients from a given solver on a given problem is largely undocumented. Integration effort, computational cost, gradient accuracy, and numerical conditioning vary widely across solvers and are discoverable only by trial and error. We introduce Mosaic, an extensible benchmarking framework for differentiable PDE solvers that standardizes access to solver gradients. Each solver is packaged as a containerized component (Tesseract) exposing a uniform gradient API regardless of language or automatic differentiation (AD) strategy, enabling researchers to evaluate, compare, and build on non-trivial physical solvers. Our evaluation of 14 solvers across fluid dynamics, structural mechanics, and heat transfer demonstrates that the benchmark surfaces practically relevant differences: order-of-magnitude variation in computational cost and Jacobian conditioning, alongside structural incompatibilities that eliminate solvers from realistic tasks entirely. Despite this variation, all solvers that produce gradients converge to similar optima, indicating that the practical barriers are memory limits, numerical stability, and setup compatibility rather than gradient accuracy alone. Mosaic is open-source and available at https://anonymous.4open.science/r/mosaic-bench.

Poster
MorphoHELM: A Comprehensive Benchmark for Evaluating Representations for Microscopy-Based Morphology Assays

Emre Hayir, Lorin Crawford, Alex X Lu

Expand abstract

Microscopy images contain rich information about how cells respond to perturbations, making them essential to applications like drug screening. To quantify images, researchers often use representation extraction methods, and recent years have seen a proliferation of deep learning methods. While measuring the quality of these representations is essential, evaluation remains fragmented, with each proposed model evaluated on different tasks and datasets, using custom pipelines and metrics, making it difficult to fairly compare models. Here, we introduce MorphoHELM, a comprehensive open benchmark for evaluating feature extraction methods for cell painting, the most widely-used morphological profiling assay. MorphoHELM consolidates evaluation standards in the field, extends and corrects them to be more robust, and evaluates on the widest range of methods to date. A defining feature of the benchmark is that each task is evaluated at different degrees of batch effects (or technical noise), directly quantifying how the ability of methods to detect biological signal degrades as noise increases. Together, these properties enable MorphoHELM to detect trade-offs between methods, and we demonstrate that models that excel at certain kinds of biological signal are weaker at others. We show that no existing model outperforms classic computer vision analytic strategies across all settings, which remain the strongest general use-case representations. All datasets, code, and evaluation tools are publicly available.

Poster
Morpho-Temporal Decoupling: How Primate Neurons Expand Dendrites Without Losing Speed

Jiawei Zhang, Haoyu Wang, Jialun Ma, Wei Dai, yuguo yu

Expand abstract

During evolution, biological neurons scale computation by expanding dendritic branches. However, this expansion creates a biophysical dilemma: increased membrane area typically imposes a capacitive load that slows somatic dynamics and narrows temporal bandwidth. Here, we investigate how primate cortical neurons resolve this tradeoff. While primary basal dendrite number is relatively conserved across mouse cortical areas, it increases systematically along the primate cortical hierarchy. Using biophysically constrained multicompartment modeling, we identify a key dimensionless control parameter---the dendritic-to-somato-apical specific membrane resistance ratio, $\rho \equiv R_{m,\mathrm{dend}}/R_{m,\mathrm{sa}}$, that governs the coupling between dendritic morphology and the effective somatic time constant $\tau$. In mouse-like neurons ($\rho>1$), adding basal dendrites progressively increases $\tau$ and promotes reliable low-pass integration. Conversely, human-like neurons operate near a balanced-resistance regime ($\rho \approx 1$), where resistive reweighting counterbalances morphology-driven capacitive loading. This allows $\tau$ to remain nearly invariant despite expanded dendritic topology, shifting the structure-function tradeoff toward faster, more flexible coding with improved high-frequency tracking. Our results reveal a biophysical scaling principle linking evolutionary dendritic architecture to hierarchy-dependent coding, offering potential design rules for preserving temporal bandwidth in dendrite-inspired neuromorphic architectures.

Poster
MorphGen: Controllable Cell-Image Generation with Biological Representation Alignment

Berker Demirel, Marco Fumero, Theofanis Karaletsos, Francesco Locatello

Expand abstract

Simulating in silico cellular responses to interventions is a promising direction to accelerate high-content image-based assays, critical for advancing drug discovery and gene editing. To support this, we introduce MorphGen, a state-of-the-art diffusion-based generative model for fluorescent microscopy that enables controllable generation across multiple cell types and perturbations. MorphGen is trained with an alignment loss that matches its representations to phenotypic embeddings from a biology foundation model, encouraging meaningful morphological patterns consistent with real cell images. Unlike prior approaches that compress multichannel stains into RGB images, sacrificing organelle-specific detail and focusing on a single cell type, MorphGen generates the complete set of fluorescent channels jointly, preserving per-organelle structure and enabling post-generation interpretation. We demonstrate biological consistency with real images via CellProfiler features, and MorphGen attains an FID score over 35% lower than the prior state-of-the-art MorphoDiff. Finally, in a compositional generalization test that holds out cell type--perturbation combinations during training, MorphGen achieves in-distribution quality on 55% of unseen pairs under a seen-calibrated criterion.

Poster
More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding

Noam Elata, Itay Lamprecht, Mikey Shechter, Daniel Ohayon, Itay Hubara, Daniel Soudry

Expand abstract

Autoregressive generation in Large Language Models (LLMs) is constrained by the memory and computational demands of attention mechanisms. Sparse attention methods mitigate this cost by selecting only high-probability entries of the attention matrix. We observe that in many such methods, this renders the probability-value multiplication negligible, shifting the bottleneck to the query-key step. Key heads can therefore be reduced to accelerate inference, while value heads preserve capacity at no additional cost. We introduce Sparse Asymmetric Group-Query Attention (SAGA), which decouples key and value head counts to exploit this principle, and pair it with approximate top-N (Atop-N) attention, a simple GPU-agnostic sparse attention method designed to isolate sparsity's effect on decoding-optimized architectures. We formalize the benefits of this asymmetry theoretically and validate it empirically, achieving end-to-end decoding speedups exceeding $2\times$ at long contexts across multiple model scales and benchmarks. Models trained from scratch with SAGA nearly match the quality of comparable GQA variants while delivering substantial efficiency gains. To facilitate adoption, we introduce an efficient fine-tuning method that converts pretrained models to the SAGA architecture at a small to negligible cost to quality, enabling practitioners to benefit from our approach without costly retraining.

Poster
More Than Meets the Eye? Uncovering the Reasoning-Planning Disconnect in Training Vision-Language Driving Models

Xurui Song, Shuo Huai, Jingjing Jiang, Jiayi Kong, Jun Luo

Expand abstract

Vision-Language Model (VLM) driving agents promise explainable end-to-end autonomy by first producing natural-language reasoning and then predicting trajectory planning. However, whether planning is **causally** driven by this reasoning remains a critical but unverified assumption. To investigate this, we build DriveMind, a large-scale driving Visual Question Answering corpus with plan-aligned Chain-of-Thought (CoT), automatically generated from nuPlan. Our data generation process converts sensors and annotations into structured inputs and, crucially, separates priors from to-be-reasoned signals, enabling clean information ablations. Using DriveMind, we train representative VLM agents with Supervised Fine-Tuning and Group Relative Policy Optimization and evaluate them with nuPlan’s metrics. Our results, unfortunately, indicate a consistent **causal disconnect** in reasoning-planning: removing ego/navigation priors causes large drops in planning scores, whereas removing CoT produces only minor changes. Attention analysis further shows that planning primarily focuses on priors rather than the CoT. Based on this evidence, we propose the Reasoning-Planning Decoupling Hypothesis, positing that the training-yielded reasoning is an ancillary byproduct rather than a causal mediator. To enable efficient diagnosis, we introduce a novel, training-free probe that measures an agent's reliance on priors by evaluating its planning robustness against minor input perturbations. In summary, we provide the community with a new dataset and a diagnostic tool to evaluate the **causal fidelity** of future models.

Poster
Monroe: A Molecular Foundation Model for In-context Probabilistic Inference

Błażej Banaszewski, Andrew Fitzgibbon

Expand abstract

Bioassay activity prediction is often data-limited because drug-discovery datasets rely on time-consuming and expensive wet-lab experiments for data generation and evaluation. This challenge has inspired recent research into molecular foundation models (MFMs), which aim to encode general-purpose chemical knowledge into molecular representations that generalize well in data-constrained scenarios. This paper presents Monroe, a new MFM with several innovations over the existing state of the art: increased scale allowing pre-training on over 80 million molecules from the PM6 quantum chemistry dataset; improved graph representation of stereochemistry; improved training losses including conformer denoising and embedding decorrelation; improved multi-task learning; and the use of a prior-data-fitted model (TabPFN) for downstream in-context prediction. Our evaluations use a principled pairwise comparison framework that measures statistically significant performance differences. Across established Polaris benchmarks, Monroe matches or exceeds existing MFMs, while on activity cliff benchmarks, designed to assess utility for molecular discovery, it achieves significant improvements over prior methods. Finally, ablation and transfer experiments show that PFN-based downstream predictors also substantially improve two leading existing models, MiniMol and CheMeleon, yielding new state-of-the-art variants we call MiniMol_PFN and Chemeleon_PFN, suggesting that our downstream adaptation strategy generalizes beyond Monroe. Full source code is at https://anonymous.4open.science/r/monroe-0E0D, and will be openly released.

Poster
Monotone Inclusion Approach to Weakly Monotone Discrete-Time Finite-Horizon Mean-Field Games

Ugur Aydin, Tamer Basar, Naci Saldi

Expand abstract

We revisit the problem of computing mean-field equilibria (MFEs) in discrete-time, monotone, finite-horizon mean-field games (MFGs). We show that, when the transition kernel is independent of the state-measure term and the reward function satisfies the usual weak monotonicity condition and is Lipschitz continuous, anchored proximal gradient descent methods can be used to compute a monotone MFE. We also establish last-iterate convergence results for these methods. Our approach relies on formulating the computation problem as an optimization problem over the space of occupation measures. Using this formulation, we show that the problem is equivalent to a class of constrained Lipschitz monotone inclusion problems. We then apply iterative methods for this monotone inclusion formulation to derive a tractable algorithm. The resulting algorithm achieves a convergence rate of $O(1/\sqrt{T})$ after $T$ iterations, without requiring any regularization. This rate holds even in the absence of a uniqueness assumption for the corresponding MFE.

Poster
Monoculture Robust Learning

Lee Cohen, Jon Kleinberg, Omer Reingold

Expand abstract

We introduce a notion of monoculture robustness in Probably Approximately Correct (PAC) learning, as a way to mitigate the problem of *algorithmic monoculture* (Kleinberg & Raghavan, 2021): the tendency of models sharing the same components (e.g., training data) to fail on the same inputs. Specifically, we ask whether a learner with access to a single dataset can output $k$ individually accurate hypotheses whose joint failure probability is comparable to training on $k$ independent datasets. We say that such a learning procedure achieves *monoculture robustness*, and we quantify this by bounding the *conjunctive error* that all $k$ hypotheses misclassify points in an individual sense. We design methods for achieving this type of guarantee by drawing a surprising connection between replicability (Impagliazzo et al., 2022) and monoculture robustness. Specifically, we show that if the base learner is $\rho$-replicable, then running it independently $k$ times on the same sample yields individual monoculture error that shrinks exponentially with $k$, controlled by the single-run error of the base learner at the point and $\rho$. We complement these positive results with a lower bound: for individual monoculture robustness, we prove a worst-case lower bound showing that, for learners, driving the individual conjunctive error down to $\gamma$ can require a sample-size multiplier of $t = \Omega(\log(1/\gamma))$.

Poster
MonoPhysics: Estimating Geometry, Appearance, and Physical Parameters from Monocular Videos

Daniel Rho, Jun M Choi, Matthew Thornton, Biswadip Dey, Roni Sengupta

Expand abstract

Existing inverse physics methods recover physical parameters from multi-view videos, where geometric constraints across views resolve scale and 3D structure. In monocular settings, however, such constraints are absent, leading to severe scale ambiguity, inaccurate geometry, and weak coupling between appearance optimization and physical simulation. We propose MonoPhysics, a framework for monocular inverse physics estimation of deformable objects using differentiable MPM simulation and 3D Gaussian Splatting, which jointly optimizes geometry, appearance, and physical parameters from a single camera view. We address these challenges through three visual-physical bridges: global scale alignment, physics-aware geometry refinement, and a differentiable position map, which together enable accurate optimization from image-space losses alone. We evaluate on Vid2Sim and our new dataset of elastic and plastic objects, showing that MonoPhysics outperforms existing baselines in monocular settings and achieves performance comparable to multi-view baselines using only a single camera.

Poster
Monitoring the Internal Monologue: Probe Trajectories Reveal Reasoning Dynamics

Maciej Chrabaszcz, Aleksander Szymczyk, Marcin Sendera, Tomasz Trzcinski, Sebastian Cygert

Expand abstract

Large Reasoning Models (LRMs) introduce new opportunities for safety monitoring through their Chain of Thought (CoT) reasoning. However, CoT is not always faithful to the model's final output, undermining its reliability as a monitoring tool. To address this, we investigate the hidden representations of LRMs to determine whether future behavior can be predicted from prompt and CoT representations. By evaluating a probe at each generated token, we construct a probe trajectory, the continuous evolution of a concept's probability across the reasoning process. We find that future model behavior is more distinguishable when examined over the full trajectory than from a single static prediction. To characterize these temporal dynamics, we extract signal-processing features that capture volatility, trend, and steady-state behavior, significantly improving the separation of future model states. We also present two methodological insights. First, template-based training data achieves near-parity with dynamically generated model responses, eliminating the need for a costly initial inference and labeling. Second, the choice of pooling operation is critical: average-pooling and last-token methods collapse to near-random performance, while max-pooling achieves up to $95\%$ AUROC and yields stable probe trajectories. Using four datasets and four reasoning models across the domains of safety and mathematics, we demonstrate that trajectory features encode task-specific dynamics that improve outcome separability. These findings establish probe trajectories as a complementary framework for monitoring LRM behavior. Warning: This article contains potentially harmful content.

Poster
MonarchRT: Efficient Attention for Real-Time Video Generation

Krish Agarwal, Zhuoming Chen, Cheng Luo, Yongqi Chen, Haizhong Zheng, Xun Huang, Atri Rudra, Beidi Chen

Expand abstract

Real-time video generation with Diffusion Transformers is bottlenecked by the quadratic cost of 3D self-attention, especially in real-time regimes that are both few-step and autoregressive, where errors compound across time and each denoising step must carry substantially more information. In this setting, we find that prior sparse-attention approximations break down, despite showing strong results for bidirectional, many-step diffusion. Specifically, we observe that video attention is not reliably sparse, but instead combines pronounced periodic structure driven by spatiotemporal position with dynamic, sparse semantic correspondences and dense mixing, exceeding the representational capacity of even oracle top-k attention. Building on this insight, we propose MonarchRT, a structured attention parameterization for video diffusion models that factorizes attention using Monarch matrices. Through appropriately aligned block structure and our extended tiled Monarch parameterization, we achieve high expressivity while preserving computational efficiency. We further overcome the overhead of parameterization through finetuning, with custom Triton kernels. We first validate the high efficacy of MonarchRT over existing sparse baselines designed only for bidirectional models. We further observe that MonarchRT attains up to 95% attention sparsity with no loss in quality when applied to the state-of-the-art model Self-Forcing, making MonarchRT a pioneering work on highly-capable sparse attention parameterization for real-time video generation. Our optimized implementation outperforms FlashAttention-2, FlashAttention-3, and FlashAttention-4 kernels on Nvidia RTX 5090, H100, and B200 GPUs respectively, providing kernel speedups in the range of 1.4-11.8X. This enables us, for the first time, to achieve true real-time video generation with Self-Forcing at 16 FPS on a single RTX 5090.

Poster
MolSpecFlow: Modality-Incomplete Molecular--Spectral Learning for MS/MS

Yu Wang, Fan Yang, Kaikun Xu, Li Hao, Li Yuan, Jun Zhu, Jingjie Zhang, Zhenchao Tang, Yatao Bian, Cheng Chang, Jianhua Yao

Expand abstract

Untargeted metabolomics produces massive numbers of tandem mass spectra, yet most spectra remain difficult to assign to molecular structures because each spectrum is only a partial, instrument-dependent observation of the underlying molecule. We study MS/MS annotation as modality-incomplete molecular--spectral inference: molecules, spectra, and fingerprints provide complementary views of the same chemical entity, but fully paired molecule--spectrum data are limited and many training examples contain only molecular or only spectral information. We introduce MolSpecFlow, a mask-conditioned heterogeneous flow framework that represents these modalities as coupled components of a shared molecular--spectral state. Availability masks encode which components exist in each example, and observation masks specify which components are conditioned on or generated, instantiating de novo annotation, molecular retrieval, and spectrum simulation within one backbone. MolSpecFlow uses a shared Mol-Spec Transformer, hybrid discrete--continuous flow dynamics, and expected mass/formula regularization over soft molecular token distributions. Experiments show consistent gains from modality-incomplete pretraining, multi-interface reuse, and external de novo transfer.

Poster
MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models

Hojung Jung, Rodrigo Hormazabal, Jaehyeong Jo, Youngrok Park, Kyunggeun Roh, Se-Young Yun, Sehui Han, Dae-Woong Jeong

Expand abstract

Molecular generation with diffusion models has emerged as a promising direction for AI-driven drug discovery and materials science. While graph diffusion models have been widely adopted due to the discrete nature of 2D molecular graphs, existing models suffer from low chemical validity and struggle to meet the desired properties compared to 1D modeling. In this work, we introduce MolHIT, a powerful molecular graph generation framework that overcomes long-standing performance limitations in existing methods. MolHIT is based on the Hierarchical Discrete Diffusion Model, which generalizes discrete diffusion to additional categories that encode chemical priors, and decoupled atom encoding that splits the atom types according to their chemical roles. Overall, MolHIT achieves new state-of-the-art performance on the MOSES dataset with near-perfect validity for the first time in graph diffusion, surpassing strong 1D baselines across multiple metrics. We further demonstrate strong performance in downstream tasks, including multi-property guided generation and scaffold extension.

Poster
Module-Aware Optimization for Graph Neural Networks

Guy Hadad, Haggai Roitman, Moshe Eliasof, Bracha Shapira

Expand abstract

Different optimizers impose different geometries on parameter updates. Adaptive methods such as Adam and AdamW use coordinate-wise updates, while matrix-aware optimizers such as Muon utilize weight-matrix structure through orthogonalized updates. For Graph Neural Networks (GNNs), however, the effective geometry of its weight matrices is shaped not only by its learned channel-mixing weights, but also by the graph-propagated representations it acts on. Thus, even if an update is well normalized as a matrix, its effect on the layer output can be amplified when applied to graph-propagated features. To address this gap, we introduce GraphMuon, a graph-aware optimizer that adapts Muon's matrix updates to the covariance structure of graph-propagated layer inputs. GraphMuon performs orthogonalization in a metric defined by these propagated representations and normalizes step sizes according to the update's induced effect on the layer representation. We show that theoretically, vanilla Muon can suffer from graph-induced update amplification, while our GraphMuon can mitigate this effect. As such, our GraphMuon yields a graph-metric analogue of Muon's polar update step. Empirically, we show that GraphMuon achieves $2.13\times$-$7.38\times$ speedups in time-to-target wall-clock runtime over AdamW across diverse graph-learning regimes, while maintaining competitive or superior predictive performance.

Poster