Lexing Ying

Jikai Jin, Yiping Lu, Jose Blanchet et al. · stanford

Learning mappings between infinite-dimensional function spaces has achieved empirical success in many disciplines of machine learning, including generative modeling, functional data analysis, causal inference, and multi-agent reinforcement learning. In this paper, we study the statistical limit of learning a Hilbert-Schmidt operator between two infinite-dimensional Sobolev reproducing kernel Hilbert spaces. We establish the information-theoretic lower bound in terms of the Sobolev Hilbert-Schmidt norm and show that a regularization that learns the spectral components below the bias contour and ignores the ones that are above the variance contour can achieve the optimal learning rate. At the same time, the spectral components between the bias and variance contours give us flexibility in designing computationally feasible machine learning algorithms. Based on this observation, we develop a multilevel kernel operator learning algorithm that is optimal when learning linear operators between infinite-dimensional function spaces.

Yinuo Ren, Hongli Zhao, Yuehaw Khoo et al. · stanford

We propose the tensorizing flow method for estimating high-dimensional probability density functions from the observed data. The method is based on tensor-train and flow-based generative modeling. Our method first efficiently constructs an approximate density in the tensor-train form via solving the tensor cores from a linear system based on the kernel density estimators of low-dimensional marginals. We then train a continuous-time flow model from this tensor-train density to the observed empirical distribution by performing a maximum likelihood estimation. The proposed method combines the optimization-less feature of the tensor-train with the flexibility of the flow-based generative models. Numerical results are included to demonstrate the performance of the proposed method.

Yiping Lu, Jose Blanchet, Lexing Ying · stanford

In this paper, we study the statistical limits in terms of Sobolev norms of gradient descent for solving inverse problem from randomly sampled noisy observations using a general class of objective functions. Our class of objective functions includes Sobolev training for kernel regression, Deep Ritz Methods (DRM), and Physics Informed Neural Networks (PINN) for solving elliptic partial differential equations (PDEs) as special cases. We consider a potentially infinite-dimensional parameterization of our model using a suitable Reproducing Kernel Hilbert Space and a continuous parameterization of problem hardness through the definition of kernel integral operators. We prove that gradient descent over this objective function can also achieve statistical optimality and the optimal number of passes over the data increases with sample size. Based on our theory, we explain an implicit acceleration of using a Sobolev norm as the objective function for training, inferring that the optimal number of epochs of DRM becomes larger than the number of PINN when both the data size and the hardness of tasks increase, although both DRM and PINN can achieve statistical optimality.

Emmanuel Candes, Laurent Demanet, Lexing Ying

This paper is concerned with the fast computation of Fourier integral operators of the general form $\int_{\R^d} e^{2\piıΦ(x,k)} f(k) d k$, where $k$ is a frequency variable, $Φ(x,k)$ is a phase function obeying a standard homogeneity condition, and $f$ is a given input. This is of interest for such fundamental computations are connected with the problem of finding numerical solutions to wave equations, and also frequently arise in many applications including reflection seismology, curvilinear tomography and others. In two dimensions, when the input and output are sampled on $N \times N$ Cartesian grids, a direct evaluation requires $O(N^4)$ operations, which is often times prohibitively expensive. This paper introduces a novel algorithm running in $O(N^2 \log N)$ time, i. e. with near-optimal computational complexity, and whose overall structure follows that of the butterfly algorithm [Michielssen and Boag, IEEE Trans Antennas Propagat 44 (1996), 1086-1093]. Underlying this algorithm is a mathematical insight concerning the restriction of the kernel $e^{2\piıΦ(x,k)}$ to subsets of the time and frequency domains. Whenever these subsets obey a simple geometric condition, the restricted kernel has approximately low-rank; we propose constructing such low-rank approximations using a special interpolation scheme, which prefactors the oscillatory component, interpolates the remaining nonoscillatory part and, lastly, remodulates the outcome. A byproduct of this scheme is that the whole algorithm is highly efficient in terms of memory requirement. Numerical results demonstrate the performance and illustrate the empirical properties of this algorithm.

Yiping Lu, Jiajin Li, Lexing Ying et al. · stanford

The optimal design of experiments typically involves solving an NP-hard combinatorial optimization problem. In this paper, we aim to develop a globally convergent and practically efficient optimization algorithm. Specifically, we consider a setting where the pre-treatment outcome data is available and the synthetic control estimator is invoked. The average treatment effect is estimated via the difference between the weighted average outcomes of the treated and control units, where the weights are learned from the observed data. {Under this setting, we surprisingly observed that the optimal experimental design problem could be reduced to a so-called \textit{phase synchronization} problem.} We solve this problem via a normalized variant of the generalized power method with spectral initialization. On the theoretical side, we establish the first global optimality guarantee for experiment design when pre-treatment data is sampled from certain data-generating processes. Empirically, we conduct extensive experiments to demonstrate the effectiveness of our method on both the US Bureau of Labor Statistics and the Abadie-Diemond-Hainmueller California Smoking Data. In terms of the root mean square error, our algorithm surpasses the random design by a large margin.

Björn Engquist, Lexing Ying

This paper introduces a new sweeping preconditioner for the iterative solution of the variable coefficient Helmholtz equation in two and three dimensions. The algorithms follow the general structure of constructing an approximate $LDL^t$ factorization by eliminating the unknowns layer by layer starting from an absorbing layer or boundary condition. The central idea of this paper is to approximate the Schur complement matrices of the factorization using moving perfectly matched layers (PMLs) introduced in the interior of the domain. Applying each Schur complement matrix is equivalent to solving a quasi-1D problem with a banded LU factorization in the 2D case and to solving a quasi-2D problem with a multifrontal method in the 3D case. The resulting preconditioner has linear application cost and the preconditioned iterative solver converges in a number of iterations that is essentially indefinite of the number of unknowns or the frequency. Numerical results are presented in both two and three dimensions to demonstrate the efficiency of this new preconditioner.

Björn Engquist, Lexing Ying

The paper introduces the sweeping preconditioner, which is highly efficient for iterative solutions of the variable coefficient Helmholtz equation including very high frequency problems. The first central idea of this novel approach is to construct an approximate factorization of the discretized Helmholtz equation by sweeping the domain layer by layer, starting from an absorbing layer or boundary condition. Given this specific order of factorization, the second central idea of this approach is to represent the intermediate matrices in the hierarchical matrix framework. In two dimensions, both the construction and the application of the preconditioners are of linear complexity. The GMRES solver with the resulting preconditioner converges in an amazingly small number of iterations, which is essentially independent of the number of unknowns. This approach is also extended to the three dimensional case with some success. Numerical results are provided in both two and three dimensions to demonstrate the efficiency of this new approach.

Xun Tang, Haoxuan Chen, Yuehaw Khoo et al. · stanford

We introduce Sketch Tomography, an efficient procedure for quantum state tomography based on the classical shadow protocol used for quantum observable estimations. The procedure applies to the case where the ground truth quantum state is a matrix product state (MPS). The density matrix of the ground truth state admits a tensor train ansatz as a result of the MPS assumption, and we estimate the tensor components of the ansatz through a series of observable estimations, thus outputting an approximation of the density matrix. The procedure is provably convergent with a sample complexity that scales quadratically in the system size. We conduct extensive numerical experiments to show that the procedure outputs an accurate approximation to the quantum state. For observable estimation tasks involving moderately large subsystems, we show that our procedure gives rise to a more accurate estimation than the classical shadow protocol. We also show that sketch tomography is more accurate in observable estimation than quantum states trained from the maximum likelihood estimation formulation.

Neeraj Gangwar, Rishabh Deshmukh, Michael Shavlovsky et al.

As model sizes continue to grow, parameter-efficient fine-tuning has emerged as a powerful alternative to full fine-tuning. While LoRA is widely adopted among these methods, recent research has explored vector-based adaptation methods due to their extreme parameter efficiency. However, these methods typically require substantially higher ranks than LoRA to match its performance, leading to increased training costs. This work introduces GiVA, a gradient-based initialization strategy for vector-based adaptation. It achieves training times comparable to LoRA and maintains the extreme parameter efficiency of vector-based adaptation. We evaluate GiVA across diverse benchmarks, including natural language understanding, natural language generation, and image classification. Experiments show that our approach consistently outperforms or achieves performance competitive with existing vector-based adaptation methods and LoRA while reducing rank requirements by a factor of eight ($8\times$).

Lexing Ying

We introduce a fast algorithm for computing sparse Fourier transforms supported on smooth curves or surfaces. This problem appear naturally in several important problems in wave scattering and reflection seismology. The main observation is that the interaction between a frequency region and a spatial region is approximately low rank if the product of their radii are bounded by the maximum frequency. Based on this property, equivalent sources located at Cartesian grids are used to speed up the computation of the interaction between these two regions. The overall structure of our algorithm follows the recently-introduced butterfly algorithm. The computation is further accelerated by exploiting the tensor-product property of the Fourier kernel in two and three dimensions. The proposed algorithm is accurate and has an $O(N \log N)$ complexity. Finally, we present numerical results in both two and three dimensions.

Fei Liu, Lexing Ying

We introduce a new additive sweeping preconditioner for the Helmholtz equation based on the perfect matched layer (PML). This method divides the domain of interest into thin layers and proposes a new transmission condition between the subdomains where the emphasis is on the boundary values of the intermediate waves. This approach can be viewed as an effective approximation of an additive decomposition of the solution operator. When combined with the standard GMRES solver, the iteration number is essentially independent of the frequency. Several numerical examples are tested to show the efficiency of this new approach.

Yuwei Fan, Jordi Feliu-Faba, Lin Lin et al.

In recent years, deep learning has led to impressive results in many fields. In this paper, we introduce a multi-scale artificial neural network for high-dimensional non-linear maps based on the idea of hierarchical nested bases in the fast multipole method and the $\mathcal{H}^2$-matrices. This approach allows us to efficiently approximate discretized nonlinear maps arising from partial differential equations or integral equations. It also naturally extends our recent work based on the generalization of hierarchical matrices [Fan et al. arXiv:1807.01883] but with a reduced number of parameters. In particular, the number of parameters of the neural network grows linearly with the dimension of the parameter space of the discretized PDE. We demonstrate the properties of the architecture by approximating the solution maps of non-linear Schr{ö}dinger equation, the radiative transfer equation, and the Kohn-Sham map.

Chao Ma, Daniel Kunin, Lei Wu et al.

A quadratic approximation of neural network loss landscapes has been extensively used to study the optimization process of these networks. Though, it usually holds in a very small neighborhood of the minimum, it cannot explain many phenomena observed during the optimization process. In this work, we study the structure of neural network loss functions and its implication on optimization in a region beyond the reach of a good quadratic approximation. Numerically, we observe that neural network loss functions possesses a multiscale structure, manifested in two ways: (1) in a neighborhood of minima, the loss mixes a continuum of scales and grows subquadratically, and (2) in a larger region, the loss shows several separate scales clearly. Using the subquadratic growth, we are able to explain the Edge of Stability phenomenon [5] observed for the gradient descent (GD) method. Using the separate scales, we explain the working mechanism of learning rate decay by simple examples. Finally, we study the origin of the multiscale structure and propose that the non-convexity of the models and the non-uniformity of training data is one of the causes. By constructing a two-layer neural network problem we show that training data with different magnitudes give rise to different scales of the loss function, producing subquadratic growth and multiple separate scales.

Lexing Ying

The Lippmann-Schwinger equation is an integral equation formulation for acoustic and electromagnetic scattering from an inhomogeneous media and quantum scattering from a localized potential. We present the sparsifying preconditioner for accelerating the iterative solution of the Lippmann-Schwinger equation. This new preconditioner transforms the discretized Lippmann-Schwinger equation into sparse form and leverages the efficient sparse linear algebra algorithms for computing an approximate inverse. This preconditioner is efficient and easy to implement. When combined with standard iterative methods, it results in almost frequency-independent iteration counts. We provide 2D and 3D numerical results to demonstrate the effectiveness of this new preconditioner.

Fei Liu, Lexing Ying

This paper presents an efficient preconditioner for the Lippmann-Schwinger equation that combines the ideas of the sparsifying and the sweeping preconditioners. Following first the idea of the sparsifying preconditioner, this new preconditioner starts by transforming the dense linear system of the Lippmann-Schwinger equation into a nearly sparse system. The key novelty is a newly designed perfectly matched layer (PML) stencil for the boundary degrees of freedoms. The resulting sparse system gives rise to fairly accurate solutions and hence can be viewed as an accurate discretization of the Helmholtz equation. This new PML stencil also paves the way for applying the moving PML sweeping preconditioner to invert the resulting sparse system approximately. When combined with the standard GMRES solver, this new preconditioner for the Lippmann-Schwinger equation takes only a few iterations to converge for both 2D and 3D problems, where the iteration numbers are almost independent of the frequency. To the best of our knowledge, this is the first method that achieves near-linear cost to solve the 3D Lippmann-Schwinger equation in high frequency cases.

Lexing Ying, Sergey Fomel

We introduce two efficient algorithms for computing the partial Fourier transforms in one and two dimensions. Our study is motivated by the wave extrapolation procedure in reflection seismology. In both algorithms, the main idea is to decompose the summation domain of into simpler components in a multiscale way. Existing fast algorithms are then applied to each component to obtain optimal complexity. The algorithm in 1D is exact and takes $O(N\log^2 N)$ steps. Our solution in 2D is an approximate but accurate algorithm that takes $O(N^2 \log^2 N)$ steps. In both cases, the complexities are almost linear in terms of the degree of freedom. We provide numerical results on several test examples.

Lin Lin, Ze Xu, Lexing Ying

Phonon calculations based on first principle electronic structure theory, such as the Kohn-Sham density functional theory, have wide applications in physics, chemistry and material science. The computational cost of first principle phonon calculations typically scales steeply as $\mathcal{O}(N_e^4)$, where $N_e$ is the number of electrons in the system. In this work, we develop a new method to reduce the computational complexity of computing the full dynamical matrix, and hence the phonon spectrum, to $\mathcal{O}(N_e^3)$. The key concept for achieving this is to compress the polarizability operator adaptively with respect to the perturbation of the potential due to the change of the atomic configuration. Such adaptively compressed polarizability operator (ACP) allows accurate computation of the phonon spectrum. The reduction of complexity only weakly depends on the size of the band gap, and our method is applicable to insulators as well as semiconductors with small band gaps. We demonstrate the effectiveness of our method using one-dimensional and two-dimensional model problems.

Laurent Demanet, Lexing Ying

This paper presents a numerical compression strategy for the boundary integral equation of acoustic scattering in two dimensions. These equations have oscillatory kernels that we represent in a basis of wave atoms, and compress by thresholding the small coefficients to zero. This phenomenon was perhaps first observed in 1993 by Bradie, Coifman, and Grossman, in the context of local Fourier bases \cite{BCG}. Their results have since then been extended in various ways. The purpose of this paper is to bridge a theoretical gap and prove that a well-chosen fixed expansion, the nonstandard wave atom form, provides a compression of the acoustic single and double layer potentials with wave number $k$ as $O(k)$-by-$O(k)$ matrices with $O(k^{1+1/\infty})$ nonnegligible entries, with a constant that depends on the relative $\ell_2$ accuracy $\eps$ in an acceptable way. The argument assumes smooth, separated, and not necessarily convex scatterers in two dimensions. The essential features of wave atoms that enable to write this result as a theorem is a sharp time-frequency localization that wavelet packets do not obey, and a parabolic scaling wavelength $\sim$ (essential diameter)${}^2$. Numerical experiments support the estimate and show that this wave atom representation may be of interest for applications where the same scattering problem needs to be solved for many boundary conditions, for example, the computation of radar cross sections.

Fei Liu, Lexing Ying

This paper introduces the recursive sweeping preconditioner for the numerical solution of the Helmholtz equation in 3D. This is based on the earlier work of the sweeping preconditioner with the moving perfectly matched layers (PMLs). The key idea is to apply the sweeping preconditioner recursively to the quasi-2D auxiliary problems introduced in the 3D sweeping preconditioner. Compared to the non-recursive 3D sweeping preconditioner, the setup cost of this new approach drops from $O(N^{4/3})$ to $O(N)$, the application cost per iteration drops from $O(N\log N)$ to $O(N)$, and the iteration count only increases mildly when combined with the standard GMRES solver. Several numerical examples are tested and the results are compared with the non-recursive sweeping preconditioner to demonstrate the efficiency of the new approach.

Yiping Lu, Wenlong Ji, Zachary Izzo et al. · stanford

Although overparameterized models have shown their success on many machine learning tasks, the accuracy could drop on the testing distribution that is different from the training one. This accuracy drop still limits applying machine learning in the wild. At the same time, importance weighting, a traditional technique to handle distribution shifts, has been demonstrated to have less or even no effect on overparameterized models both empirically and theoretically. In this paper, we propose importance tempering to improve the decision boundary and achieve consistently better results for overparameterized models. Theoretically, we justify that the selection of group temperature can be different under label shift and spurious correlation setting. At the same time, we also prove that properly selected temperatures can extricate the minority collapse for imbalanced classification. Empirically, we achieve state-of-the-art results on worst group classification tasks using importance tempering.

Lexing Ying

We introduce a new coarse-graining algorithm, tensor network skeletonization, for the numerical computation of tensor networks. This approach utilizes a structure-preserving skeletonization procedure to remove short-range correlations effectively at every scale. This approach is first presented in the setting of 2D statistical Ising model and is then extended to higher dimensional tensor networks and disordered systems. When applied to the Euclidean path integral formulation, this approach also gives rise to new efficient representations of the ground states for 1D and 2D quantum Ising models.

Jianfeng Lu, Lexing Ying

We propose an efficient algorithm for density fitting of Bloch waves for Hamiltonian operators with periodic potential. The algorithm is based on column selection and random Fourier projection of the orbital functions. The computational cost of the algorithm scales as $\mathcal{O}\bigl(N_{\text{grid}} N^2 + N_{\text{grid}} NK \log (NK)\bigr)$, where $N_{\text{grid}}$ is number of spatial grid points, $K$ is the number of sampling $k$-points in first Brillouin zone, and $N$ is the number of bands under consideration. We validate the algorithm by numerical examples in both two and three dimensions.

Xun Tang, Yoonhaeng Hur, Yuehaw Khoo et al.

In this paper, we present a density estimation framework based on tree tensor-network states. The proposed method consists of determining the tree topology with Chow-Liu algorithm, and obtaining a linear system of equations that defines the tensor-network components via sketching techniques. Novel choices of sketch functions are developed in order to consider graphical models that contain loops. Sample complexity guarantees are provided and further corroborated by numerical experiments.

Yinuo Ren, Tesi Xiao, Tanmay Gangwani et al. · stanford

Multi-objective optimization (MOO) aims to optimize multiple, possibly conflicting objectives with widespread applications. We introduce a novel interacting particle method for MOO inspired by molecular dynamics simulations. Our approach combines overdamped Langevin and birth-death dynamics, incorporating a "dominance potential" to steer particles toward global Pareto optimality. In contrast to previous methods, our method is able to relocate dominated particles, making it particularly adept at managing Pareto fronts of complicated geometries. Our method is also theoretically grounded as a Wasserstein-Fisher-Rao gradient flow with convergence guarantees. Extensive experiments confirm that our approach outperforms state-of-the-art methods on challenging synthetic and real-world datasets.

Lexing Ying

The boundary integral method is an efficient approach for solving time-harmonic acoustic obstacle scattering problems. The main computational task is the evaluation of an oscillatory boundary integral at each discretization point of the boundary. This paper presents a new fast algorithm for this task in two dimensions. This algorithm is built on top of directional low-rank approximations of the scattering kernel and uses oscillatory Chebyshev interpolation and local FFTs to achieve quasi-linear complexity. The algorithm is simple, fast, and kernel-independent. Numerical results are provided to demonstrate the effectiveness of the proposed algorithm.

Lexing Ying

The boundary integral method is an efficient approach for solving time-harmonic obstacle scattering problems by a bounded scatterer. This paper presents the directional preconditioner for the iterative solution of linear systems of the boundary integral method. This new preconditioner builds a data-sparse approximation of the integral operator, transforms it into a sparse linear system, and computes an approximate inverse with efficient sparse and hierarchical linear algebra algorithms. This preconditioner is efficient and results in small and almost frequency-independent iteration counts when combined with standard iterative solvers. Numerical results are provided to demonstrate the effectiveness of the new preconditioner.

Lexing Ying

This paper introduces the sparsifying preconditioner for the pseudospectral approximation of highly indefinite systems on periodic structures, which include the frequency-domain response problems of the Helmholtz equation and the Schrödinger equation as examples. This approach transforms the dense system of the pseudospectral discretization approximately into an sparse system via an equivalent integral reformulation and a specially-designed sparsifying operator. The resulting sparse system is then solved efficiently with sparse linear algebra algorithms and serves as a reasonably accurate preconditioner. When combined with standard iterative methods, this new preconditioner results in small iteration counts. Numerical results are provided for the Helmholtz equation and the Schrödinger in both 2D and 3D to demonstrate the effectiveness of this new preconditioner.

Björn Engquist, Lexing Ying

This paper introduces a directional multiscale algorithm for the two dimensional $N$-body problem of the Helmholtz kernel with applications to high frequency scattering. The algorithm follows the approach in [Engquist and Ying, SIAM Journal on Scientific Computing, 29 (4), 2007] where the three dimensional case was studied. The main observation is that, for two regions that follow a directional parabolic geometric configuration, the interaction between the points in these two regions through the Helmholtz kernel is approximately low rank. We propose an improved randomized procedure for generating the low rank representations. Based on these representations, we organize the computation of the far field interaction in a multidirectional and multiscale way to achieve maximum efficiency. The proposed algorithm is accurate and has the optimal $O(N\log N)$ complexity for problems from two dimensional scattering applications. We present numerical results for several test examples to illustrate the algorithm and its application to two dimensional high frequency scattering problems.

Xin Ye, Jianlin Xia, Lexing Ying

It has been known in potential theory that, for some kernels matrices corresponding to well-separated point sets, fast analytical low-rank approximation can be achieved via the use of proxy points. This proxy point method gives a surprisingly convenient way of explicitly writing out approximate basis matrices for a kernel matrix. However, this elegant strategy is rarely known or used in the numerical linear algebra community. It still needs clear algebraic understanding of the theoretical background. Moreover, rigorous quantifications of the approximation errors and reliable criteria for the selection of the proxy points are still missing. In this work, we use contour integration to clearly justify the idea in terms of a class of important kernels. We further provide comprehensive accuracy analysis for the analytical compression and show how to choose nearly optimal proxy points. The analytical compression is then combined with fast rank-revealing factorizations to get compact low-rank approximations and also to select certain representative points. We provide the error bounds for the resulting overall low-rank approximation. This work thus gives a fast and reliable strategy for compressing those kernel matrices. Furthermore, it provides an intuitive way of understanding the proxy point method and bridges the gap between this useful analytical strategy and practical low-rank approximations. Some numerical examples help to further illustrate the ideas.

Yuhua Zhu, Zachary Izzo, Lexing Ying

This paper revisits the bandit problem in the Bayesian setting. The Bayesian approach formulates the bandit problem as an optimization problem, and the goal is to find the optimal policy which minimizes the Bayesian regret. One of the main challenges facing the Bayesian approach is that computation of the optimal policy is often intractable, especially when the length of the problem horizon or the number of arms is large. In this paper, we first show that under a suitable rescaling, the Bayesian bandit problem converges toward a continuous Hamilton-Jacobi-Bellman (HJB) equation. The optimal policy for the limiting HJB equation can be explicitly obtained for several common bandit problems, and we give numerical methods to solve the HJB equation when an explicit solution is not available. Based on these results, we propose an approximate Bayes-optimal policy for solving Bayesian bandit problems with large horizons. Our method has the added benefit that its computational cost does not increase as the horizon increases.

Samarth Gupta, Daniel N. Hill, Lexing Ying et al.

In most applications of model-based Markov decision processes, the parameters for the unknown underlying model are often estimated from the empirical data. Due to noise, the policy learnedfrom the estimated model is often far from the optimal policy of the underlying model. When applied to the environment of the underlying model, the learned policy results in suboptimal performance, thus calling for solutions with better generalization performance. In this work we take a Bayesian perspective and regularize the objective function of the Markov decision process with prior information in order to obtain more robust policies. Two approaches are proposed, one based on $L^1$ regularization and the other on relative entropic regularization. We evaluate our proposed algorithms on synthetic simulations and on real-world search logs of a large scale online shopping store. Our results demonstrate the robustness of regularized MDP policies against the noise present in the models.

Chao Ma, Lexing Ying

In this paper, we show that structures similar to self-attention are natural to learn many sequence-to-sequence problems from the perspective of symmetry. Inspired by language processing applications, we study the orthogonal equivariance of seq2seq functions with knowledge, which are functions taking two inputs -- an input sequence and a ``knowledge'' -- and outputting another sequence. The knowledge consists of a set of vectors in the same embedding space as the input sequence, containing the information of the language used to process the input sequence. We show that orthogonal equivariance in the embedding space is natural for seq2seq functions with knowledge, and under such equivariance the function must take the form close to the self-attention. This shows that network structures similar to self-attention are the right structures to represent the target function of many seq2seq problems. The representation can be further refined if a ``finite information principle'' is considered, or a permutation equivariance holds for the elements of the input sequence.

Zhenning Cai, Yuwei Fan, Lexing Ying

We propose an entropic Fourier method for the numerical discretization of the Boltzmann collision operator. The method, which is obtained by modifying a Fourier Galerkin method to match the form of the discrete velocity method, can be viewed both as a discrete velocity method and as a Fourier method. As a discrete velocity method, it preserves the positivity of the solution and satisfies a discrete version of the H-theorem. As a Fourier method, it allows one to readily apply the FFT-based fast algorithms. A second-order convergence rate is validated by numerical experiments

Lexing Ying

This note considers the unstructured sparse recovery problems in a general form. Examples include rational approximation, spectral function estimation, Fourier inversion, Laplace inversion, and sparse deconvolution. The main challenges are the noise in the sample values and the unstructured nature of the sample locations. This note proposes the eigenmatrix, a data-driven construction with desired approximate eigenvalues and eigenvectors. The eigenmatrix offers a new way for these sparse recovery problems. Numerical results are provided to demonstrate the efficiency of the proposed method.

Hongrui Chen, Lexing Ying

Diffusion models have achieved huge empirical success in data generation tasks. Recently, some efforts have been made to adapt the framework of diffusion models to discrete state space, providing a more natural approach for modeling intrinsically discrete data, such as language and graphs. This is achieved by formulating both the forward noising process and the corresponding reversed process as Continuous Time Markov Chains (CTMCs). In this paper, we investigate the theoretical properties of the discrete diffusion model. Specifically, we introduce an algorithm leveraging the uniformization of continuous Markov chains, implementing transitions on random time points. Under reasonable assumptions on the learning of the discrete score function, we derive Total Variation distance and KL divergence guarantees for sampling from any distribution on a hypercube. Our results align with state-of-the-art achievements for diffusion models in $\mathbb{R}^d$ and further underscore the advantages of discrete diffusion models in comparison to the $\mathbb{R}^d$ setting.

Yinuo Ren, Haoxuan Chen, Yuchen Zhu et al. · gatech, stanford

Discrete diffusion models have emerged as a powerful generative modeling framework for discrete data with successful applications spanning from text generation to image synthesis. However, their deployment faces challenges due to the high dimensionality of the state space, necessitating the development of efficient inference algorithms. Current inference approaches mainly fall into two categories: exact simulation and approximate methods such as $τ$-leaping. While exact methods suffer from unpredictable inference time and redundant function evaluations, $τ$-leaping is limited by its first-order accuracy. In this work, we advance the latter category by tailoring the first extension of high-order numerical inference schemes to discrete diffusion models, enabling larger step sizes while reducing error. We rigorously analyze the proposed schemes and establish the second-order accuracy of the $θ$-trapezoidal method in KL divergence. Empirical evaluations on GPT-2 level text and ImageNet-level image generation tasks demonstrate that our method achieves superior sample quality compared to existing approaches under equivalent computational constraints.

Rajat Vadiraj Dwaraknath, Lexing Ying

Metric graphs are structures obtained by associating edges in a standard graph with segments of the real line and gluing these segments at the vertices of the graph. The resulting structure has a natural metric that allows for the study of differential operators and stochastic processes on the graph. Brownian motions in these domains have been extensively studied theoretically using their generators. However, less work has been done on practical algorithms for simulating these processes. We introduce the first algorithm for simulating Brownian motions on metric graphs through a timestep splitting Euler-Maruyama-based discretization of their corresponding stochastic differential equation. By applying this scheme to Langevin diffusions on metric graphs, we also obtain the first algorithm for sampling on metric graphs. We provide theoretical guarantees on the number of timestep splittings required for the algorithm to converge to the underlying stochastic process. We also show that the exit probabilities of the simulated particle converge to the vertex-edge jump probabilities of the underlying stochastic differential equation as the timestep goes to zero. Finally, since this method is highly parallelizable, we provide fast, memory-aware implementations of our algorithm in the form of custom CUDA kernels that are up to ~8000x faster than a GPU implementation using PyTorch on simple star metric graphs. Beyond simple star graphs, we benchmark our algorithm on a real cortical vascular network extracted from a DuMuX tissue-perfusion model for tracer transport. Our algorithm is able to run stable simulations with timesteps significantly larger than the stable limit of the finite volume method used in DuMuX while also achieving speedups of up to ~1500x.

Haoxuan Chen, Yinuo Ren, Martin Renqiang Min et al. · stanford

Diffusion models (DMs) have proven to be effective in modeling high-dimensional distributions, leading to their widespread adoption for representing complex priors in Bayesian inverse problems (BIPs). However, current DM-based posterior sampling methods proposed for solving common BIPs rely on heuristic approximations to the generative process. To exploit the generative capability of DMs and avoid the usage of such approximations, we propose an ensemble-based algorithm that performs posterior sampling without the use of heuristic approximations. Our algorithm is motivated by existing works that combine DM-based methods with the sequential Monte Carlo (SMC) method. By examining how the prior evolves through the diffusion process encoded by the pre-trained score function, we derive a modified partial differential equation (PDE) governing the evolution of the corresponding posterior distribution. This PDE includes a modified diffusion term and a reweighting term, which can be simulated via stochastic weighted particle methods. Theoretically, we prove that the error between the true posterior distribution can be bounded in terms of the training error of the pre-trained score function and the number of particles in the ensemble. Empirically, we validate our algorithm on several inverse problems in imaging to show that our method gives more accurate reconstructions compared to existing DM-based methods.

Haoxuan Chen, Lexing Ying · stanford

Sampling from a multimodal distribution is a fundamental and challenging problem in computational science and statistics. Among various approaches proposed for this task, one popular method is Annealed Importance Sampling (AIS). In this paper, we propose an ensemble-based version of AIS by combining it with population-based Monte Carlo methods to improve its efficiency. By keeping track of an ensemble instead of a single particle along some continuation path between the starting distribution and the target distribution, we take advantage of the interaction within the ensemble to encourage the exploration of undiscovered modes. Specifically, our main idea is to utilize either the snooker algorithm or the genetic algorithm used in Evolutionary Monte Carlo. We discuss how the proposed algorithm can be implemented and derive a partial differential equation governing the evolution of the ensemble under the continuous time and mean-field limit. We also test the efficiency of the proposed algorithm on various continuous and discrete distributions.

Lexing Ying

This note considers the multidimensional unstructured sparse recovery problems. Examples include Fourier inversion and sparse deconvolution. The eigenmatrix is a data-driven construction with desired approximate eigenvalues and eigenvectors proposed for the one-dimensional problems. This note extends the eigenmatrix approach to multidimensional problems. Numerical results are provided to demonstrate the performance of the proposed method.

Xun Tang, Holakou Rahmanian, Michael Shavlovsky et al.

Entropic optimal transport (OT) and the Sinkhorn algorithm have made it practical for machine learning practitioners to perform the fundamental task of calculating transport distance between statistical distributions. In this work, we focus on a general class of OT problems under a combination of equality and inequality constraints. We derive the corresponding entropy regularization formulation and introduce a Sinkhorn-type algorithm for such constrained OT problems supported by theoretical guarantees. We first bound the approximation error when solving the problem through entropic regularization, which reduces exponentially with the increase of the regularization parameter. Furthermore, we prove a sublinear first-order convergence rate of the proposed Sinkhorn-type algorithm in the dual space by characterizing the optimization procedure with a Lyapunov function. To achieve fast and higher-order convergence under weak entropy regularization, we augment the Sinkhorn-type algorithm with dynamic regularization scheduling and second-order acceleration. Overall, this work systematically combines recent theoretical and numerical advances in entropic optimal transport with the constrained case, allowing practitioners to derive approximate transport plans in complex scenarios.

Kaizhao Liu, Jose Blanchet, Lexing Ying et al.

Bootstrap is a popular methodology for simulating input uncertainty. However, it can be computationally expensive when the number of samples is large. We propose a new approach called \textbf{Orthogonal Bootstrap} that reduces the number of required Monte Carlo replications. We decomposes the target being simulated into two parts: the \textit{non-orthogonal part} which has a closed-form result known as Infinitesimal Jackknife and the \textit{orthogonal part} which is easier to be simulated. We theoretically and numerically show that Orthogonal Bootstrap significantly reduces the computational cost of Bootstrap while improving empirical accuracy and maintaining the same width of the constructed interval.

Lexing Ying

Differential privacy is a framework for protecting the identity of individual data points in the decision-making process. In this note, we propose a new form of differential privacy called tangent differential privacy. Compared with the usual differential privacy that is defined uniformly across data distributions, tangent differential privacy is tailored towards a specific data distribution of interest. It also allows for general distribution distances such as total variation distance and Wasserstein distance. In the case of risk minimization, we show that entropic regularization guarantees tangent differential privacy under rather general conditions on the risk function.

Yinuo Ren, Chao Ma, Lexing Ying

Why do neural networks trained with large learning rates for a longer time often lead to better generalization? In this paper, we delve into this question by examining the relation between training and testing loss in neural networks. Through visualization of these losses, we note that the training trajectory with a large learning rate navigates through the minima manifold of the training loss, finally nearing the neighborhood of the testing loss minimum. Motivated by these findings, we introduce a nonlinear model whose loss landscapes mirror those observed for real neural networks. Upon investigating the training process using SGD on our model, we demonstrate that an extended phase with a large learning rate steers our model towards the minimum norm solution of the training loss, which may achieve near-optimal generalization, thereby affirming the empirically observed benefits of late learning rate decay.

Xun Tang, Michael Shavlovsky, Holakou Rahmanian et al.

Computing the optimal transport distance between statistical distributions is a fundamental task in machine learning. One remarkable recent advancement is entropic regularization and the Sinkhorn algorithm, which utilizes only matrix scaling and guarantees an approximated solution with near-linear runtime. Despite the success of the Sinkhorn algorithm, its runtime may still be slow due to the potentially large number of iterations needed for convergence. To achieve possibly super-exponential convergence, we present Sinkhorn-Newton-Sparse (SNS), an extension to the Sinkhorn algorithm, by introducing early stopping for the matrix scaling steps and a second stage featuring a Newton-type subroutine. Adopting the variational viewpoint that the Sinkhorn algorithm maximizes a concave Lyapunov potential, we offer the insight that the Hessian matrix of the potential function is approximately sparse. Sparsification of the Hessian results in a fast $O(n^2)$ per-iteration complexity, the same as the Sinkhorn algorithm. In terms of total iteration count, we observe that the SNS algorithm converges orders of magnitude faster across a wide range of practical cases, including optimal transportation between empirical distributions and calculating the Wasserstein $W_1, W_2$ distance of discretized densities. The empirical performance is corroborated by a rigorous bound on the approximate sparsity of the Hessian matrix.

Yinuo Ren, Yiping Lu, Lexing Ying et al.

Inferring a diffusion equation from discretely-observed measurements is a statistical challenge of significant importance in a variety of fields, from single-molecule tracking in biophysical systems to modeling financial instruments. Assuming that the underlying dynamical process obeys a $d$-dimensional stochastic differential equation of the form $$\mathrm{d}\boldsymbol{x}_t=\boldsymbol{b}(\boldsymbol{x}_t)\mathrm{d} t+Σ(\boldsymbol{x}_t)\mathrm{d}\boldsymbol{w}_t,$$ we propose neural network-based estimators of both the drift $\boldsymbol{b}$ and the spatially-inhomogeneous diffusion tensor $D = ΣΣ^{T}$ and provide statistical convergence guarantees when $\boldsymbol{b}$ and $D$ are $s$-Hölder continuous. Notably, our bound aligns with the minimax optimal rate $N^{-\frac{2s}{2s+d}}$ for nonparametric function estimation even in the presence of correlation within observational data, which necessitates careful handling when establishing fast-rate generalization bounds. Our theoretical results are bolstered by numerical experiments demonstrating accurate inference of spatially-inhomogeneous diffusion tensors.

Jiahao Yao, Haoya Li, Marin Bukov et al.

Variational quantum algorithms stand at the forefront of simulations on near-term and future fault-tolerant quantum devices. While most variational quantum algorithms involve only continuous optimization variables, the representational power of the variational ansatz can sometimes be significantly enhanced by adding certain discrete optimization variables, as is exemplified by the generalized quantum approximate optimization algorithm (QAOA). However, the hybrid discrete-continuous optimization problem in the generalized QAOA poses a challenge to the optimization. We propose a new algorithm called MCTS-QAOA, which combines a Monte Carlo tree search method with an improved natural policy gradient solver to optimize the discrete and continuous variables in the quantum circuit, respectively. We find that MCTS-QAOA has excellent noise-resilience properties and outperforms prior algorithms in challenging instances of the generalized QAOA.

Haoya Li, Hsiang-fu Yu, Lexing Ying et al.

Entropy regularized Markov decision processes have been widely used in reinforcement learning. This paper is concerned with the primal-dual formulation of the entropy regularized problems. Standard first-order methods suffer from slow convergence due to the lack of strict convexity and concavity. To address this issue, we first introduce a new quadratically convexified primal-dual formulation. The natural gradient ascent descent of the new formulation enjoys global convergence guarantee and exponential convergence rate. We also propose a new interpolating metric that further accelerates the convergence significantly. Numerical results are provided to demonstrate the performance of the proposed methods under multiple settings.

Chao Ma, Lexing Ying

In this paper, we study the problem of finding mixed Nash equilibrium for mean-field two-player zero-sum games. Solving this problem requires optimizing over two probability distributions. We consider a quasistatic Wasserstein gradient flow dynamics in which one probability distribution follows the Wasserstein gradient flow, while the other one is always at the equilibrium. Theoretical analysis are conducted on this dynamics, showing its convergence to the mixed Nash equilibrium under mild conditions. Inspired by the continuous dynamics of probability distributions, we derive a quasistatic Langevin gradient descent method with inner-outer iterations, and test the method on different problems, including training mixture of GANs.

Zachary Izzo, James Zou, Lexing Ying

A recent line of work has focused on training machine learning (ML) models in the performative setting, i.e. when the data distribution reacts to the deployed model. The goal in this setting is to learn a model which both induces a favorable data distribution and performs well on the induced distribution, thereby minimizing the test loss. Previous work on finding an optimal model assumes that the data distribution immediately adapts to the deployed model. In practice, however, this may not be the case, as the population may take time to adapt to the model. In many applications, the data distribution depends on both the currently deployed ML model and on the "state" that the population was in before the model was deployed. In this work, we propose a new algorithm, Stateful Performative Gradient Descent (Stateful PerfGD), for minimizing the performative loss even in the presence of these effects. We provide theoretical guarantees for the convergence of Stateful PerfGD. Our experiments confirm that Stateful PerfGD substantially outperforms previous state-of-the-art methods.