论文检索

Optimization · Non-Convex

Leonardo Galli, Curtis Fox, Wiebke Bartolomaeus, Mark Schmidt, Holger Rauhut

The training of neural networks often entails objective functions that are not globally $L$-smooth. For these functions, it is both theoretically and practically difficult to reply to the question: what is the largest possible step size that ensures the convergence of gradient descent (GD)? We address this longstanding open question in deep learning by providing a unifying definition of "large'' step sizes that requires only local Lipschitz (or even Hölder) continuity of the gradient. We design first-order adaptive methods that provably yield large step sizes and show that they operate at the edge of stability (EoS) right from the start of the training. In particular, the loss decreases nonmonotonically and the product between the step size and sharpness, i.e., the largest eigenvalue of the hessian, stays above the EoS threshold of 2 throughout training. Using our method, we are also able to minimize the sharpness all the way down to its global minimum. Contrary to expectation, we find that encountering globally-flat regions too early in the training may both slow down convergence and jeopardize the generalization ability of the network. Exploiting a self-stabilization argument, we allow GD to enter slightly sharper valleys and turn unsuccessful training runs into very successful ones.

Optimization · Non-Convex

Yutong Chao, Michal Ciebielski, Jalal Etesami, Majid Khadiv

In this paper, we study a class of non-convex optimization problems known as multi-affine quadratic equality constrained problems, which appear in various applications--from generating feasible force trajectories in robotic locomotion and manipulation to training neural networks. Although these problems are generally non-convex, they exhibit convexity or related properties when all variables except one are fixed. Under mild assumptions, we prove that the alternating direction method of multipliers (ADMM) converges when applied to this class of problems. Furthermore, when the "degree" of non-convexity in the constraints remains within certain bounds, we show that ADMM achieves a linear convergence rate. We validate our theoretical results through practical examples in robotic locomotion.

Theory · Everything Else

Hien Dang, Pratik Patil, Alessandro Rinaldo

Self-distillation (SD), retraining a student on a mixture of ground-truth labels and a teacher’s own predictions using the same architecture and training data, often improves generalization empirically, but it is unclear when improvement is guaranteed. We study SD for ridge regression with an unconstrained mixing weight $\xi \in \mathbb{R}$. Conditional on the training data, without any distributional assumptions, we prove that for any squared prediction risk $R$ (including out-of-distribution), the optimally mixed student strictly improves upon the ridge teacher at every regularization level $\lambda$ where the teacher risk is not stationary ($R'(\lambda) \neq 0$). We also characterize the optimal mixing weight $\xi^\star$ in terms of the risk derivative $R'$, showing that it can surprisingly be negative. To further quantify SD risk improvements, we derive exact risk asymptotics in the proportional asymptotics regime for general anisotropic covariance and deterministic signals. Finally, we propose a consistent one-shot tuning method to estimate $\xi^\star$ without retraining, sample splitting, or grid search. Experiments on real-world tasks and pre-trained neural network features validate our theory and tuning method.

Deep Learning · Graph Neural Networks

Junshu Sun, Wanxing Chang, Qingming Huang, Shuhui Wang

Graph neural networks (GNNs) tightly couple their input-output parameters to dataset-specific feature spaces and target sets, exhibiting limited transferability across different datasets. In contrast, language models (LMs) generalize flexibly via a unified input-output interface, motivating recent attempts to adapt LMs to graph tasks. However, existing methods struggle to encode whole-graph information, leading to potential information loss and suboptimal graph understanding. In this work, we propose a novel weight-level information injection paradigm for adapting LMs to graph tasks. This paradigm injects whole-graph information by generating task-specific weight updates that interact directly with hidden representations. Instantiating this paradigm following low-rank adaptation (LoRA), we introduce GaRA, a Graph-aware LoRA generation model. GaRA constructs low-rank weight updates conditioned on the original graph structures and constrains the norm of the generated updates, thus injecting whole-graph information and avoiding the optimization bias in the weight generation. Empirical studies demonstrate that GaRA consistently outperforms baselines on zero-shot graph learning tasks. Codes are provided in the supplementary material.

Optimization · Non-Convex

Chanwoong Park, Uijeong Jang, Ernest Ryu, Insoon Yang

Sharpness-Aware Minimization (SAM) is widely used to seek flatter minima—often linked to better generalization. In its standard implementation, SAM updates the current iterate using the loss gradient evaluated at a point perturbed by distance $\rho$ along the normalized gradient direction. We show that, for some choices of $\rho$, SAM can stall at points where this shifted (perturbed-point) gradient vanishes despite a nonzero original gradient, and therefore, they are not stationary points of the original loss. We call these points hallucinated minimizers, prove their existence under simple nonconvex landscape conditions (e.g., the presence of a local minimizer and a local maximizer), and establish sufficient conditions for local convergence of the SAM iterates to them. We corroborate this failure mode in neural network training and observe that it aligns with SAM's performance degradation often seen at large $\rho$. Finally, as a practical safeguard, we find that a short initial SGD warm-start before enabling SAM mitigates this failure mode and reduces sensitivity to the choice of $\rho$.

Deep Learning · Generative Models and Autoencoders

Xianglu Wang, Bangxian Han, Hu Ding

Diffusion models have achieved state-of-the-art performance in image generation tasks. However, training powerful diffusion models remains time-consuming, which limits their practical deployment. In this paper, we revisit the learning dynamics of diffusion models through the lens of *spectral bias*, a phenomenon in which deep neural networks prioritize learning low-frequency modes. Through an empirical analysis of diffusion training, we observe that diffusion models exhibit a **dual** spectral bias. First, over training iterations, they fit low-frequency components earlier than high-frequency details. Second, along the diffusion timesteps, early denoising steps mainly reconstruct coarse low-frequency content, while high-frequency details emerge in later steps. Motivated by this observation, we propose Frequency-aware Patch Selection **(FaPS)**, a general and fast training method for diffusion models that can be applied to both UNet and DiT backbones. Specifically, FaPS introduces a *frequency-aware gating* that adaptively selects image patches based on their frequency information and focuses computation only on the selected patches. Since the selection decisions are discrete and thus non-differentiable, we model the gating as a stochastic policy network and optimize it end-to-end using a policy gradient method. Our experiments demonstrate that FaPS achieves up to $\mathbf{3}\times$ faster training while maintaining comparable or superior generation quality, and improves the performance of diffusion models in limited-data settings.

Deep Learning · Graph Neural Networks

Yuting Zhang, Zhitao Xiao, Zhitao Xiao, Lei Geng, Yanwei Pang, Xiao Wang

Self-supervised Continual Graph Learning (CGL) aims to successively learn from a graph sequence with different tasks without label supervision—a paradigm that has attracted widespread attention. Most existing self-supervised CGL methods rely on instance-level consistency objectives that enforce stability of individual node (or node-pair) embeddings. Due to optimizing nodes in isolation, these methods fail to maintain global relational structure, causing inter-node correspondences to progressively distort under continual learning. To this end, we propose a novel Structure-Aware Optimal Transport (SAOT) framework that explicitly captures and preserves relational structure within graph representations across sequential tasks. Specifically, SAOT leverages optimal transport theory to capture global inter-node correspondences, thereby facilitating and enhancing graph representation learning. Simultaneously, SAOT incorporates a cross-task knowledge distillation mechanism to preserve the previous structural knowledge. Extensive experiments on four CGL benchmark datasets demonstrate that SAOT outperforms existing self-supervised baselines. In particular, SAOT achieves significant performance gains, improving average accuracy by up to 6% on CoraFull-CL and over 15% on Products-CL compared with state-of-the-art methods in the Class-IL setting.

Applications · Chemistry, Physics, and Earth Sciences

Chen-Yang Dai, Che-Chia Chang, Te-Sheng Lin, Ming-Chih Lai, Chieh-Hsin Lai

Physics-informed neural networks (PINNs) solve time-dependent partial differential equations (PDEs) by learning a mesh-free, differentiable solution that can be evaluated anywhere in space and time. However, standard space-time PINNs take time as an input but reuse a single network with shared weights across all times, forcing the same features to represent markedly different dynamics. This coupling degrades accuracy and can destabilize training when enforcing PDE, boundary, and initial constraints jointly. We propose *Time-Induced Neural Networks (TINNs)*, a novel architecture that parameterizes the network weights as a learned function of time, allowing the effective spatial representation to evolve over time while maintaining shared structure. The resulting formulation naturally yields a nonlinear least-squares problem, which we optimize efficiently using a Levenberg-Marquardt method. Experiments on various time-dependent PDEs show up to $4\times$ improved accuracy and $10\times$ faster convergence compared to PINNs and strong baselines.

Theo X. Olausson, Joao Monteiro, Michal Klein, Marco Cuturi

Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring identification of database vectors that align most strongly with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of search across queries drawn from a known distribution. Our key insight is that the MIPS value function - the maximum inner product as a function of the query - is convex (as the pointwise maximum of linear functions), and its gradient at each query equals the optimal database vector. We explore two complementary architectures: (1) Input Convex Neural Networks (ICNNs) that learn the convex value function and recover the optimal match via gradient computation, and (2) VectorICNNs that directly regress the argmax, bypassing gradient computation entirely at inference time. For ICNNs, we combine score regression with gradient matching losses; for VectorICNNs, we introduce a score consistency loss derived from Euler's theorem for homogeneous functions. We further propose homogenization wrappers that enforce positive 1-homogeneity, theoretically linking function values to gradients. Our experiments on retrieval benchmarks demonstrate that convexity provides an effective inductive bias, with learned potentials achieving high match rates while requiring only a single forward pass at inference.

General Machine Learning · Everything Else

Jiayi Dai, Randy Goebel

Because of the pervasive use of deep neural networks (DNNs), especially in high-stakes domains, the interpretability of DNNs has received increased attention. The general idea of rationale extraction (RE) is to provide an interpretable-by-design framework for DNNs via a select-predict architecture where two neural networks learn jointly to perform feature selection and prediction, respectively. Given only the remote supervision from the final task prediction, the process of learning to select subsets of features (or *rationales*) requires searching in the space of all possible feature combinations, which is computationally challenging and even harder when the base neural networks are not sufficiently capable. To improve the predictive performance of RE models that are based on less capable or smaller neural networks (i.e., the students), we propose **REKD** (**R**ationale **E**xtraction with **K**nowledge **D**istillation) where a student RE model learns from the rationales and predictions of a teacher (i.e., a *rationalist*) in addition to the student's own RE optimization. This structural adjustment to RE aligns well with how humans could learn effectively from interpretable and verifiable knowledge. Because of the neural-model agnostic nature of the method, any black-box neural network could be integrated as a backbone model. To demonstrate the viability of REKD, we conduct experiments with multiple variants of BERT and vision transformer (ViT) models. Our experiments across language and vision classification datasets (i.e., IMDB movie reviews, CIFAR 10 and CIFAR 100) show that REKD significantly improves the predictive performance of the student RE models.

Optimization · Large Scale, Parallel and Distributed

Rustem Islamov, Michael Crawshaw, Jeremy Cohen, Robert Gower

The Edge of Stability (EoS) is a phenomenon where the sharpness (largest eigenvalue) of the Hessian converges to $2/\eta$ during training with gradient descent (GD) with a step-size $\eta$. Despite violating classical smoothness assumptions, EoS has been widely observed in deep learning, but its theoretical foundations remain incomplete. We propose a framework for analyzing EoS of non-Euclidean GD using directional smoothness (Mishkin et al., 2024), which naturally extends to non-Euclidean norms. This approach allows us to characterize EoS beyond the standard Euclidean setting, encompassing methods such as $\ell_{\infty}$-descent, Block CD, Spectral GD, and Muon without momentum. We derive the appropriate measure of the generalized sharpness under an arbitrary norm. Our generalized sharpness measure includes previously studied vanilla GD and preconditioned GD as special cases. Through analytical results and experiments on neural networks, we show that non-Euclidean GD also exhibits progressive sharpening followed by oscillations around the threshold $2/\eta$. Practically, our framework provides a single, geometry-aware spectral measure that works across optimizers, bridging the gap between empirical observations and deep learning theory.

Optimization · Non-Convex

Rustem Islamov, Michael Crawshaw, Jeremy Cohen, Robert Gower

The Edge of Stability (EoS) is a phenomenon where the sharpness (largest eigenvalue) of the Hessian converges to $2/\eta$ during training with gradient descent (GD) with a step-size $\eta$. Despite violating classical smoothness assumptions, EoS has been widely observed in deep learning, but its theoretical foundations remain incomplete. We propose a framework for analyzing EoS of non-Euclidean GD using directional smoothness (Mishkin et al., 2024), which naturally extends to non-Euclidean norms. This approach allows us to characterize EoS beyond the standard Euclidean setting, encompassing methods such as $\ell_{\infty}$-descent, Block CD, Spectral GD, and Muon without momentum. We derive the appropriate measure of the generalized sharpness under an arbitrary norm. Our generalized sharpness measure includes previously studied vanilla GD and preconditioned GD as special cases. Through analytical results and experiments on neural networks, we show that non-Euclidean GD also exhibits progressive sharpening followed by oscillations around the threshold $2/\eta$. Practically, our framework provides a single, geometry-aware spectral measure that works across optimizers, bridging the gap between empirical observations and deep learning theory.

Deep Learning · Graph Neural Networks

Geon Lee, Sunwoo Kim, Kyungho Kim, Kijung Shin

Graph collaborative filtering (GCF) is a dominant paradigm in recommender systems, where contrastive learning (CL) objectives such as the Sampled Softmax (SSM) loss are widely used for optimization. However, it remains unclear how CL interacts with the prediction mechanism of GCF. By unfolding the prediction mechanism of GCF, we show that the user-item prediction score is computed by aggregating learnable weights over a large number of neighbor pairs formed by the multi-hop neighbors of the user and the item. This analysis implies that effective optimization critically depends on which neighbor pairs are upweighted during training. Empirically, we find that effective recommendation is achievable by selectively upweighting only a small subset of neighbor pairs whose constituent neighbors are structurally similar to the target user and item, and that the effect of such selective upweighting varies across different neighbor pair types. Based on these findings, we analyze SSM and identify key limitations in its neighbor pair weight update dynamics. To address these limitations, we propose NT-SSM, an effective and principled CL objective that induces type-aware neighbor pair weight update dynamics. Experiments demonstrate consistent performance improvements over SSM across multiple datasets and GCF models.

Deep Learning · Graph Neural Networks

Chendi Qian, Christopher Morris, Stefanie Jegelka, Christian Sohler

There has been a growing interest in using neural networks, especially message-passing neural networks (MPNNs), to solve hard combinatorial optimization problems heuristically. However, existing learning-based approaches for hard combinatorial optimization tasks often rely on supervised training data, reinforcement learning, or gradient estimators, leading to significant computational overhead, unstable training, or a lack of provable performance guarantees. In contrast, classical approximation algorithms offer such performance guarantees under worst-case inputs but are non-differentiable and unable to adaptively exploit structural regularities in natural input distributions. We address this dichotomy with the fundamental example of Uniform Facility Location (UniFL), a variant of the combinatorial facility location problem with applications in clustering, data summarization, logistics and supply chain design. We develop a fully differentiable MPNN model that embeds approximation-algorithmic principles while avoiding the need for solver supervision or discrete relaxations. Our approach admits provable approximation and size generalization guarantees to much larger instances than seen during training. Empirically, we show that our approach outperforms standard non-learned approximation algorithms in terms of solution quality, closing the gap with computationally intensive integer linear programming approaches. Overall, this work provides a step toward bridging learning-based methods and approximation algorithms for discrete optimization.

General Machine Learning · Representation Learning

Yulin Zhang, Boxuan Hu, Huimin Li, Junlin Hu

Hyperbolic metric learning is highly effective in embedding hierarchical data structures. However, past work has predominantly focused on the conformal Poincaré model, leaving other geometries like the Klein model largely under-explored. In addition, the curved geodesics of the Poincaré model present a fundamental geometric misalignment with the linear projections dominating the feature transformation steps in the modern neural network backbones. In this paper, we investigate the Klein model, a projective model of hyperbolic geometry whose straight-line geodesics offer a structurally aligned alternative in linear encoders, for hyperbolic metric learning. By formalizing a framework based on Einstein gyrovector operations, we derive a numerically stable metric learning approach that mitigates the inherent optimization challenges of the Klein model. Extensive experiments on multiple image datasets for fine-grained image classification task show that the Klein model not only serves as a viable alternative to the Poincaré model but also achieves highly competitive performance by leveraging its unique geometric properties, without increasing parameter complexity. Our empirical findings establish the Klein model as an efficient geometric prior for hyperbolic metric learning.

Semih Cayci

Ill-conditioned kernel matrices and loss landscapes can make first-order methods for training neural networks converge slowly. We establish non-asymptotic convergence bounds for the Gauss–Newton method in both under- and overparameterized regimes, showing it avoids these conditioning bottlenecks. In the underparameterized setting, Gauss–Newton gradient flow in parameter space induces a Riemannian gradient flow on a low-dimensional submanifold of function space. Using tools from Riemannian optimization, we show that, under an appropriate output scaling, the loss satisfies geodesic Polyak–Lojasiewicz and Lipschitz-smoothness conditions, implying geometric convergence to the optimal in-class predictor at an explicit rate independent of Gram-matrix conditioning. In the overparameterized setting, we identify adaptive, curvature-aware regularization schedules and prove fast geometric convergence to a global optimum for both Gauss–Newton gradient flow and discrete-time Gauss–Newton iterates, with rates independent of the minimum eigenvalue of the neural tangent kernel and, locally, independent of the strong-convexity modulus. Overall, Gauss–Newton can be provably faster in ill-conditioned regimes where first-order methods slow down.

Applications · Time Series

Yihang Lu, Xianwei Meng, Enhong Chen

Neural Forecasters (NFs) have become a cornerstone of Long-term Time Series Forecasting (LTSF). However, recent progress has been hampered by an overemphasis on architectural complexity at the expense of fundamental forecasting principles. In this work, we revisit the principles of LTSF. We begin by formulating a Variance Reduction Hypothesis (VRH), positing that generating and combining multiple forecasts is essential to reducing the inherent uncertainty of NFs. Guided by this, we propose Boosted Direct Output (BDO), a streamlined paradigm that synergistically hybridizes the causal structure of Auto-Regressive (AR) with the stability of Direct Output (DO), while implicitly realizing the principle of forecast combination within a single network. Furthermore, we address the critical validation-test generalization gap by employing parameter smoothing to stabilize optimization. Extensive experiments demonstrate that these trivial yet principled improvements enable a direct temporal MLP to outperform recent, complex state-of-the-art models in nearly all benchmarks, without relying on intricate inductive biases. Finally, we empirically verify our hypothesis, establishing a dynamic performance bound that highlights promising directions for future research. The code for review is available at: \url{https://anonymous.4open.science/r/ReNF-A151}.

Deep Learning · Theory

Viktoriia Chekalina, Daniil Moskovskiy, Tatyana Matveeva, Andrey Kuznetsov, Evgeny Frolov

The Fisher Information Matrix (FIM) provides a principled geometric framework for parameter sensitivity in neural networks, but directly computing and using the full FIM is infeasible in high-dimensional models. As a result, most existing methods rely on diagonal approximations that discard important correlation structure. We introduce Matrix-free Fisher Factorization (MFF), a GPU-tractable algorithm that captures both diagonal and off-diagonal dependencies without materializing the full matrix. For post-training neural network layer compression, we prove that under Matrix-Variate Normal assumptions, MFF yields GFWSVD, a unique closed-form linear layer decomposition that optimally minimizes the expected second-order loss increase. Experiments on controlled numerical benchmarks with large neural networks show that GFWSVD achieves up to 50\% compression while matching or exceeding state-of-the-art diagonal and activation-based baselines across most tasks, and it reliably avoids collapse in dense architectures such as Llama 3. Moreover, when used to initialize existing optimization pipelines (e.g., Dobi-SVD), GFWSVD better preserves accuracy at 40\% parameter reduction in regimes where standard methods substantially degrade. Together, these results position MFF and GFWSVD as foundational algorithmic primitives for scalable, second-order-aware neural network approximation and parameter sensitivity.

Deep Learning · Theory

Tianyu Pang, Vignesh Kothapalli, Shenyang Deng, Haohui Wang, Dawei Zhou, Yaoqing Yang

We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn a single-index target function. In particular, we derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent, enabling a precise characterization of early training dynamics. We characterize how learning rates should scale under the gradient approximation in the first two steps, and prove that performing updates with this approximation yields a tractable surrogate loss with a tight, small approximation error. This formulation enables the theoretical analysis of layer-wise learning rates and reveals a distinct early-training regime: test loss can be minimized by unequal learning rates at the initial step, while equal learning rates become optimal in subsequent steps. Our numerical experiments validate these theoretical predictions and demonstrate the importance of balancing layer-wise learning-rate during early training.

Deep Learning · Theory

Yuxin Ma, Nan Chen, Mateo Diaz, Soufiane Hayou, Dmitriy Kunisky, Soledad Villar

Modern large-scale neural networks are often trained and released in multiple sizes to accommodate diverse inference budgets. To improve efficiency, recent work has explored *model upscaling*: initializing larger models from trained smaller ones in order to transfer knowledge and accelerate convergence. However, this method can be sensitive to hyperparameters that need to be tuned at the target upscaled model size, which is prohibitively costly to do directly. It remains unclear whether the most common workaround—tuning on smaller models and extrapolating via hyperparameter scaling laws—is still sound when using upscaling. We address this with principled approaches to upscaling with respect to model widths and efficiently tuning hyperparameters in this setting. First, motivated by $\mu$P and any-dimensional architectures, we introduce a general upscaling method applicable to a broad range of architectures and optimizers, backed by theory guaranteeing that models are equivalent to their widened versions and allowing for rigorous analysis of infinite-width limits. Second, we extend the theory of $\mu$Transfer to a hyperparameter transfer technique for models upscaled using our method and empirically demonstrate that this method is effective on realistic datasets and architectures.