NeurIPS 2026 论文列表

Conference on Neural Information Processing Systems

录用类型

6,437篇论文

按官方录用论文列表统计,可能与会议公布的数字略有出入。

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

浏览全部论文

Anatomy-Activated Mixture-of-Experts for 3D Medical Vision-Language Pre-training

Szymon Płotka, Gizem Mert, Pedro R.A.S. Bassi, Wenxuan Li, Zongwei Zhou, Anna M Marcinkiewicz, Wiktoria Romańczyk, Jarosław B Ćwikła, Łukasz Struski, Ewa Szczurek, Jacek Tabor, Arkadiusz Sitek

展开摘要

Medical Vision-Language Pre-training (VLP) leverages the semantic richness of radiology reports to learn 3D representations, yet current architectures remain fundamentally anatomy-agnostic. While clinical diagnostics rely on organ-specific context, standard encoders apply a uniform, dense parameterization to all volumetric patches, and existing Mixture-of-Experts (MoE) routers gate on isolated local features without explicit access to the organ composition of the scan or crop. We propose Anatomy-Activated Mixture-of-Experts (A$^{2}$MoE), a framework that incorporates anatomical inductive biases within the model's computational path. A$^{2}$MoE introduces two key components: (1) a Histogram Representation Router (HRR) that conditions expert selection on global visual context and a predicted anatomy-presence histogram, ensuring that parameters specialize in specific anatomical domains; and (2) a mask-free Organ-Query Multi-Scale Projector (OQMSP) that utilizes learnable queries to extract organ-level embeddings for fine-grained alignment without requiring manual annotations at inference. By routing at the patch level rather than the token level, A$^{2}$MoE reduces routing complexity by orders of magnitude while satisfying load balancing by construction. Pre-trained on over 60,000 multi-phase CT scans, A$^{2}$MoE achieves state-of-the-art performance across segmentation, classification, and detection benchmarks. Furthermore, we demonstrate that our anatomically-grounded inductive biases support robust generalization to MRI and ultrasound, establishing a scalable foundation for medical representation learning.

Poster
Analyzing Learning-Dynamics Across Loss and Prediction Choices for Diffusion Models

Dongwoo Kim

展开摘要

Denoising generative models predict clean data ($x$), noise ($\epsilon$), or velocity ($v$), and are trained with a squared loss that can be computed in any of these three spaces. Prior work treats these choices as interchangeable up to loss reweighting, but practitioners observe sharp qualitative differences that remain unexplained. We systematically decouple the \emph{prediction space}, what the network outputs, from the \emph{loss space}, where supervision is applied, creating a $3{\times}3$ design matrix. We show that loss-space changes only reweight how much each noise level contributes to training, while prediction-space changes alter what the network is asked to learn: two fundamentally different mechanisms. This decoupling exposes two regimes. Without a representational bottleneck, the loss choice primarily controls denoising quality (MSE) while the prediction choice primarily controls sample quality (FID). The reason is that ODE sampling must convert the network's output back into a clean image, and this conversion amplifies errors differently across prediction types, turning sub-one-percent gaps in training error into 22-fold gaps in FID that are invisible to standard training diagnostics. Under a representational bottleneck, the prediction choice instead dominates both metrics: recovering noise or velocity requires reconstructing high-dimensional information that the bottleneck has destroyed, while recovering the clean image remains feasible because real images lie on a low-dimensional manifold. We validate these findings on a U-Net and the full 130M-parameter JiT architecture across CIFAR-10 and Imagenette, and give architecture-dependent recommendations for both regimes.

Poster
AnalogToBi: Device-Level Analog Circuit Topology Generation via Bipartite Graph and Grammar Guided Decoding

Seungmin Kim, Mingun Kim, Yuna Lee, Yulhwa Kim

展开摘要

Analog circuit design remains highly dependent on expert knowledge due to the complexity of device-level interactions and topology design. Recent transformer-based approaches for device-level topology generation have shown promise, yet they suffer from low electrical validity without human-in-the-loop (HITL) training and severe memorization caused by sequence-based circuit representations. In this work, we propose **AnalogToBi**, a framework for device-level analog circuit topology generation. AnalogToBi introduces *circuit-type conditioning* for categorizing heterogeneous multi-type topology datasets, *device renaming augmentation* to mitigate memorization, *a bipartite graph representation* for improved structural generalization, and *grammar-guided decoding* to enforce structural validity during bipartite graph generation. Experimental results demonstrate that AnalogToBi achieves high validity and novelty without HITL training while effectively avoiding memorization of training topologies.

Poster
AnaDiffusion: Anatomically Compositional Latent Diffusion for Controllable 3D Brain MRI Generation

Tracy Han, Lulin Liu, Bangya Liu, Yuanhao Cai, Nuo Chen, Jade Wang, Ziqian Xie, Chenyu You, Shuiwang Ji, Degui Zhi, Zhiwen Fan

展开摘要

3D brain MRI generation has made significant advances for medical imaging, simulation, and controllable anatomical analysis. However, existing generative models typically synthesize 3D volumes monolithically, often overlooking regional anatomical structures and limiting local controllability. To address these limitations, we introduce AnaDiffusion, an anatomically compositional latent diffusion framework that factorizes the generation process into distinct, anatomically meaningful regions followed by part-to-whole assembly and global refinement. Our approach trains part diffusion models first to capture local structural priors. Then, we inject assembled anatomical composite of parts into whole-brain latent and continue denoising. This mechanism enables the model to resolve global context while preserving the injected anatomy. As a result, AnaDiffusion outputs both explicit part assets and a globally coherent volume, enabling controllable part editing without requiring additional dense segmentation masks at inference time while maintaining coherent part-to-whole brain structure. On ADNI, AnaDiffusion achieves the lowest FID across the whole brain, left/right hemispheres, cerebellar-brainstem complex, and seam regions, while also reducing Cohen's d for ventricles, cerebellum, and brainstem. In localized editing experiments, paired MS-SSIM shows high target transfer and off-target preservation, supporting controllable part replacement with limited non-target anatomical drift.

Poster
An exact information theory of generalization phase transitions in Bayesian diffusion models

Henry Hunt, Mason Kamb, Surya Ganguli

展开摘要

How diffusion models circumvent the curse of dimensionality to learn complex distributions over high dimensional spaces from a finite training set, instead of memorizing it, remains a fundamental mystery. To address this, we introduce analytically tractable Bayesian information restricted diffusion (BIRD) models, in which each pixel observes restricted information about noisy data. A BIRD model time-reverses diffusion by inferring which past training sample produced its current restricted observation using the Bayesian posterior. This model class generalizes existing analytical diffusion models that use spatially local information restriction. We show that spatially local BIRD models closely approximate trained diffusion models early in training, across different architectures such as UNets and DiTs. Under minimal assumptions on the data distribution, we identify an information-theoretic phase boundary between memorization and generalization in the joint space of amount of training data, time in the reverse generative process, and amount of information restriction: a BIRD model memorizes when the mutual information between its restricted noisy observations and the training data exceeds the log number of training points, and it generalizes otherwise. Experiments across a range of datasets confirm our theoretically predicted location for the transition. We find that generation proceeds near the edge of memorization: both spatially local BIRD models and early-training diffusion models track the memorization-generalization phase boundary by increasingly restricting information over time. Overall, our results reveal a fundamental role for information restriction in generative AI to circumvent the curse of dimensionality.

Poster
An Open-Source Training Dataset for Foundation Models for Black-box Optimization

Aaron Klein, Herilalaina Rakotoarison, Luca Thale-Bombien, David Salinas

展开摘要

Most black-box optimization methods require extensive hyperparameter tuning, often limiting their ability to generalize across different optimization domains. Foundation models for black-box optimization that learn optimization principles from a large collection of optimization trajectories offer a promising alternative, with the potential to outperform manually designed methods across diverse problem classes. However, prior work has either relied on non-public datasets or on purely synthetic data, limiting reproducibility and generalization to real-world problems. As a result, progress in this area has been constrained by the lack of large-scale, real-world, publicly available pre-training data. We introduce BBO-Pile, the first open-source dataset comprising over 500K optimization trajectories evaluated across 3095 different black-boxes for different optimizers which represents by far the largest public dataset for this task. Using this dataset, we train a family of foundation models at multiple scales, ranging from 2M to 80M parameters and from 200M to 2B training tokens, and study their scaling behavior with respect to compute. Our results demonstrate that large-scale pre-training is a viable and effective approach to imitate black-box optimization methods, paving the way for future research in this direction.

Poster
An Information-theoretic Framework for Auditing Unfairness in Training Data

Mohamed Nafea, Sokrat Aldarmini

展开摘要

Quantifying bias in training data---independently of any downstream model---is critical for fair machine learning, especially when data collection and model design occur in silos. Yet most fairness research focuses on model-level mitigation, leaving the data side underexplored. We address this gap by deriving multiple information-theoretic measures for auditing unfairness bias in data without requiring access to model predictions. We define our measures over feature sets and use Shapley values to deduce feature-level contributions to global unfairness bias. For some of these measures, we prove that Shapley values moderate redundancy across feature coalitions, matching empirical evidence in prior work. We develop our measures within an axiomatic framework that embeds a fairness notion, such as statistical parity (SP) or equalized odds (EO), into desired and undesired independence properties. We characterize baseline measures that guide those we develop, for example, by showing that a baseline upper bounds the bias of any downstream restricted predictor. Finally, we show that SP- and EO-aligned independence properties can conflict, paralleling Kleinberg et al.'s result on incompatible fairness notions. We validate our measures through an empirical study on real and synthetic datasets, assessing how well they predict the true bias of an extensively tuned neural network trained on the same data, and via a simple feature selection experiment. For synthetic data, we propose a parametric linear structural causal model that enables controlled generation of diverse correlation structures and bias types. Overall, our analysis provides a theoretically and empirically validated guideline for selecting an unfairness measure given a group fairness notion and testable data conditions.

Poster
An Equivariance Principle for Optimizer Design: Symmetry-Compatible Updates for Embeddings, LM Heads, and MoE Routers

Tim Tsz-Kit Lau, Weijie Su

展开摘要

Most modern deep neural networks are trained with Adam and related coordinate-wise adaptive optimizers, which are efficient but do not respect the natural matrix symmetries, geometry, and spectral structure of neural network parameters. We develop a symmetry-based principle for matrix-gradient optimization, showing that orthogonally equivariant optimizers, especially spectral optimizers such as stochastic spectral descent, Muon, Scion, and polar gradient methods, are better aligned with matrix layers than coordinate-wise updates. We show that coordinate-wise adaptivity breaks orthogonal equivariance, discards gradient spectral structure, and can be viewed as applying spectral updates to a pathological diagonal lifting of the parameter space. Motivated by permutation and orthogonal symmetries, we derive symmetry-compatible optimizers for different layer types, including linear layers, embedding and LM head matrices, and MoE routers. This yields one-sided spectral optimizers, row-norm optimizers, and hybrid row-norm/one-sided-spectral variants. Pre-training experiments on dense and sparse MoE language models, including Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style, and downsized gpt-oss models, show that replacing AdamW on large vocabulary-indexed matrices and router matrices with symmetry-compatible optimizers consistently improves final validation loss and sometimes training stability.

Poster
An Enigma of Artificial Reason: Investigating the Production-Evaluation Gap in Large Reasoning Models

Mingzhong Sun, Teresa Yeo, Armando Solar-Lezama, Tan Zhi-Xuan

展开摘要

Studies of human reasoning have shown that people are typically stronger at evaluating reasoning than producing it from scratch. In contrast, large reasoning models (LRMs) are trained to excel at producing long chains of reasoning to solve complex problems. How then do LRMs perform at evaluating reasons? We investigate this with the Valid-Answer-Invalid-Reasoning (VAIR) dataset: math problems and solutions with trivial reasoning flaws but valid answers, designed to isolate reasoning evaluation from the confound of reasoning production. Unlike humans, who we find are only 6\% worse at grading than solving such problems, we find a substantial production-evaluation gap in LRMs: frontier models score as low as 48\% when evaluating VAIR solutions, despite near-perfect solution production. Why this enigma? Through chain-of-thought (CoT) analysis, we find evidence of an answer confirmation bias: LRMs often produce then check for the correct answer instead of carefully verifying each step, fabricating rationalizations even when noticing anomalous reasoning. Linear probes corroborate this, showing that while LRM activations encode some representation of valid reasoning, they fail to robustly represent VAIR solutions as invalid. Causal patching of the final answer's representations causes LRM verdicts and activations to flip, demonstrating that answer validity is responsible for models' confirmation biases. These findings indicate an outstanding limitation in dominant approaches to reasoning training, which incentivize LRMs to produce and confirm reasoning towards correct answers, but not to robustly evaluate the underlying reasons.

Poster
An Empirical Study on Noisy Data and LLM Pretraining Loss Divergence

Qizhen (Irene) Zhang, Ankush Garg, Jakob Foerster, Niladri S. Chatterji, Kshitiz Malik, Mike Lewis

展开摘要

Large-scale pretraining datasets drive the success of large language models (LLMs). However, these web-scale corpora inevitably contain large amounts of noisy data due to unregulated web content or randomness inherent in data. Although LLM pretrainers often speculate that such noise contributes to instabilities in large-scale pretraining and, in the worst cases, loss divergence, this phenomenon remains poorly understood. In this work, we present a systematic empirical study of whether noisy data causes LLM pretraining divergences, costing more than 100,000 GPU hours in total. By injecting controlled, synthetic, uniform random noise into otherwise clean datasets, we analyze training dynamics across model sizes ranging from 480M to 5.2B parameters. We show that noisy data indeed induce training loss divergence, and that the probability of divergence depends strongly on the noise type, the amount of noise, and the model scale. We further find that noise-induced divergences exhibit activation patterns distinct from those caused by high learning rates, and we provide diagnostics that differentiate these two failure modes. Together, these results provide a large-scale, controlled characterization of how noisy data affects loss divergence in LLM pretraining.

Poster
An Embarrassingly Simple Graph Heuristic Reveals Shortcut-Solvable Benchmarks for Sequential Recommendation

Haoyu Han, Li Ma, Hanbing Wang, Bingheng Li, Daochen Zha, Chun How Tan, Huiji Gao, Xin Liu, Stephanie Moyerman, Sanjeev Katariya, Hui Liu, Jiliang Tang

展开摘要

Sequential recommendation is a central task in recommender systems, and recent research has increasingly shifted toward generative recommenders that leverage both sequential patterns and semantic item information. However, these methods are often evaluated on a small set of widely used benchmarks. This raises a natural question: do these benchmarks actually require the advanced modeling capabilities of modern generative recommenders? We conduct a benchmark audit using an intentionally simple graph heuristic: starting from only the most recent one or two interacted items, it retrieves candidates from a few-hop item-transition graph and ranks them with item-feature similarity. Surprisingly, despite its simplicity, this heuristic matches or outperforms a broad set of modern baselines on a variety of popular sequential recommendation benchmarks. For example, it achieves relative NDCG@10 improvements of 38.10\% and 44.18\% over the best competing baseline on the widely used Amazon Review Sports and CDs datasets, respectively. We further show that this phenomenon is not merely an artifact of a particular heuristic, but reflects shortcut solvability in existing benchmarks. Specifically, we identify three potential shortcut structures that could make next-item prediction easier than expected: low-branching local transition structure, feature-smooth transitions, and limited dependence on long user histories. These shortcuts do not need to appear simultaneously. Depending on the dataset, even one or two strong shortcut signals can make simple local retrieval highly competitive, while weakening the relevant signals allows more sophisticated models to show clearer benefits. Our broader evaluation across 14 diverse datasets further shows that model rankings change substantially with dataset properties, while the simple graph heuristic remains competitive on 10 out of 14 datasets. These findings suggest that strong performance on several standard sequential recommendation benchmarks may not faithfully reflect whether recent methods achieve the advanced modeling capabilities they aim to demonstrate. Rather than treating datasets as interchangeable leaderboards, we argue for more careful dataset selection and dataset-level diagnostic analysis when using benchmarks to support claims about the benefits of new recommendation models.

Poster
An Efficient Cross-modal Feature Reconstruction Model for Multimodal Multi-class Anomaly Detection

Zhenghan Wang, Zhihui Luo, Haitao Hu, Wengang Cheng

展开摘要

Industrial anomaly detection is evolving from category-specific detectors toward unified models capable of handling multiple categories simultaneously. However, in multi-class setting, greater inter-category diversity not only compels enhanced reconstruction capacity to model normal patterns but also exacerbates the “identical shortcut” problem, wherein anomalies are likewise well reconstructed. Moreover, existing unified frameworks are typically confined to single-modal (RGB) inputs and lack multimodal capability. To address these, we propose an Efficient Cross-modal Feature Reconstruction (ECFR) model for unified multi-class anomaly detection, which harnesses the inherent difficulty of cross-modal reconstruction to alleviate the shortcut issue and amplify anomaly discrimination. The framework is built on two core modules: an Adaptive Feature Interaction and Recalibration (AFIR) module and a Hybrid Attention Convolution (HAC) module. Through feature interaction, fusion, and reconstruction, our model achieves two key outcomes: learning normal patterns of multi-class objects while forcing reconstruction failures for anomalous inputs. Extensive experiments on the MVTec 3D-AD and Eyecandies datasets validate our approach, which achieves state-of-the-art performance in multimodal multi-class anomaly detection while requiring significantly fewer model parameters (44.63 M) and lower computational complexity (15.77 GFLOPs) compared to previous multimodal anomaly detection methods.

Poster
An Efficient Algorithm for Thresholding Monte Carlo Tree Search

Shoma Nameki, Atsuyoshi Nakamura, Junpei Komiyama, Koji Tabata

展开摘要

We introduce the Thresholding Monte Carlo Tree Search problem, in which, given a tree $\mathcal{T}$ and a threshold $\theta$, a player must answer whether the root node value of $\mathcal{T}$ is at least $\theta$ or not. In the given tree, 'MAX' or 'MIN' is labeled on each internal node, and the value of a 'MAX'-labeled ('MIN'-labeled) internal node is the maximum (minimum) of its child values. The value of a leaf node is the mean reward of an unknown distribution, from which the player can sample rewards. For this problem, we develop a $\delta$-correct sequential sampling algorithm based on the Track-and-Stop strategy that has asymptotically optimal sample complexity. We show that a ratio-based modification of the D-Tracking arm-pulling strategy leads to a substantial improvement in empirical sample complexity, as well as reducing the per-round computational cost from linear to logarithmic in the number of arms.

Poster
An Educated Guess: Deriving Statistically Aligned Gradient Estimates for Zeroth-Order Optimization

Justus F Hübotter, Serge Thill, Marcel A. J. van Gerven, Nasir Ahmad

展开摘要

As novel hardware architectures are explored for the purposes of machine learning, including neuromorphic chips, asynchronous systems, and ASICs, effective learning algorithms are highly desirable. While backpropagation is the standard for differentiation on standard hardware, its requirement for global synchrony and exact differentiability limits the exploration of exotic algorithms on non-traditional substrates and devices. Zeroth-order (ZO) optimization methods, such as SPSA and weight perturbation, offer a compelling alternative by estimating gradients through inference-only passes; however, these methods have historically suffered due to the "curse of dimensionality," where random perturbations become increasingly misaligned with the true gradient in high-dimensional spaces. In this work, we move beyond blind random noise by deriving zeroth-order search directions from activation geometry available during the forward pass. This leads to a family of Spectral Zeroth-Order methods, SZO-$\kappa$, whose directions solve local response-based gradient-guessing objectives under explicit stochastic assumptions. These methods spend perturbation budget on feature-supported directions that are more likely to produce informative loss responses. To further bridge the gap with first-order methods, we study covariance corrections and probe stabilizers, including search direction orthogonalization and bias corrections that reduce estimator error. Using multilayer and convolutional layers and networks, we analyze gradient alignment, bias-variance structure, and training stability of these methods relative to standard perturbation baselines and backpropagation.

Poster
An Axiomatic Analysis of DPO and NLHF as Reference-Dependent Probabilistic Voting Rules

Wesley H. Holliday, Adam Lesnikowski, Roberto-Rafael Maura-Rivero

展开摘要

Post-training methods for large language models (LLMs) implicitly aggregate diverse human feedback in ways that can be studied explicitly using social choice theory. In this paper, we adopt the perspective that for each post-training method, there is a corresponding reference-dependent probabilistic voting rule. Taking Direct Preference Optimization (DPO) and Nash Learning from Human Feedback (NLHF) as our case studies, we study axiomatic properties of the associated voting rules. First, we show that NLHF violates one of the central axioms of social choice, namely monotonicity, if and only if its KL regularization parameter is below a bound, whereas DPO satisfies monotonicity regardless of its regularization parameter. Second, we show that while DPO is often associated with Borda-style preference aggregation, when DPO is regarded as a probabilistic voting rule, it violates two variable population axioms that the standard probabilistic version of Borda satisfies. Third, we show that unlike DPO, NLHF violates a commutativity axiom from Bayesian epistemology, which has practical implications for staged post-training. Finally, we conduct an empirical study of the frequency and magnitude of axiom violations in the CoVal dataset using real base models.

Poster
An Assessment of Human vs. Model Uncertainty in Soft-Label Learning and Calibration

Maja Pavlovic, Silviu Paun, Massimo Poesio

展开摘要

Central to human-aligned AI is understanding the benefits of human-elicited labels over synthetic alternatives. While human soft-labels improve calibration by capturing uncertainty, prior studies conflate these benefits with the implicit correction of mislabeled data (mode shifts), obscuring true effects of soft-labels. We present a controlled audit of soft-label learning across MNIST and a synthetic variant, re-annotating subsets to extract human uncertainty. By decoupling soft-label supervision from underlying label mode shifts, we show that while human soft-labels do provide accuracy gains, their larger value lies in acting as a regularizer that improves model calibration on difficult samples and promotes stable convergence across training runs. Dataset cartography reveals models trained on human soft-labels mirror human uncertainty, whereas those trained on synthetic labels fail to align with humans. Broadly, this work provides a diagnostic testbed for human-AI uncertainty alignment.

Poster
An Analytical Model of Compute-limited Multistage Training Pipelines

Nishil Patel, Jin Hwa Lee, Basile Confavreux, Andrew Saxe

展开摘要

Modern foundation models are trained through multistage pipelines, typically beginning with unsupervised pretraining, followed by supervised fine-tuning (SFT) on step-by-step solutions, and reinforcement learning (RL) with bulk feedback. Empirically, the design of these stages, particularly the compute allocated to each, strongly affects performance. Yet our theoretical understanding of multistage training remains limited. Here, we introduce a simple, analytically tractable model of multistage training. The model reproduces several qualitative phenomena: SFT and RL can be complementary, pretraining quality imposes a performance ceiling with power-law scaling, optimal compute allocation varies with data quality and pretraining level, and RL can selectively amplify task-relevant pretrained skills. Our model distills the complex behaviors into key ingredients, allowing clearer understanding of how and where they arise. Overall, our model provides a theoretical starting point for explaining the phenomenology of multistage training pipelines in the compute-limited regime.

Poster
Amortized Structured Stochastic Variational Inference for Gaussian Process Latent Variable Models

Maksym Tretiakov, Sarah Filippi, Vincent Fortuin, Ruth Misener, Ruby Sedgwick, James Odgers

展开摘要

Many machine learning methods aim to approximate the lower-dimensional manifold on which the data lives. A desirable feature of such methods is that they should capture the epistemic uncertainty of this learned manifold. One model that achieves this is the Gaussian Process Latent Variable Model, in which a Gaussian Process (GP) mapping from the latent space provides an estimate of the uncertainty of the manifold. However, the effectiveness of this uncertainty estimation is limited by the mean-field variational approximation between the GP inducing points and the latent variables. In this work, we use Amortized Structured Stochastic Variational Inference to allow the variational posterior for the latent space to be conditionally dependent on the value of the inducing points. We demonstrate that this more flexible variational posterior improves several metrics related to the reconstruction of points on the data manifold.

Poster
Amortized Optimal Transport from Sliced Potentials

Minh-Phuc Truong, Khai Nguyen

展开摘要

We propose a novel amortized optimization method for predicting optimal transport (OT) plans across multiple pairs of measures by leveraging Kantorovich potentials derived from sliced OT. We introduce two amortization strategies: regression-based amortization (RA-OT) and objective-based amortization (OA-OT). In RA-OT, we formulate a functional regression model that treats Kantorovich potentials from the original OT problem as responses and those obtained from sliced OT as predictors, and estimate these models via least-squares methods. In OA-OT, we estimate the parameters of the functional model by optimizing the Kantorovich dual objective. In both approaches, the predicted OT plan is subsequently recovered from the estimated potentials. As amortized OT methods, both RA-OT and OA-OT enable efficient solutions to repeated OT problems across different measure pairs by reusing information learned from prior instances to rapidly approximate new solutions. Moreover, by exploiting the structure provided by sliced OT, the proposed models are more parsimonious, independent of specific structures of the measures, such as the number of atoms in the discrete case, while achieving high accuracy. We demonstrate the effectiveness of our approaches on tasks including MNIST digit transport, color transfer, supply-demand transportation on spherical data, and mini-batch OT conditional flow matching.

Poster
Amortized Molecular Optimization via Group Relative Policy Optimization

Muhammad Bin Javaid, Hasham Hussain, Ashima Khanna-Reiter, Berke Kisin, Jonathan Pirnay, Alexander Mitsos, Dominik G Grimm, Martin Grohe

展开摘要

In structurally constrained molecular optimization, state-of-the-art methods restart an expensive oracle-driven search from scratch for every new input structure, scaling poorly to settings with many starting structures or expensive oracles. While amortized approaches that learn a transferable policy could in principle remove this bottleneck, existing methods struggle to generalize to diverse structural constraints at inference time. We present AMORTIX, an amortized Graph Transformer model that natively supports such constraints, optimizing molecular structures in a single forward pass with zero inference-time oracle calls. A central challenge for amortized training in this domain is that optimization difficulty varies drastically across starting structures. We show that, under this heterogeneity, standard reinforcement learning methods fail to stabilize training, and address this by normalizing rewards within groups of completions sharing the same starting structure. We evaluate on structurally constrained single- and multi-target kinase inhibitor design, and on a few-shot prodrug case study. AMORTIX outperforms both amortized and instance-optimization baselines on goal-directed scaffold decoration and ranks first among amortized methods on the PMO benchmark; the prodrug case study further demonstrates transfer of a learned modification rule to unseen drug structures.

Poster