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

Topic: Simulation for Driving · 940 papers · Reset filters

MiCo: Microstructure-Consistent Flow Matching for Diffusion MRI Angular Super-Resolution

Zhixuan Zhou, Tingting Dan, Guorong Wu

Expand abstract

Diffusion-weighted MRI (dMRI) requires densely sampling $q$-space across many directions and shells, leading to long scan times that limit clinical use. Angular super-resolution (ASR) recovers a dense signal from a sparse acquisition, enabling substantially shorter scans. We propose a physics-guided flow matching framework for ASR built on a basic fact about dMRI: under the cumulant expansion, the diffusion signal naturally splits into a known linear function of the $q$-vectors plus a residual carrying higher-order angular detail. We exploit this twice. First, we initialize the flow from a closed-form dense DTI estimate, recasting ASR as residual flow matching from a physically meaningful starting point rather than uninformative noise. Second, we impose a microstructure-consistency loss, rooted in the physical principle that any two diffusion-weighted signals from the same voxel must be explained by a single underlying set of microstructural parameters. Together, these designs anchor predictions wherever the acquisition is informative while leaving the learned prior to model crossing-fiber and higher-order angular detail. On the Human Connectome Project (HCP-YA) and UK Biobank, our method matches or surpasses analytical baselines and recent deep learning methods on signal reconstruction, kurtosis scalar maps, and fiber-orientation reconstruction under aggressive subsampling.

Poster
MetaLoop: Benchmarking the Full Metacognitive Loop in LLMs

Nora Petrova, John Burden, Jerome Wynne

Expand abstract

We introduce MetaLoop, a seven-task benchmark that evaluates the full metacognitive loop in large language models: monitoring one's own uncertainty, translating that signal into action (abstaining, switching strategies, correcting errors), and updating from feedback. Existing evaluations test these components in isolation; MetaLoop produces integrated per-model profiles across all three. Evaluating 12 frontier models alongside 460 human participants, we find that (1) accuracy does not predict metacognitive ability (high-accuracy models routinely approve their own errors); (2) models universally produce calibration signals but fail to act on them; (3) forced self-explanation distinguishes genuine self-monitoring from brittle pattern-matching; and (4) humans and LLMs both show monitoring–control gaps but with different failure modes (humans are loss-averse, models overconfident), with clarification detection as the one clear human advantage. We release the benchmark, scoring code, model outputs, and human data.

Poster
MetaKE: Meta-Learning for Knowledge Editing Toward a Better Accuracy-Editability Trade-off

Shuxin Liu, Di Gao, Ou Wu

Expand abstract

Existing locate-then-edit Knowledge Editing (KE) methods typically decompose editing into two stages: upstream target representation optimization and downstream constrained parameter optimization. The optimization across the two stages is disconnected: upstream applies uniform regularization without observing downstream realization of the planned residual, hindering a refined accuracy–editability trade-off. Since this realization is request-specific and depends on downstream constraints, uniform regularization can over-shrink high-association requests, causing insufficient editing, while it can under-regularize low-association requests, producing over-large planned residuals that reduce downstream editability. To bridge this disconnect, we propose $\textbf{MetaKE}$ ($\textbf{Meta}$-learning for $\textbf{K}$nowledge $\textbf{E}$diting), a new framework that unifies upstream and downstream stages into a bi-level optimization problem. The inner level optimizes parameter updates for the target representation, while the outer level optimizes representation using feedback from downstream constraints, achieving a better semantic accuracy-editability trade-off. To avoid costly multi-layer backpropagation, we introduce a Structural Gradient Proxy to approximate and propagate this feedback. Extensive experiments show that MetaKE outperforms strong baselines, offering a new perspective on KE. Code at: https://anonymous.4open.science/r/MetaKE-C160.

Poster
Meta-TTRL: A Metacognitive Framework for Self-Improving Test-Time Reinforcement Learning for T2I Generation in Unified Multimodal Models

Lit Sin Tan, Junzhe Chen, Xiaolong Fu, Lichen Ma, Junshi Huang, Jianzhong Shi, Yan Li, Lijie Wen

Expand abstract

Test-time scaling (TTS) improves text-to-image (T2I) generation in unified multimodal models (UMMs), but existing TTS methods typically rely on frozen-parameter search or sampling, yielding only ephemeral, instance-level improvements. We propose Meta-TTRL, a metacognitive test-time reinforcement learning (TTRL) framework that converts test-time generation experience into learning signals, enabling UMMs to self-improve during T2I generation without external reward models. Meta-TTRL formulates test-time learning as a closed-loop interaction between an object-level generator and a meta-level introspector. The introspector decomposes prompts into structured verification rubrics and produces confidence-enhanced intrinsic monitoring signals, which are used to optimize the generator policy through reinforcement learning (RL). Extensive experiments demonstrate that Meta-TTRL generalizes well across three representative UMMs, including Janus-Pro-7B, BAGEL, and Qwen-Image, achieving significant gains on compositional reasoning tasks and multiple T2I benchmarks with limited data. Further analyses with external introspectors, RL leakage, and alternative monitoring signals reveal a key principle for effective TTRL: metacognitive synergy, where monitoring signals must align with the model's own optimization regime.

Poster
MentalHospital: A Virtual Environment for Evaluating Psychiatric Clinical Encounters

Yuming Yang, Xiao Sun, Yuanwei Zou, Zhengxiao Wu, Jiang Zhong, Haoyang Zeng, Jingwang Huang, Kaiwen Wei

Expand abstract

Large language models (LLMs) have shown strong performance on isolated psychiatric tasks, including dialogue, diagnosis, and treatment planning, yet existing benchmarks rarely simulate complete psychiatric clinical encounters. We introduce **MentalHospital**, a virtual evaluation environment for LLM-based psychiatric clinical encounters. MentalHospital instantiates the Subjective Interviewing, Objective Examination, Diagnostic Assessment, and Treatment Planning (S.O.A.P.) workflow, using skill-augmented standardized patients constructed from 1,193 de-identified psychiatric electronic health record (EHR) cases spanning all major ICD-11 categories and 76 disorders. Each encounter is assessed through a dual-track protocol that combines objective comparison against EHR-derived references with subjective assessment of clinical process quality. To scale specialist judgment, we develop **MentalEval**, five domain-specific evaluators covering communication empathy, interviewing professionalism, clinical-note quality, diagnostic rigor, and treatment appropriateness, trained with rubric-grounded SFT and expert-guided DPO. Survey responses from 22 clinicians support MentalHospital's clinical fidelity (3.88/5), while MentalEval achieves strong expert alignment with an average QWK of 0.944. Benchmarking shows that even the strongest LLM trails clinicians by 37.28 percentage points in objective psychiatric competence, with mental status assessment as a key bottleneck.

Poster
Memory Type Varies: Empowering LLM Agents for Long-Term Memory with Diverse Strategies

Yi Wen, Derong Xu, Pengyue Jia, Yichao Wang, Yingyi Zhang, Maolin Wang, Junyi Li, Wenlin Zhang, Xiaopeng Li, Yong Liu, Xiangyu Zhao

Expand abstract

The memory capabilities of Large Language Models (LLMs) have garnered increasing attention recently. Despite great success achieved, existing retrieval-based memory approaches typically overlook the differences between memories and employ a unified strategy to process all memories, leading to suboptimal performance. Thus, an intuitive question arises: can we categorize memory into different types and select appropriate strategies? However, given the topic-rich, scenario-complex, and boundary-blurred nature of memory scenarios, achieving precise classification of memories is not easy. To address this challenge, we propose a memory multi-class dataset in this paper, termed TriMEM, which provides precise annotations for memory types across diverse scenarios. Building upon this foundation, we propose a novel memory framework, named MemoType, which can adaptively recognize each memory and query type with the learned router model. With the memory and query routing, MemoType can retrieve the memory with corresponding query types and design tailored retrieval strategies, thereby enhancing the retrieval performance. Moreover, we theoretically prove that any single retrieval strategy is subject to a fundamental upper bound on its expected retrieval precision in multi-class corpora, leading to systematic precision degradation. Extensive experiments on three datasets demonstrate that MemoType consistently outperforms existing methods, achieving up to 16.18% improvement in Recall@1.

Poster
MemeEconomy : Do LLM Agents Trade Ethics for Survival?

Syed Nazmus Sakib, Nafiul Haque, Ahnaf Manan, M. M MORSHED, Shifat E. Arman

Expand abstract

Alignment-trained large language model (LLM) agents are increasingly deployed in autonomous, multi-round settings where ethical compliance can come into direct tension with task success. While prior safety literature extensively documents failure modes such as sycophancy and reward hacking, it largely overlooks a critical vulnerability: how agents behave under simulated existential economic pressure. We introduce **MemeEconomy**, a multimodal agentic market simulation in which LLM agents operate as meme investors across 300 real-world events. Agents select content from a five-tier harm taxonomy and target specific online communities under both low-pressure and high-pressure ("survival") environments. Each decision is recorded using a Belief--Desire--Intention (BDI) schema, which we employ as a diagnostic instrument to externalize the agent's event understanding, declared priorities, and risk acknowledgment. Across ten models spanning four frontier model families, harmful selection rates remain consistently high and increase by an average of **+15 percentage points** under economically induced survival pressure. Frontier alignment-trained models frequently commit to harm-tolerant selections at initialization while simultaneously suppressing the explicit harm acknowledgment that previously accompanied such decisions. In competitive tournament settings, ethically aligned agents achieve higher overall rankings, yet moderation penalties fail to meaningfully suppress harmful behavior. Harmful selection rates remain elevated even in rounds immediately following moderation removals, indicating that the moderation mechanism does not function as an effective deterrent. We further introduce **MemeAgent**, a 2B-parameter verifier trained on the simulation's BDI logs. MemeAgent substantially outperforms zero-shot frontier verifiers on in-distribution auditing tasks, achieving **87.8\%** accuracy compared to **63--74\%** for frontier baselines, while also generalizing to external multimodal harm benchmarks (**78\%** vs. **50--56\%**). Our findings demonstrate that the same structured reasoning traces that expose the failure mode can also be leveraged to train systems capable of detecting it.

Poster
MedCache: Training-Free Spatially Aware Caching for Accelerated Medical Video Generation

Ufaq Khan, Umair Nawaz, Sathira Silva, Numan Saeed, Muhammad U Sheikh, Muhammad Bilal, Junaid Qadir, Lena Maier-Hein, Yutong Xie, Muhammad Haris Khan

Expand abstract

Video world models are now being adapted to medical applications such as surgical simulation, procedural rehearsal, and synthetic data generation, but their inference cost remains a major obstacle to practical use. Training-free caching is attractive because it can be applied to existing checkpoints, yet current methods face a clear trade-off. Lightweight rules track only the average change in the video and treat all regions of the frame as equally important, which is a poor fit for medical video where the meaningful motion is often confined to a small region such as a tool tissue interaction. Richer rules recover spatial awareness by partially running the network on every step to inform the skip decision, but this decision cost is paid whether the step is skipped or not, and it eats into the speedup. We introduce \textbf{MedCache}, a training-free and probe-free cache for rectified-flow video generation. MedCache reads a spatial risk signal directly from the latent using cheap arithmetic, with no help from the network itself. It combines this risk signal with a prompt-derived domain prior and a lazy consistency check that runs only when reuse is no longer clearly safe. Across three open-world model families in Text2World and Image2World settings, MedCache improves both speed and quality over the strongest training free baselines in the medical regime it targets. On Cosmos-H-Surgical 2B Text2World, it raises Total quality from 0.767 to 0.837 over EasyCache while reducing latency from 185 seconds to 95 seconds. Ablations show that risk-weighted drift drives the runtime gain, while the prior and guard recover fidelity.

Poster
Mechanistic Critics for Sample-Efficient NPU Design-Space Exploration

Yicheng He, Yuqi Xue, Tianyang Wu, Jian Huang, Bin Hu

Expand abstract

NPU design-space exploration for LLM inference mainly relies on an expensive simulator that often returns only scalar feedback such as latency and power. Such feedback can rank evaluated configurations, but gives little information for deciding why a configuration is slow or which architecture parameter should be changed next. This limitation is especially severe because prefill and decode workloads can be limited by different resources, including compute throughput, HBM bandwidth, and inter-chip communication. We propose CriticNPU, a sample-efficient DSE framework built around an online-calibrated mechanistic critic for NPU architecture design. The critic uses explicit hardware limits to evaluate NPU configurations: before simulation, it estimates whether a proposed configuration is likely to improve over the current best and filters weak candidates; after simulation, it aggregates operator-level runtimes into prefill/decode bottleneck summaries and predicts the effect of candidate parameter changes. We instantiate the critic with a Roofline-style model that derives compute, memory, and communication limits from architecture parameters and calibrates per-operator-type effective limits using simulator observations. And CriticNPU combines this critic with a workload extractor and an architecture planner to propose configurations. Across six LLM workloads and 16 deployment configurations per model, CriticNPU achieves best-or-tied-best latency in 74% of prefill and 87.5% of decode settings while using roughly half the simulator calls of competitive DSE baselines. We conduct extensive ablation studies and detailed analyses to validate the effectiveness of our method from multiple perspectives.

Poster
Mean-Field Parallel Decoding for Discrete Diffusion Language Models

Tamim Zoabi, Ameen Ali, Liran Ringel, Lior Wolf

Expand abstract

Discrete diffusion language models enable parallel token generation, offering a pathway to low-latency decoding. However, selecting tokens independently by marginal confidence limits effective parallelism: tokens that appear reliable in isolation can form incompatible configurations when several positions are updated at once. We introduce a training-free decoding framework that coordinates these parallel updates. At each forward pass, the method assigns a commit score to each masked position and refines these scores using pairwise interactions derived from the model's predictive distributions. A variational relaxation yields a simple fixed-point update that suppresses conflicting simultaneous commitments within a single forward pass.This mechanism allows the decoder to commit more tokens in parallel while maintaining competitive generation quality. The method is lightweight, requires no auxiliary model or retraining, and drops into existing diffusion decoding pipelines without modification. Experiments on reasoning and code-generation benchmarks show consistent improvements in the quality-latency trade-off.

Poster
Margin Dynamics for Large Language Model Alignment

Xingzi Xu, Saygin Seyfioglu, Karim Bouyarmane

Expand abstract

Practitioners choose among LLM alignment methods largely by trial and error. We provide a theoretical foundation: each pairwise method defines a vector field on the probability simplex that governs how the model redistributes mass between preferred and dispreferred completions. Projecting onto a single preferred/dispreferred pair reduces this to a one-dimensional margin ODE with a closed-form solution. These solutions reveal that DPO margins grow as $O(\log t)$ while IPO margins converge exponentially to $1/(2\tau)$; that label noise induces a qualitative phase transition in DPO (unbounded $\to$ bounded) but affects IPO only quantitatively; and that gradient concentration under asymmetric weighting governs trainability at scale. We confirm all three predictions across five model families spanning four benchmarks. The closed-form solutions also enable principled method design: we prove that adding \emph{any} label smoothing $\varepsilon > 0$ to SimPO creates a finite margin attractor at $\Delta^* = \beta^{-1}\log((1{-}\varepsilon)/\varepsilon)$, while $\varepsilon{=}0$ (standard SimPO) has no attractor. We test this across three model scales (3.8B--9B) and confirm it: the theory-derived configuration (TPO-S, $\varepsilon{=}0.05$) achieves the highest mean IFEval score on all three models.

Poster
MapPolicy: Structure-Aware Imitation Learning for Robot Manipulation via Physically Constrained Scene Map

Zehao Du, Zheng Huang, Rui Li, Hongyu Liu, Jiude Wei, Jiachun Bao, Cewu Lu, Jianhua Sun

Expand abstract

Robot manipulation under partial observability requires policies to act from incomplete visual evidence while still reasoning about object structure, spatial layout, and physical relations. Existing imitation learning methods mainly encode currently visible observations, which makes them brittle when task-critical parts, contact regions, or support relations are occluded. We introduce Scene Map, a structured scene representation that models object shape, part structure, spatial layout, and physical constraints as a graph of manipulation-relevant primitives and relations. Scene Map turns occluded but task-relevant structure into explicit prediction targets and provides structural priors for inferring unseen components from partial observations. Building on Scene Map, we propose MapPolicy, an imitation learning framework that fuses structured map features with visual and robot state features for action prediction. MapPolicy uses structure-modulated graph attention to propagate global structural and spatial information through primitive features, and a physical constraint loss to regularize learned features with geometric and kinematic consistency. Across 49 manipulation tasks in three simulation benchmarks and five real world tasks, MapPolicy consistently improves over strong imitation learning baselines, with especially clear gains under severe occlusion. Our code will be released publicly to facilitate reproducibility and further research.

Poster
Map-Guided Caching: A Global Perspective for Efficient Diffusion Transformer

Jiaqi Ji, Ran Yang, Bo Wei, Hui Li, Hee Min Choi

Expand abstract

Feature caching is a promising way to accelerate Diffusion Transformers (DiTs) without modifying the model backbone, but existing methods largely rely on local heuristics or input-agnostic rules when deciding which intermediate features to cache and reuse. We present Map-Guided Caching, a feature caching framework that predicts a Temporal Variation Map (TVMap) for each denoising trajectory and uses the map to derive a global caching policy. The TVMap summarizes block-wise temporal variation across timesteps, while a lightweight Map Predictor estimates this variation from the prompt and initial noise before denoising starts. Given the predicted map, we formulate policy generation as a global path planning problem and combine a dynamic-programming-based orientation planner with a quantity planner that targets a desired speedup ratio. Across image generation, image editing, and video generation, the proposed method achieves up to $3.3\times$ speedup on FLUX.1 [dev] and HunyuanVideo, and $2.7\times$ on FLUX.1 Kontext [dev], while preserving competitive output quality. The advantage is especially pronounced on complex prompts, where local caching policies degrade more sharply. These results suggest that input-adaptive, globally planned caching is an effective alternative to purely local feature reuse policies.

Poster
Manifold-Guided Stereo-Monocular Refinement for Endoscopic Stereo Disparity Estimation

Jiewen Yang, Yi Zhong, Bin Pu, Xingbo Dong, Ming Zheng, Qika Lin, Kai He, Xuechen Zhang

Expand abstract

Dense stereo disparity estimation provides the geometric basis for 3D perception in computer-assisted endoscopy. Endoscopic stereo matching is difficult because reliable tissue correspondences often appear next to weak-texture, specular, occluded, or deformed regions within the same frame. Existing stereo and stereo-monocular methods improve cost aggregation or introduce monocular priors, but they often treat monocular guidance as a global feature source or an independently predicted candidate. Without evaluating the local support of the current stereo match during refinement, the update can reinforce corrupted correspondences in ambiguous regions or allow monocular priors to override reliable metric matches. To tackle these problems, we introduce Manifold-Stereo, a reliability-gated latent-flow framework for endoscopic stereo disparity estimation. The method refines a monocular disparity candidate through a conditional latent residual flow on an image-aligned correction space induced by stereo cost-volume geometry, then calibrates the corrected candidate to the current stereo disparity scale. A geometric consistency gate computed from left-right feature alignment regulates the candidate contribution during recurrent refinement, preserving stereo estimates in well-matched regions while supplying monocular context where local matching is ambiguous. Experiments on in-domain SCARED and zero-shot SERV-CT and Hamlyn show strong average disparity accuracy and competitive thresholded outlier rates against recent stereo and endoscopic baselines.

Poster
Manifold-Aligned Adversarial Perturbation for Anti-Customization under Diffusion-based Purification

Ruotian Liu, Zhen Chen, Zhiguo Yang, Xiangyu Yin, Peipei Xu, Wenjie Ruan

Expand abstract

Anti-customization techniques aim to protect visual content from unauthorized diffusion model personalization by adding imperceptible perturbations to training images. However, such protection can be significantly weakened by diffusion-based purification, which may remove adversarial perturbations before customization. We study representative anti-customization methods under diffusion purification and observe that existing methods exhibit similar off-manifold perturbation structures that are largely suppressed by purification. We then provide a principled geometric analysis that characterizes purification as a local contraction toward the data manifold, primarily suppressing perturbation components along normal directions. Motivated by this insight, we propose Manifold-Aligned Customization protection **MAC**, a practical purification-aware framework that optimizes downstream anti-customization objectives through a differentiable purifier and uses a structured perturbation parameterization to bias updates toward purification-retentive, manifold-aligned directions. Across experiments on CelebA-HQ, VGGFace2-HQ, and WikiArt under diffusion-based purification, MAC substantially improves protection effectiveness across most prompts and domains. Finally, frequency analysis and ImageNet and MS-COCO retention studies show that spectral signatures cannot predict purification outcomes, providing additional evidence for the proposed manifold-alignment perspective beyond a single data domain.

Poster
Manifold Prior Guided Deep Unfolding for Hyperspectral Image Reconstruction

Shuqi LIU, Yingkai Zhang, Lin Gu, Tao Zhang, Ying Fu

Expand abstract

Hyperspectral Image (HSI) reconstruction has made substantial progress with the deep unfolding framework by decomposing the problem into a projection module and a denoiser module. Nevertheless, existing methods still exhibit limitations of insufficient matching with HSI data. The issues lie in two aspects: 1) the data-consistency projection module applies a fixed gradient descent step while ignoring the spectral content adaptivity of HSI; 2)the denoiser module either suffers from high computational complexity or is limited to the local receptive field, failing to efficiently capture the global spatial-spectral dependencies. In this paper, we propose a deep unfolding framework centered on a two-stage manifold learning strategy to obtain a degradation-free structural prior, which is then strategically injected into both the projection and denoiser modules. In the data-consistency projection, a prior-guided spectral weight mechanism is introduced to adaptively modulate the weight of spectral band. In the denoiser module, a manifold-guided Mamba is proposed to achieve the structure-aware extraction of spatial-spectral features through prior modulation. Quantitative and visual comparisons on synthetic and real-world datasets show that the proposed method significantly outperforms existing state-of-the-art approaches, while maintaining the structural interpretability of the unfolding framework. To support reproducibility, the code will be publicly released upon acceptance.

Poster
Majority Bit-Aware Watermarking for Large Language Models

Jiahao Xu, Rui Hu, Olivera Kotevska, Zikai Zhang

Expand abstract

The growing deployment of Large Language Models (LLMs) has raised concerns about their misuse in generating harmful or deceptive content. To address this issue, watermarking methods have been proposed to embed identifiable multi-bit messages into generated text for misuse tracing. However, existing methods often suffer from a fundamental trade-off between text quality and decoding accuracy. In particular, they have to restrict the size of the preferred token set (i.e., green list) during encoding to maintain a detectable watermark signal for decoding, which inevitably degrades generation quality. To improve this trade-off, we propose a novel message encoding paradigm called \textit{majority bit-aware encoding}, which relaxes the watermark signal strength from the green list size. This strategy allows for a strong watermark signal to be preserved in generated texts even when using a large green list. We introduce two instantiations of this paradigm: MajorMark and MajorMark$^{+}$, where the latter is specifically optimized for long messages. Extensive experiments on state-of-the-art LLMs demonstrate that our methods achieve higher decoding accuracy and superior text quality compared to prior baselines. The code of the proposed methods is available \href{https://anonymous.4open.science/r/MajorMark}{here} for review.

Poster
MainFL: Intertemporal Data Valuation for Robust Auction-based Federated Learning

Bangqi Pan, Yun Xin, Jianfeng Lu, Gang Li, Shuqin Cao, Zhou Chendi

Expand abstract

Auction-based Federated Learning (AFL) provides a principled framework for incentivizing data owners and allocating decentralized training resources to data consumers. However, practical AFL systems face a coupled valuation and allocation challenge: data quality is uncertain before training, post-training feedback can be strategically distorted, and heterogeneous deadlines render late model updates ineffective for learning. We propose MainFL, an incentive-compatible intertemporal data valuation mechanism for robust auction-based federated learning. MainFL integrates three components: (i) a bounded mutual-information peer prediction rule that elicits truthful pre-training quality reports without requiring ground-truth labels; (ii) an intertemporal reliability calibration scheme that updates future valuations using post-session validation signals; and (iii) a deadline-induced layered clinching mechanism that allocates quality-weighted data while enforcing budget and feasibility constraints.We theoretically establish Bayes–Nash truthfulness for data owners, allocation feasibility and budget safety, and truthful residual-demand equilibrium for data consumers under standard clinching-auction regularity conditions. Empirical results on widely used federated learning benchmarks show that MainFL improves social welfare by up to 12.1% and model accuracy by up to 7.4% compared to state-of-the-art AFL baselines, while substantially reducing straggler-induced task timeouts.

Poster
Macrocanonical Generator Networks: data-efficient neural surrogates for amortized physics simulation

Niall Jeffrey, Benjamin Wandelt

Expand abstract

Many scientific simulations are computationally expensive, limiting their use in simulation-based inference, uncertainty quantification, and decision making. We introduce *macrocanonical generator networks*, a framework for learning fast, data-efficient neural surrogates for *amortized* simulation of stationary physical processes by matching multiscale statistics of a target process. Inspired by microcanonical maximum-entropy synthesis, which matches feature statistics per sample by inner-loop optimization, our approach instead trains a neural generator to satisfy the same constraints in expectation, preserving physically realistic sample-to-sample fluctuations. We prove that gradient-descent training induces preconditioned gradient-descent dynamics in sample space that inherit the symmetry-preservation properties of the microcanonical gradient descent of Bruna and Mallat, amortizing its per-sample inner-loop optimization into a single training run after which every sample is generated in one forward pass. We further show that residual generator parameterizations initialized near the identity, trained with a small learning rate, preserve entropy locally during early optimization and mitigate mode collapse. Experiments on cosmology and fluid simulations show that macrocanonical generators recover realistic samples from extremely small training sets, in some cases a single training example, while generating new realizations more than $10^6$ times faster per sample than the reference simulator and more than $10^3$ times faster than microcanonical gradient descent. The trained generators can be reused for downstream pretraining, including simulation-based inference, providing a data-efficient route to multi-fidelity scientific machine learning workflows.

Poster
MaRK: Markov-adapted Recurrent Kernels for Dynamic Operator Conditioning in State Space Models

Syed Ibrahim Omer, Ginny Wong, Xiangyu Zhao

Expand abstract

State Space Models (SSMs) offer an efficient alternative to Transformers for sequence modeling, yet conditioning pre-trained SSMs for iterative generation typically operates outside the recurrent operator, through input injection or activation modulation. While such mechanisms expose the model to conditioning information, they leave the underlying temporal dynamics fixed. We introduce MaRK (Markov-adapted Recurrent Kernels), a dynamic operator-conditioning framework that maps context vectors directly into bounded modulations of a frozen SSM's recurrence ($A$), read-in ($B$), read-out ($C$), skip ($D$), and discretization ($\Delta$) parameters. Viewed through the lens of Linear Parameter-Varying systems, MaRK induces a context-indexed family of Markov parameter sequences, allowing each diffusion timestep to reshape the model's input-output memory kernel. We instantiate MaRK on a frozen 111M-parameter Hydra SSM backbone and study three adapter geometries: Hypernet, Chebyshev polynomial, and Discrete Cosine Transform kernels. Since these adapters modify the Markov parameter sequence through low-rank auxiliary maps on the frozen backbone, parameter-efficient fine-tuning arises as a structural consequence of the adaptation mechanism itself, requiring only 7.5--13M trainable auxiliary parameters to transition from a bidirectional objective to an iterative diffusion regime. The bounded recurrence parameterization further yields an analytic Affine Quadratic Stability certificate for the modulated recurrence. Through synthetic LPV recovery experiments and Markov-operator diagnostics, we show that MaRK recovers coordinate-invariant temporal operators under matched assumptions and produces distinct, stable timestep-conditioned memory profiles. Empirically, the Chebyshev variant yields the strongest performance, achieving an average validation loss of 2.63, followed by the DCT (2.66) and Hypernet (3.60) geometries. Together, these results provide initial evidence that dynamic operator modulation is a principled operator-level conditioning mechanism for adapting SSMs beyond input-stream injection and adaptive normalization.

Poster