论文检索

Applications · Time Series

Haiqi Jiang, Hui Xiong

Conventional time-series discriminative forecasting relies on point-wise regression, which inherently induces over-smoothing and fails to capture stochastic volatility in complex systems. While first-order generative flow matching methods mitigate this issue, they ignore system inertia, resulting in phase-space ambiguities and high sensitivity to noise. We introduce KineFlow, a generative time-series forecasting framework that augments flow matching with a phase-space Neural Acceleration Field, treating exogenous inputs as driving forces that produce gradual momentum shifts rather than abrupt state perturbations. This second-order formulation serves as a structural filter via double integration, suppressing high-frequency noise and producing robust, physically consistent predictions. Extensive experiments on six real-world benchmarks demonstrate that KineFlow achieves an average 15% MSE improvement over discriminative baselines and an 8% gain in CRPS compared to state-of-the-art generative methods.

Optimization · Everything Else

Zihao Zhao, Kai-Chia Mo, Shing-Hei Ho, Brandon Amos, Kai Wang

Differentiable optimization studies how to embed a mathematical program as a differentiable layer in machine learning pipelines. However, existing approaches typically rely on implicit differentiation, involving expensive Hessian computation while differentiating through optimality conditions. To address this challenge, we formulate the differentiable optimization problem as a bilevel optimization instance. We construct a new active-set Lagrangian as a proxy to compute an $\epsilon$-approximate hypergradient using only near-constant $O(\log (1/\epsilon))$ first-order information. We also show that applying this efficient hypergradient oracle to constrained bilevel optimization improves the overall gradient complexity to $\tilde{O}(\delta^{-1}\epsilon^{-3})$ to reach a $(\delta, \epsilon)$-Goldstein stationary point. We implement our method `FFOLayer`, as a drop-in Python library compatible with existing differentiable optimization solvers. Our algorithm shows significantly faster computation with similar convergence compared to other existing solvers. Our code is available [here](https://anonymous.4open.science/anonymize/FFOLayer-B78B).

Reinforcement Learning · Everything Else

Fang YUAN, Quanjun Yin, Siqi Shen, Yuxiang Xie, Junqiang Yang, Long Qin, Junjie Zeng, Qinglun Li

Unsupervised Environment Design (UED) offers a promising paradigm for improving reinforcement learning generalization by adaptively shaping training environments, but it requires reliable environment evaluation to remain effective. However, existing UED methods evaluate environments using indirect proxy signals such as regret, value-based errors, or Monte Carlo, which suffer from bias, high variance, or substantial computational overhead and fail to reflect agent realized learning progress. To address these limitations, we propose Parameter Change Environment Design (PACE), which evaluates an environment through the policy parameter change induced by training on that environment, directly grounding environment selection in realized learning progress. Specifically, PACE assigns environment value using a first-order approximation of the policy optimization objective, where the improvement induced by an environment is proportional to the squared $\ell_2$ norm of the corresponding parameter update, enabling low-variance and computation-efficient evaluation without additional rollouts. Experiments on MiniGrid and Craftax show that PACE consistently outperforms established UED baselines, achieving higher IQM and smaller Optimality Gap on OOD evaluations, including an IQM of 96.4% and an Optimality Gap of 17.2% on MiniGrid.

Deep Learning · Graph Neural Networks

Pablo Barcelo, Fabian Jogl, Alexander Kozachinskiy, Matthias Lanzinger, Stefan Neumann, Cristobal Rojas

Graph neural networks (GNNs) are widely used in graph learning and most architectures propagate information by passing messages between vertices. In this work, we shift our attention to GNNs that perform message passing on *edges* and introduce EB-1WL, an edge-based color-refinement test, and a corresponding architecture, EB-GNN. Our EB-GNN architecture is inspired by the classic triangle-counting algorithm of Chiba and Nishizeki and passes messages along edges and triangles. Our contributions are as follows: 1. Theoretically, we show that EB-1WL is significantly more expressive than 1WL. We provide a complete logical characterization of EB-1WL in first-order logic, along with distinguishability results via homomorphism counting. To the best of our knowledge, EB-GNN has the strongest theoretical expressivity guarantees among edge-based message-passing GNNs in the literature. 2. Unlike many GNN architectures that are more expressive than 1WL, we prove that EB-1WL and EB-GNN admit near-linear time and memory usage on practical graph learning workloads. 3. We show in experiments that EB-GNN is a highly efficient general-purpose architecture: it substantially outperforms simple MPNNs and remains competitive with task-specialized state-of-the-art GNNs at substantially lower computational cost.

Optimization · Large Scale, Parallel and Distributed

Yaoshuai Ma, Xiao Wang, Wei Yao, Jin Zhang

Distributed optimization over time-varying directed graphs has shown promising performance in addressing challenges posed by complex communication constraints in real-world scenarios. In many practical settings, however, the direct application of distributed optimization algorithms encounters additional difficulties, most notably hyperparameter tuning, which our empirical observations suggest can be effectively mitigated by integrating bilevel optimization. Motivated by these findings, we study distributed bilevel optimization over time-varying directed networks, a problem that remains largely unexplored due to the compounded challenges arising from consensus bias in dynamic unbalanced communication and the nested optimization structure. In this work, we propose a fully first-order distributed gradient-based algorithm that integrates the Push–Pull (also known as AB) communication strategy with a value function-based penalty method and establish its non-asymptotic convergence properties. Notably, a simplified variant of our analysis framework for nonconvex single-level distributed optimization establishes a convergence rate for the Push–Pull algorithm, thereby resolving an open question concerning its convergence over time-varying directed graphs. Empirical evaluations across diverse tasks, including hyperparameter tuning, data hyper-cleaning, and reinforcement learning, validate the effectiveness and efficiency of the proposed algorithm.

Social Aspects · Accountability, Transparency, and Interpretability

Jinyang Liu, Munir Hiabu

Interpretable machine learning requires models that are accurate and structurally faithful to the data. Existing explainability methods rely heavily on additive representations (e.g., GAMs, SHAP, functional ANOVA), which can suffer from signal cancellation and extrapolation in presence of strong interactions. We propose Tensor Separation Learning (TSL), a regression model that learns a sum of separable (rank-1) tensor products via an orthogonal greedy algorithm. By enforcing separability, TSL avoids the information loss inherent in additive projections caused by marginalizing higher-order interactions. The learned TSL model can be fully reconstructed from first-order partial dependence functions of its fitted factors. We establish approximation-rate guarantees for functions with bounded mixed $ p $-th order partial derivatives and demonstrate that TSL competes with black-box models on regression benchmarks. Crucially, TSL improves interpretability by factorizing interactions, allowing users to explicitly disentangle the magnitude of an effect from its direction directly via the fitted factors.

Deep Learning · Theory

Yuling Jiao, Yanming Lai, Defeng Sun, Yang Wang, Bokai Yan

We develop approximation and statistical theory for standard Transformer networks in sequence modeling. Given a sequence-to-sequence target on $[0,1]^{d_x \times n}$ whose entries are $\gamma$-H\"older for $\gamma \in (0,1]$ or belong to a first-order Sobolev class, we establish explicit $L^p$-approximation bounds for all $p \in [1,\infty]$, including the previously elusive endpoint $p=\infty$ under softmax attention. In particular, achieving error $\varepsilon$ in $L^p$-norm requires $\mathcal{O}(\varepsilon^{-d_x n/\gamma})$ parameters for $\gamma$-H\"older targets and $\mathcal{O}(\varepsilon^{-d_x n})$ parameters for Sobolev targets, matching the best known scalings in ambient dimension $d_x n$. We further study nonparametric regression with sequential and dependent observations using Transformer networks. Assuming stationary $\beta$-mixing covariates whose temporal dependence weakens over time, we analyze a sliding-window empirical risk minimization procedure and establish excess-risk guarantees for the Transformer-based estimators. Our analysis clarifies the role of attention and enables extensions beyond softmax.

General Machine Learning · Causality

Medha Agarwal, Alex Luedtke

We introduce the Sinkhorn treatment effect, an optimal transport measure of divergence between counterfactual distributions. Unlike classical quantities such as the average treatment effect, this measure captures differences across entire distributions. We analyze this divergence as a statistical functional and show it can be written as a smooth transformation of counterfactual mean embeddings with an appropriate kernel. This characterization allows us to establish first-order pathwise differentiability in general, and second-order pathwise differentiability under the null hypothesis of equal counterfactual distributions. Leveraging this smoothness, we construct debiased estimators and use them to obtain asymptotically valid tests for distributional treatment effects. Experiments on simulated and image data demonstrate the practical advantages of our estimator and testing procedure.

Optimization · Zero-order and Black-box Optimization

Zhe Li, Bicheng Ying, Zidong Liu, Haibo Yang

Classical optimization theory establishes that zeroth-order (ZO) algorithms suffer from a dimension-dependent slowdown, with convergence rates typically scaling with the model dimension compared to first-order methods. However, in contrast to these theoretical expectations, a growing body of recent work demonstrates the successful application of ZO methods to fine-tuning Large Language Models (LLMs) with billions of parameters. To explain this paradox, we derive the one-step learning dynamics of ZO SGD, where the empirical Neural Tangent Kernel (eNTK) naturally emerges as the key term governing the learning behavior. Inspection of the eNTK produced by ZO-SGD reveals that each element corresponds to the inner product of neural tangent vectors projected onto a random low-dimensional subspace. Thus, by invoking the Johnson-Lindenstrauss Lemma, our analysis shows that the fidelity of the ZO eNTK is governed primarily by the number of perturbations. Crucially, the approximation error depends on the model output size rather than the massive parameter dimension. This dimension-free property provides a theoretical justification for the scalability of ZO methods to LLMs finetuning tasks. We believe that this kernel-based framework offers a novel perspective for understanding ZO methods within the context of learning dynamics.

Optimization · Large Scale, Parallel and Distributed

Daniel Berg Thomsen, Adrien Taylor, Aymeric Dieuleveut

Communication costs are a major bottleneck in distributed learning and first-order optimization. A common approach to alleviate this issue is to compress the gradient information exchanged between agents. However, such compression typically degrades the convergence guarantees of gradient-based methods. Error feedback mechanisms provide a simple and computationally cheap remedy for this issue, but numerous variants have been proposed, and their relative performance remains poorly understood. In this work, we provide tight convergence analyses for two of the main error feedback algorithms in the literature, namely error feedback and EF21. Our results hold independently of the number of participating agents and rely on the construction of novel Lyapunov functions that recover the known best guarantees in the single-agent regime.

Probabilistic Methods · Everything Else

Peiman Mohseni, Nick Duffield, Raymond K. W. Wong

Many phenomena in science and engineering involve unknown latent functions observed through finite, irregularly sampled measurements. Neural processes (NPs) offer a powerful framework for probabilistic functional inference by bridging stochastic processes with deep learning. In many domains, these functions exhibit symmetries---most notably translation equivariance---that can be exploited to improve sample efficiency and generalization. Existing translation-equivariant NPs, however, have two key limitations: (i) they are constructed by stacking generic components with nonlinearities, obscuring the induced function class, thus limiting interpretability; and (ii) convolutional designs are based on localized receptive fields and require dense discretization, while attention-based methods avoid these issues but scale quadratically with the number of observations. We address these challenges through two contributions. First, we characterize continuous translation-equivariant operators through their Volterra expansions, representing them as sums of higher-order convolutions. This yields analytical transparency while remaining amenable to efficient approximation by first-order convolution operators. Second, we introduce set Fourier convolutions (SFConvs), a frequency-domain parameterization that operates directly on irregularly sampled sets. SFConvs achieve global receptive fields without spatial discretization and scale linearly in the number of observations. Building on these ideas, we propose two families of conditional NPs (CNPs): SFConvCNPs, constructed by stacking SFConv blocks with nonlinearies, and SFVConvCNPs, which integrate the Volterra formulation. Experiments on synthetic and real-world datasets demonstrate the efficacy of our methods compared to the state-of-the-art baselines.

Theory · Probabilistic Methods

Zijing Ou, Chinmay Pani, Yingzhen Li

Discrete diffusion models have become highly effective across various domains. However, real-world applications often require the generative process to adhere to certain constraints. To this end, we propose a Sequential Monte Carlo (SMC) framework that enables scalable inference-time control of discrete diffusion models through principled importance weighting and optimal proposal construction. Specifically, our approach derives tractable importance weights for a range of intermediate targets and characterises the optimal proposal, for which we develop two practical approximations: a first-order gradient-based approximation and an amortised proposal trained to minimise the log-variance of the importance weights. Empirical results across synthetic tasks, language modelling, biology design, and text-to-image generation demonstrate that our framework enhances controllability and sample quality, highlighting the effectiveness of SMC as a versatile recipe for scaling discrete diffusion models at inference time.

General Machine Learning · Transfer, Multitask and Meta-learning

Filippo Rinaldi, Aniello Panariello, Giacomo Salici, Fengyuan Liu, Marco Ciccone, Angelo Porrello, Simone Calderara

When a new release of a foundation model is published, practitioners typically need to repeat fine-tuning, even if the same task was already tackled in the previous version. A promising alternative is to reuse the parameter changes (i.e., task vectors) that capture how a model adapts to a specific task. However, these vectors often fail to transfer across different pre-trained models because their parameter spaces are misaligned. In this work, we show that successful transfer depends strongly on the gradient-sign structure of the new model. Based on this insight, we propose GradFix, which approximates the ideal sign structure and leverages it to transfer knowledge using only a handful of labeled samples. Notably, this requires no additional fine-tuning: we only compute a few target-model gradients without parameter updates and mask the source task vector accordingly. This yields an update that is locally aligned with the target loss landscape, effectively rebasing the task vector onto the new pre-training. We provide a theoretical guarantee that our method ensures first-order descent. Empirically, we demonstrate significant performance gains on vision and language benchmarks, consistently outperforming naive task vector addition and few-shot fine-tuning. We further show that transporting task vectors improves multi-task and multi-source model merging. Code is available at https://github.com/fillo-rinaldi/GradFix.

Theory · Interpretability and Visualization

Garrett Baker, George Wang, Jesse Hoogland, Vinayak Pathak, Daniel Murfet

We develop a linear response framework for interpretability that treats a neural network as a Bayesian statistical mechanical system. A small perturbation of the data distribution, for example shifting the Pile toward GitHub or legal text, induces a first-order change in the posterior expectation of an observable localized on a chosen component of the network. The resulting susceptibility can be estimated efficiently with local SGLD samples and factorizes into signed, per-token contributions that serve as attribution scores. We combine these susceptibilities into a response matrix whose low-rank structure separates functional modules such as multigram and induction heads in a 3M-parameter transformer.

Theory · Interpretability and Visualization

Zhenshuo Zhang, Minxuan Duan, Hongyang Zhang

Modern AI agents such as large language models are trained on diverse tasks---translation, code generation, mathematical reasoning, and text prediction---simultaneously. A key question is to quantify how each individual training task influences performance on a target task, a problem we refer to as task attribution. The direct approach, leave-one-out retraining, measures the effect of removing each task, but is computationally infeasible at scale. An alternative approach that builds surrogate models to predict a target task's performance for any subset of training tasks has emerged in recent literature. Prior work focuses on linear surrogate models, which capture first-order relationships, but miss nonlinear interactions such as synergy, antagonism, or XOR-type effects. In this paper, we first consider a unified task weighting framework for analyzing task attribution methods, and show a new connection between linear surrogate models and influence functions through a second-order analysis. Then, we introduce kernel surrogate models, which more effectively represent second-order task interactions. To efficiently learn the kernel surrogate, we develop a gradient-based estimation procedure that leverages a first-order approximation of pretrained models; empirically, this yields accurate surrogate estimates with less than 2% relative error without repeated retraining. Experiments across multiple domains---including mathematical reasoning in transformers, in-context learning, and multi-objective reinforcement learning---demonstrate the effectiveness of kernel surrogate models. They achieve a 25% higher correlation with the leave-one-out ground truth than linear surrogates and influence-function baselines, enabling more accurate and scalable task attribution. When used for downstream task selection, kernel surrogate models further yield a 40% improvement in demonstration selection for in-context learning and multi-objective reinforcement learning benchmarks.

Theory · Game Theory

Denizalp Goktas, Sadie Zhao, Yiling Chen, Amy Greenwald

In this paper, we study a generalization of Markov games and pseudo-games that we call Markov pseudo-games, which like the former, captures time and uncertainty, and like the latter, allows for the players’ actions to determine the set of actions available to the other players. In the same vein as Arrow and Debreu, we intend for this model to be rich enough to encapsulate a broad mathematical framework for modeling economies. We then prove the existence of a game-theoretic equilibrium in our model, which in turn implies the existence of a general equilibrium in the corresponding economies. Finally, going beyond Arrow and Debreu, we introduce a solution method for Markov pseudo-games, and prove its polynomial-time convergence. We then provide an application of Markov pseudo-games to infinite-horizon Markov exchange economies, a stochastic economic model that extends Radner’s stochastic exchange economy and Magill and Quinzii’s infinite horizon incomplete markets model. We show that under suitable assumptions, the solutions of any infinite horizon Markov exchange economy (i.e., recursive Radner equilibria—RRE) can be formulated as the solution to a concave Markov pseudo-game, thus establishing the existence of RRE, and providing first-order methods for approximating RRE. Finally, we demonstrate the effectiveness of our approach in practice by building the corresponding generative adversarial policy neural network, and using it to compute RRE in a variety of infinite-horizon Markov exchange economies.

Optimization · Non-Convex

Francisco Patitucci, Ruichen Jiang, Aryan Mokhtari

A recent breakthrough in nonconvex optimization is the online-to-nonconvex conversion framework of Cutkosky et al. (2023), which reformulates the task of finding an $\varepsilon$-first-order stationary point as an online learning problem. When both the gradient and the Hessian are Lipschitz continuous, instantiating this framework with two different online learners achieves a complexity of $ \mathcal{O}(\varepsilon^{-1.75}\log(1/\varepsilon)) $ in the deterministic case and a complexity of $ \mathcal{O}(\varepsilon^{-3.5}) $ in the stochastic case. However, this approach suffers from several limitations: (i) the deterministic method relies on a complex double-loop scheme that solves a fixed-point equation to construct hint vectors for an optimistic online learner, introducing an extra logarithmic factor; (ii) the stochastic method assumes a bounded second-order moment of the stochastic gradient, which is stronger than standard variance bounds; and (iii) different online learning algorithms are used in the two settings. In this paper, we address these issues by introducing an online optimistic gradient method based on a novel **doubly optimistic hint function**. Specifically, we use the gradient at an extrapolated point as the hint, motivated by two optimistic assumptions: that the difference between the hint and the target gradient remains near constant, and that consecutive update directions change slowly due to smoothness. Our method eliminates the need for a double loop and removes the logarithmic factor. Furthermore, by simply replacing full gradients with stochastic gradients and under the standard assumption that their variance is bounded by $\sigma^2$, we obtain a unified algorithm with complexity $\mathcal{O}(\varepsilon^{-1.75} + \sigma^2 \varepsilon^{-3.5})$, smoothly interpolating between the best-known deterministic rate and the optimal stochastic rate.

General Machine Learning · Transfer, Multitask and Meta-learning

Shaocong Ma, Peiran Yu, Heng Huang

Fine-tuning Large Language Models (LLMs) typically involves either full fine-tuning, which updates all model parameters, or Parameter-Efficient Fine-Tuning (PEFT), which adjusts a small subset of parameters. However, both approaches have inherent limitations: full fine-tuning is computationally expensive, while PEFT often struggles to learn new knowledge and exhibits suboptimal performance. To overcome these issues, we propose a novel *hybrid fine-tuning* approach that jointly updates both LLMs and PEFT modules using a combination of zeroth-order and first-order optimization methods. To analyze our new algorithm, we develop a theoretical framework centered on the concept of *hybrid smoothness condition*, which accounts for the heterogeneous nature of the optimization landscape in joint LLM and PEFT training. We derive a rigorous convergence analysis for the convergence of reshuffling-type SGD algorithm under multiple learning rates and demonstrate its effectiveness through extensive empirical studies across various downstream tasks and model architectures. On the practical side, our results demonstrate consistent performance improvement, making the approach a viable solution for large-scale language model fine-tuning.

Jenny Zhang, Shengran Hu, Cong Lu, Robert Lange, Jeff Clune

Most of today's AI systems are constrained by human-designed, fixed architectures and cannot autonomously and continuously improve themselves. The scientific method, on the other hand, is a cumulative and open-ended system, where each innovation builds upon previous artifacts, enabling future discoveries. There is growing hope that the current manual process of advancing AI could itself be automated. If done safely, such automation would accelerate AI development and allow us to reap its benefits much sooner. This prospect raises the question of how AI systems can endlessly improve themselves while getting better at solving relevant problems. Meta-learning can automate the discovery of novel algorithms, but is limited by first-order improvements and the human design of a suitable search space. The Gödel machine proposed a theoretical alternative: a self-improving AI that repeatedly modifies itself in a provably beneficial manner. Unfortunately, proving that most changes are net beneficial is impossible in practice. We introduce the Darwin Gödel Machine (DGM), a novel self-improving system that iteratively modifies its own code (thereby also improving its ability to modify its own codebase) and empirically validates each change using coding benchmarks. Inspired by Darwinian evolution and open-endedness research, the DGM grows an archive of generated coding agents. It samples agents from this archive, which self-modify to create new, interesting versions of themselves. This open-ended exploration forms a growing tree of diverse, high-quality agents and allows the parallel exploration of many different paths through the search space. Empirically, the DGM automatically improves its coding capabilities (e.g., better code editing tools, long-context window management, peer-review mechanisms), increasing performance on SWE-bench from 20.0% to 50.0%, and on Polyglot from 14.2% to 30.7%. Furthermore, the DGM significantly outperforms baselines without self-improvement or open-ended exploration. All experiments were done with safety precautions (e.g., sandboxing, human oversight). Overall, the DGM represents a significant step toward self-improving AI, capable of gathering its own stepping stones along a path that unfolds into endless innovation.

Optimization · Non-Convex

Juan Ramirez, Simon Lacoste-Julien

Constrained optimization is a powerful framework for enforcing requirements on neural networks. These constrained deep learning problems are typically solved using first-order methods on their min-max Lagrangian formulation, but such approaches often suffer from oscillations and can fail to find all local solutions. While the Augmented Lagrangian method (ALM) addresses these issues, practitioners often favor dual optimistic ascent schemes (PI control) on the standard Lagrangian, which perform well empirically but lack formal guarantees. In this paper, we establish a previously unknown equivalence between these approaches: dual optimistic ascent on the Lagrangian is equivalent to gradient descent-ascent on the Augmented Lagrangian. This finding allows us to transfer the robust theoretical guarantees of the ALM to the dual optimistic setting, proving it converges linearly to all local solutions. Furthermore, the equivalence provides principled guidance for tuning the optimism hyper-parameter. Our work closes a critical gap between the empirical success of dual optimistic methods and their theoretical foundation in the single-step, first-order regime commonly used in constrained deep learning.