Goto

Collaborating Authors

 curran associate


StAD: Stein Amortized Divergence for Fast Likelihoods with Diffusion and Flow

arXiv.org Machine Learning

Diffusion and flow-based models are ubiquitously used for generative modelling and density estimation. They admit a deterministic probability flow ordinary differential equation (PF-ODE), analogous to continuous normalizing flows (CNFs), which describes the transport of the probability mass. Obtaining the likelihood from these models is of interest to many workflows, especially Bayesian analysis, and requires solving the trace of the Jacobian to compute the divergence of the learned PF-ODE, which is either $\mathcal{O}(D^2)$ to compute exactly or $\mathcal{O}(D)$ with a noisy estimate. We introduce StAD, a new distillation method to predict and learn the divergence of the PF-ODE using the Langevin-Stein operator without ever computing the Jacobian. We show that our method is competitive with the Hutchinson and Hutch++ on CIFAR-10, ImageNet and other density estimation tasks, consistently improving the variance and speed of the likelihood predictions compared to the Hutchinson. We additionally show our method will generalize to a varied class of generative models, and show that under some regularity conditions these learned vector fields can be made to satisfy the Stein class.


Uncovering Symmetry Transfer in Large Language Models via Layer-Peeled Optimization

arXiv.org Machine Learning

Large language models (LLMs) are pretrained by minimizing the cross-entropy loss for next-token prediction. In this paper, we study whether this optimization strategy can induce geometric structure in the learned model weights and context embeddings. We approach this problem by analyzing a constrained layer-peeled optimization program, which serves as a mathematically tractable surrogate for LLMs by treating the output projection matrix and last-layer context embeddings as optimization variables. Our analysis of this nonconvex optimization program demonstrates that symmetries in the target next-token distributions are transferred to the global minimizers of the layer-peeled model in a precise group-theoretic sense. Specifically, we prove that when the target tokens exhibit a cyclic-shift symmetry (such as the seven days of the week or the twelve months of the year), the optimal logit matrix is exactly circulant, and the Gram matrices of both the output projections and the context embeddings form circulant geometries as well. Next, for exchangeable target distributions invariant under the symmetric group and, more generally, under two-transitive group actions, we show that the global optimal output projection matrix forms a simplex equiangular tight frame, while the optimal logit matrix and context embeddings inherit the permutation symmetries present in the input data. A key technical step is to reduce the constrained nonconvex factorized problem to an explicit logit-level convex characterization for cyclic symmetry and to a symmetry-based lower bound for permutation symmetry, together with a sharp characterization of the optimal factorization. Finally, we empirically demonstrate that open-source LLMs naturally exhibit symmetries consistent with our theoretical predictions, despite being trained without any explicit regularization promoting such geometric structure.


Decentralized Proximal Stochastic Gradient Langevin Dynamics

arXiv.org Machine Learning

Decentralized learning is a learning process in which data is distributed across computational agents or collected by individual agents, and model parameters are computed as the consensus of the agents. It has gained a lot of interest for applications where agents can collaboratively learn a predictive model without sharing their own data, but sharing only their local models with their immediate neighbors to generate a global model [He et al., 2018, Hendrikx et al., 2019, Arjevani et al., 2020]. We assume there are N agents who are connected over an undirected communication network G = (V,E) where V = {1,...,N} represents the agents and E V V denotes the set of edges; i.e., if agent i and j are connected then (i,j) E implies (j,i) E. Suppose we have a collection of n independent and identically distributed (i.i.d.) data pairs zi = (ai,yi), where ai Rp is the feature vector and yi the label or response of the i-th observation. Let Z = [z1,z2,,zn] Rnp be sampled from the distribution p(Z|x) where the parameter x Rd has a common prior. The goal is to sample from the posterior distribution p(x|Z) p(Z|x)p(x) by distributing Z among N agents such that Zi = {zi1,zi2,,zini} is the subset of data exclusive to agent i.





Multi-Step Budgeted Bayesian Optimization with Unknown Evaluation Costs

Neural Information Processing Systems

Bayesian optimization (BO) is a sample-efficient approach to optimizing costly-toevaluate black-box functions. Most BO methods ignore how evaluation costs may vary over the optimization domain. However, these costs can be highly heterogeneous and are often unknown in advance. This occurs in many practical settings, such as hyperparameter tuning of machine learning algorithms or physics-based simulation optimization. Moreover, those few existing methods that acknowledge cost heterogeneity do not naturally accommodate a budget constraint on the total evaluation cost.


Uncertainty Estimation for Multi-view Data: The Power of Seeing the Whole Picture

Neural Information Processing Systems

Uncertainty estimation is essential to make neural networks trustworthy in realworld applications. Extensive research efforts have been made to quantify and reduce predictive uncertainty. However, most existing works are designed for unimodal data, whereas multi-view uncertainty estimation has not been sufficiently investigated. Therefore, we propose a new multi-view classification framework for better uncertainty estimation and out-of-domain sample detection, where we associate each view with an uncertainty-aware classifier and combine the predictions of all the views in a principled way.


Score-basedGenerativeNeuralNetworksfor Large-ScaleOptimalTransport

Neural Information Processing Systems

Comparison of statistical distances can also enable distribution testing, quantification of distribution shifts, and provide methods to correct for distribution shift through domainadaptation[12]. Optimal transport theory provides a rich set of tools for comparing distributions inWasserstein Distance.


One-Line-of-Code Data Mollification Improves Optimization of Likelihood-based Generative Models

Neural Information Processing Systems

Generative Models (GMs) have attracted considerable attention due to their tremendous success in various domains, such as computer vision where they are capable to generate impressive realistic-looking images. Likelihood-based GMs are attractive due to the possibility to generate new data by a single model evaluation. However, they typically achieve lower sample quality compared to state-of-the-art score-based Diffusion Models (DMs). This paper provides a significant step in the direction of addressing this limitation. The idea is to borrow one of the strengths of score-based DMs, which is the ability to perform accurate density estimation in low-density regions and to address manifold overfitting by means of data mollification. We propose a view of data mollification within likelihood-based GMs as a continuation method, whereby the optimization objective smoothly transitions from simple-to-optimize to the original target. Crucially, data mollification can be implemented by adding one line of code in the optimization loop, and we demonstrate that this provides a boost in generation quality of likelihood-based GMs, without computational overheads. We report results on real-world image data sets and UCI benchmarks with popular likelihood-based GMs, including variants of variational autoencoders and normalizing flows, showing large improvements in FID score and density estimation.