Showing 1–50 of 86 results for author: Chen, R T

FlowMM: Generating Materials with Riemannian Flow Matching

Authors: Benjamin Kurt Miller, Ricky T. Q. Chen, Anuroop Sriram, Brandon M Wood

Abstract: Crystalline materials are a fundamental component in next-generation technologies, yet modeling their distribution presents unique computational challenges. Of the plausible arrangements of atoms in a periodic lattice only a vanishingly small percentage are thermodynamically stable, which is a key indicator of the materials that can be experimentally realized. Two fundamental tasks in this area ar… ▽ More Crystalline materials are a fundamental component in next-generation technologies, yet modeling their distribution presents unique computational challenges. Of the plausible arrangements of atoms in a periodic lattice only a vanishingly small percentage are thermodynamically stable, which is a key indicator of the materials that can be experimentally realized. Two fundamental tasks in this area are to (a) predict the stable crystal structure of a known composition of elements and (b) propose novel compositions along with their stable structures. We present FlowMM, a pair of generative models that achieve state-of-the-art performance on both tasks while being more efficient and more flexible than competing methods. We generalize Riemannian Flow Matching to suit the symmetries inherent to crystals: translation, rotation, permutation, and periodic boundary conditions. Our framework enables the freedom to choose the flow base distributions, drastically simplifying the problem of learning crystal structures compared with diffusion models. In addition to standard benchmarks, we validate FlowMM's generated structures with quantum chemistry calculations, demonstrating that it is about 3x more efficient, in terms of integration steps, at finding stable materials compared to previous open methods. △ Less

Submitted 7 June, 2024; originally announced June 2024.

Comments: https://github.com/facebookresearch/flowmm

Journal ref: ICML 2024

Neural Optimal Transport with Lagrangian Costs

Authors: Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky T. Q. Chen, Brandon Amos

Abstract: We investigate the optimal transport problem between probability measures when the underlying cost function is understood to satisfy a least action principle, also known as a Lagrangian cost. These generalizations are useful when connecting observations from a physical system where the transport dynamics are influenced by the geometry of the system, such as obstacles (e.g., incorporating barrier f… ▽ More We investigate the optimal transport problem between probability measures when the underlying cost function is understood to satisfy a least action principle, also known as a Lagrangian cost. These generalizations are useful when connecting observations from a physical system where the transport dynamics are influenced by the geometry of the system, such as obstacles (e.g., incorporating barrier functions in the Lagrangian), and allows practitioners to incorporate a priori knowledge of the underlying system such as non-Euclidean geometries (e.g., paths must be circular). Our contributions are of computational interest, where we demonstrate the ability to efficiently compute geodesics and amortize spline-based paths, which has not been done before, even in low dimensional problems. Unlike prior work, we also output the resulting Lagrangian optimal transport map without requiring an ODE solver. We demonstrate the effectiveness of our formulation on low-dimensional examples taken from prior work. The source code to reproduce our experiments is available at https://github.com/facebookresearch/lagrangian-ot. △ Less

Submitted 31 May, 2024; originally announced June 2024.

Comments: UAI 2024

Variational Schrödinger Diffusion Models

Authors: Wei Deng, Weijian Luo, Yixin Tan, Marin Biloš, Yu Chen, Yuriy Nevmyvaka, Ricky T. Q. Chen

Abstract: Schrödinger bridge (SB) has emerged as the go-to method for optimizing transportation plans in diffusion models. However, SB requires estimating the intractable forward score functions, inevitably resulting in the costly implicit training loss based on simulated trajectories. To improve the scalability while preserving efficient transportation plans, we leverage variational inference to linearize… ▽ More Schrödinger bridge (SB) has emerged as the go-to method for optimizing transportation plans in diffusion models. However, SB requires estimating the intractable forward score functions, inevitably resulting in the costly implicit training loss based on simulated trajectories. To improve the scalability while preserving efficient transportation plans, we leverage variational inference to linearize the forward score functions (variational scores) of SB and restore simulation-free properties in training backward scores. We propose the variational Schrödinger diffusion model (VSDM), where the forward process is a multivariate diffusion and the variational scores are adaptively optimized for efficient transport. Theoretically, we use stochastic approximation to prove the convergence of the variational scores and show the convergence of the adaptively generated samples based on the optimal variational scores. Empirically, we test the algorithm in simulated examples and observe that VSDM is efficient in generations of anisotropic shapes and yields straighter sample trajectories compared to the single-variate diffusion. We also verify the scalability of the algorithm in real-world data and achieve competitive unconditional generation performance in CIFAR10 and conditional generation in time series modeling. Notably, VSDM no longer depends on warm-up initializations and has become tuning-friendly in training large-scale experiments. △ Less

Submitted 19 June, 2024; v1 submitted 8 May, 2024; originally announced May 2024.

Comments: ICML 2024

Mid-infrared 2D nonredundant optical phased array of mirror emitters in an InGaAs/InP platform

Authors: Jason Midkiff, Po-Yu Hsiao, Patrick T. Camp, Ray T. Chen

Abstract: The extension of photonic technologies such as lidar and free-space optical communications from the traditional visible and near-infrared wavelengths to longer wavelengths can improve performance in adverse environments such as haze, fog, smoke, or strong solar background. Non-mechanical beam steerers will be a critical component of the low size, weight, and power modules needed for the portable o… ▽ More The extension of photonic technologies such as lidar and free-space optical communications from the traditional visible and near-infrared wavelengths to longer wavelengths can improve performance in adverse environments such as haze, fog, smoke, or strong solar background. Non-mechanical beam steerers will be a critical component of the low size, weight, and power modules needed for the portable or unmanned systems deployed in these environments. In this work, we demonstrate the first 2D optical phased array for non-mechanical beam steering in the mid-infrared spectral region. We combine a total internal-reflection mirror emitter with a non-redundant array of 30 elements to carry out 2D beam steering at a single wavelength of 4.6 $μ$m. The experiment yielded 600 resolvable far-field points, with 2400 points over a $28^\circ \times 28^\circ$ field of view calculated theoretically. Moreover, the device was fabricated in a passive InGaAs/InP platform, contributing another advance in the ongoing development of quantum cascade laser-based photonic integration. △ Less

Submitted 12 April, 2024; originally announced April 2024.

Comments: Main document: 16 pages, 11 figures; Supplement: 3 pages, 2 figures

Leveraging Normalizing Flows for Orbital-Free Density Functional Theory

Authors: Alexandre de Camargo, Ricky T. Q. Chen, Rodrigo A. Vargas-Hernández

Abstract: Orbital-free density functional theory (OF-DFT) for real-space systems has historically depended on Lagrange optimization techniques, primarily due to the inability of previously proposed electron density ansatze to ensure the normalization constraint. This study illustrates how leveraging contemporary generative models, notably normalizing flows (NFs), can surmount this challenge. We pioneer a La… ▽ More Orbital-free density functional theory (OF-DFT) for real-space systems has historically depended on Lagrange optimization techniques, primarily due to the inability of previously proposed electron density ansatze to ensure the normalization constraint. This study illustrates how leveraging contemporary generative models, notably normalizing flows (NFs), can surmount this challenge. We pioneer a Lagrangian-free optimization framework by employing these machine learning models as ansatze for the electron density. This novel approach also integrates cutting-edge variational inference techniques and equivariant deep learning models, offering an innovative alternative to the OF-DFT problem. We demonstrate the versatility of our framework by simulating a one-dimensional diatomic system, LiH, and comprehensive simulations of H$_2$, LiH, and H$_2$O molecules. The inherent flexibility of NFs facilitates initialization with promolecular densities, markedly enhancing the efficiency of the optimization process. △ Less

Submitted 18 April, 2024; v1 submitted 12 April, 2024; originally announced April 2024.

Comments: 6 pages, 4 Figures, (SI: 15 pages, 4 figures)

Photonic-Electronic Integrated Circuits for High-Performance Computing and AI Accelerators

Authors: Shupeng Ning, Hanqing Zhu, Chenghao Feng, Jiaqi Gu, Zhixing Jiang, Zhoufeng Ying, Jason Midkiff, Sourabh Jain, May H. Hlaing, David Z. Pan, Ray T. Chen

Abstract: In recent decades, the demand for computational power has surged, particularly with the rapid expansion of artificial intelligence (AI). As we navigate the post-Moore's law era, the limitations of traditional electrical digital computing, including process bottlenecks and power consumption issues, are propelling the search for alternative computing paradigms. Among various emerging technologies, i… ▽ More In recent decades, the demand for computational power has surged, particularly with the rapid expansion of artificial intelligence (AI). As we navigate the post-Moore's law era, the limitations of traditional electrical digital computing, including process bottlenecks and power consumption issues, are propelling the search for alternative computing paradigms. Among various emerging technologies, integrated photonics stands out as a promising solution for next-generation high-performance computing, thanks to the inherent advantages of light, such as low latency, high bandwidth, and unique multiplexing techniques. Furthermore, the progress in photonic integrated circuits (PICs), which are equipped with abundant photoelectronic components, positions photonic-electronic integrated circuits as a viable solution for high-performance computing and hardware AI accelerators. In this review, we survey recent advancements in both PIC-based digital and analog computing for AI, exploring the principal benefits and obstacles of implementation. Additionally, we propose a comprehensive analysis of photonic AI from the perspectives of hardware implementation, accelerator architecture, and software-hardware co-design. In the end, acknowledging the existing challenges, we underscore potential strategies for overcoming these issues and offer insights into the future drivers for optical computing. △ Less

Submitted 11 July, 2024; v1 submitted 21 March, 2024; originally announced March 2024.

Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models

Authors: Neta Shaul, Uriel Singer, Ricky T. Q. Chen, Matthew Le, Ali Thabet, Albert Pumarola, Yaron Lipman

Abstract: This paper introduces Bespoke Non-Stationary (BNS) Solvers, a solver distillation approach to improve sample efficiency of Diffusion and Flow models. BNS solvers are based on a family of non-stationary solvers that provably subsumes existing numerical ODE solvers and consequently demonstrate considerable improvement in sample approximation (PSNR) over these baselines. Compared to model distillatio… ▽ More This paper introduces Bespoke Non-Stationary (BNS) Solvers, a solver distillation approach to improve sample efficiency of Diffusion and Flow models. BNS solvers are based on a family of non-stationary solvers that provably subsumes existing numerical ODE solvers and consequently demonstrate considerable improvement in sample approximation (PSNR) over these baselines. Compared to model distillation, BNS solvers benefit from a tiny parameter space ($<$200 parameters), fast optimization (two orders of magnitude faster), maintain diversity of samples, and in contrast to previous solver distillation approaches nearly close the gap from standard distillation methods such as Progressive Distillation in the low-medium NFE regime. For example, BNS solver achieves 45 PSNR / 1.76 FID using 16 NFE in class-conditional ImageNet-64. We experimented with BNS solvers for conditional image generation, text-to-image generation, and text-2-audio generation showing significant improvement in sample approximation (PSNR) in all. △ Less

Submitted 2 March, 2024; originally announced March 2024.

Physics-informed Meta-instrument for eXperiments (PiMiX) with applications to fusion energy

Authors: Zhehui Wang, Shanny Lin, Miles Teng-Levy, Pinghan Chu, Bradley T. Wolfe, Chun-Shang Wong, Christopher S. Campbell, Xin Yue, Liyuan Zhang, Derek Aberle, Mariana Alvarado Alvarez, David Broughton, Ray T. Chen, Baolian Cheng, Feng Chu, Eric R. Fossum, Mark A. Foster, Chengkun Huang, Velat Kilic, Karl Krushelnick, Wenting Li, Eric Loomis, Thomas Schmidt Jr., Sky K. Sjue, Chris Tomkins , et al. (2 additional authors not shown)

Abstract: Data-driven methods (DDMs), such as deep neural networks, offer a generic approach to integrated data analysis (IDA), integrated diagnostic-to-control (IDC) workflows through data fusion (DF), which includes multi-instrument data fusion (MIDF), multi-experiment data fusion (MXDF), and simulation-experiment data fusion (SXDF). These features make DDMs attractive to nuclear fusion energy and power p… ▽ More Data-driven methods (DDMs), such as deep neural networks, offer a generic approach to integrated data analysis (IDA), integrated diagnostic-to-control (IDC) workflows through data fusion (DF), which includes multi-instrument data fusion (MIDF), multi-experiment data fusion (MXDF), and simulation-experiment data fusion (SXDF). These features make DDMs attractive to nuclear fusion energy and power plant applications, leveraging accelerated workflows through machine learning and artificial intelligence. Here we describe Physics-informed Meta-instrument for eXperiments (PiMiX) that integrates X-ray (including high-energy photons such as $γ$-rays from nuclear fusion), neutron and others (such as proton radiography) measurements for nuclear fusion. PiMiX solves multi-domain high-dimensional optimization problems and integrates multi-modal measurements with multiphysics modeling through neural networks. Super-resolution for neutron detection and energy resolved X-ray detection have been demonstrated. Multi-modal measurements through MIDF can extract more information than individual or uni-modal measurements alone. Further optimization schemes through DF are possible towards empirical fusion scaling laws discovery and new fusion reactor designs. △ Less

Submitted 16 January, 2024; originally announced January 2024.

Comments: 24 pages, 46 references, and 13 Figures. Manuscript extended a recent presentation in the 29th IAEA Fusion Energy Conference (FEC), London, UK, Oct. 16 - 21, 2023

Report number: Los Alamos National Laboratory report number LA-UR-24-20196

Reflected Schrödinger Bridge for Constrained Generative Modeling

Authors: Wei Deng, Yu Chen, Nicole Tianjiao Yang, Hengrong Du, Qi Feng, Ricky T. Q. Chen

Abstract: Diffusion models have become the go-to method for large-scale generative models in real-world applications. These applications often involve data distributions confined within bounded domains, typically requiring ad-hoc thresholding techniques for boundary enforcement. Reflected diffusion models (Lou23) aim to enhance generalizability by generating the data distribution through a backward process… ▽ More Diffusion models have become the go-to method for large-scale generative models in real-world applications. These applications often involve data distributions confined within bounded domains, typically requiring ad-hoc thresholding techniques for boundary enforcement. Reflected diffusion models (Lou23) aim to enhance generalizability by generating the data distribution through a backward process governed by reflected Brownian motion. However, reflected diffusion models may not easily adapt to diverse domains without the derivation of proper diffeomorphic mappings and do not guarantee optimal transport properties. To overcome these limitations, we introduce the Reflected Schrodinger Bridge algorithm: an entropy-regularized optimal transport approach tailored for generating data within diverse bounded domains. We derive elegant reflected forward-backward stochastic differential equations with Neumann and Robin boundary conditions, extend divergence-based likelihood training to bounded domains, and explore natural connections to entropic optimal transport for the study of approximate linear convergence - a valuable insight for practical training. Our algorithm yields robust generative modeling in diverse domains, and its scalability is demonstrated in real-world constrained generative modeling through standard image benchmarks. △ Less

Submitted 6 January, 2024; originally announced January 2024.

TaskMet: Task-Driven Metric Learning for Model Learning

Authors: Dishank Bansal, Ricky T. Q. Chen, Mustafa Mukadam, Brandon Amos

Abstract: Deep learning models are often deployed in downstream tasks that the training procedure may not be aware of. For example, models solely trained to achieve accurate predictions may struggle to perform well on downstream tasks because seemingly small prediction errors may incur drastic task errors. The standard end-to-end learning approach is to make the task loss differentiable or to introduce a di… ▽ More Deep learning models are often deployed in downstream tasks that the training procedure may not be aware of. For example, models solely trained to achieve accurate predictions may struggle to perform well on downstream tasks because seemingly small prediction errors may incur drastic task errors. The standard end-to-end learning approach is to make the task loss differentiable or to introduce a differentiable surrogate that the model can be trained on. In these settings, the task loss needs to be carefully balanced with the prediction loss because they may have conflicting objectives. We propose take the task loss signal one level deeper than the parameters of the model and use it to learn the parameters of the loss function the model is trained on, which can be done by learning a metric in the prediction space. This approach does not alter the optimal prediction model itself, but rather changes the model learning to emphasize the information important for the downstream task. This enables us to achieve the best of both worlds: a prediction model trained in the original prediction space while also being valuable for the desired downstream task. We validate our approach through experiments conducted in two main settings: 1) decision-focused model learning scenarios involving portfolio optimization and budget allocation, and 2) reinforcement learning in noisy environments with distracting states. The source code to reproduce our experiments is available at https://github.com/facebookresearch/taskmet △ Less

Submitted 8 December, 2023; originally announced December 2023.

Comments: NeurIPS 2023

Stochastic Optimal Control Matching

Authors: Carles Domingo-Enrich, Jiequn Han, Brandon Amos, Joan Bruna, Ricky T. Q. Chen

Abstract: Stochastic optimal control, which has the goal of driving the behavior of noisy systems, is broadly applicable in science, engineering and artificial intelligence. Our work introduces Stochastic Optimal Control Matching (SOCM), a novel Iterative Diffusion Optimization (IDO) technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffu… ▽ More Stochastic optimal control, which has the goal of driving the behavior of noisy systems, is broadly applicable in science, engineering and artificial intelligence. Our work introduces Stochastic Optimal Control Matching (SOCM), a novel Iterative Diffusion Optimization (IDO) technique for stochastic optimal control that stems from the same philosophy as the conditional score matching loss for diffusion models. That is, the control is learned via a least squares problem by trying to fit a matching vector field. The training loss, which is closely connected to the cross-entropy loss, is optimized with respect to both the control function and a family of reparameterization matrices which appear in the matching vector field. The optimization with respect to the reparameterization matrices aims at minimizing the variance of the matching vector field. Experimentally, our algorithm achieves lower error than all the existing IDO techniques for stochastic optimal control for three out of four control problems, in some cases by an order of magnitude. The key idea underlying SOCM is the path-wise reparameterization trick, a novel technique that may be of independent interest. Code at https://github.com/facebookresearch/SOC-matching △ Less

Submitted 28 June, 2024; v1 submitted 4 December, 2023; originally announced December 2023.

Orbital-Free Density Functional Theory with Continuous Normalizing Flows

Authors: Alexandre de Camargo, Ricky T. Q. Chen, Rodrigo A. Vargas-Hernández

Abstract: Orbital-free density functional theory (OF-DFT) provides an alternative approach for calculating the molecular electronic energy, relying solely on the electron density. In OF-DFT, both the ground-state density is optimized variationally to minimize the total energy functional while satisfying the normalization constraint. In this work, we introduce a novel approach by parameterizing the electroni… ▽ More Orbital-free density functional theory (OF-DFT) provides an alternative approach for calculating the molecular electronic energy, relying solely on the electron density. In OF-DFT, both the ground-state density is optimized variationally to minimize the total energy functional while satisfying the normalization constraint. In this work, we introduce a novel approach by parameterizing the electronic density with a normalizing flow ansatz, which is also optimized by minimizing the total energy functional. Our model successfully replicates the electronic density for a diverse range of chemical systems, including a one-dimensional diatomic molecule, specifically Lithium hydride with varying interatomic distances, as well as comprehensive simulations of hydrogen and water molecules, all conducted in Cartesian space. △ Less

Submitted 22 November, 2023; originally announced November 2023.

Comments: 6 pages, 3 figures

Guided Flows for Generative Modeling and Decision Making

Authors: Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, Ricky T. Q. Chen

Abstract: Classifier-free guidance is a key component for enhancing the performance of conditional generative models across diverse tasks. While it has previously demonstrated remarkable improvements for the sample quality, it has only been exclusively employed for diffusion models. In this paper, we integrate classifier-free guidance into Flow Matching (FM) models, an alternative simulation-free approach t… ▽ More Classifier-free guidance is a key component for enhancing the performance of conditional generative models across diverse tasks. While it has previously demonstrated remarkable improvements for the sample quality, it has only been exclusively employed for diffusion models. In this paper, we integrate classifier-free guidance into Flow Matching (FM) models, an alternative simulation-free approach that trains Continuous Normalizing Flows (CNFs) based on regressing vector fields. We explore the usage of \emph{Guided Flows} for a variety of downstream applications. We show that Guided Flows significantly improves the sample quality in conditional image generation and zero-shot text-to-speech synthesis, boasting state-of-the-art performance. Notably, we are the first to apply flow models for plan generation in the offline reinforcement learning setting, showcasing a 10x speedup in computation compared to diffusion models while maintaining comparable performance. △ Less

Submitted 7 December, 2023; v1 submitted 22 November, 2023; originally announced November 2023.

Neural Network Methods for Radiation Detectors and Imaging

Authors: S. Lin, S. Ning, H. Zhu, T. Zhou, C. L. Morris, S. Clayton, M. Cherukara, R. T. Chen, Z. Wang

Abstract: Recent advances in image data processing through machine learning and especially deep neural networks (DNNs) allow for new optimization and performance-enhancement schemes for radiation detectors and imaging hardware through data-endowed artificial intelligence. We give an overview of data generation at photon sources, deep learning-based methods for image processing tasks, and hardware solutions… ▽ More Recent advances in image data processing through machine learning and especially deep neural networks (DNNs) allow for new optimization and performance-enhancement schemes for radiation detectors and imaging hardware through data-endowed artificial intelligence. We give an overview of data generation at photon sources, deep learning-based methods for image processing tasks, and hardware solutions for deep learning acceleration. Most existing deep learning approaches are trained offline, typically using large amounts of computational resources. However, once trained, DNNs can achieve fast inference speeds and can be deployed to edge devices. A new trend is edge computing with less energy consumption (hundreds of watts or less) and real-time analysis potential. While popularly used for edge computing, electronic-based hardware accelerators ranging from general purpose processors such as central processing units (CPUs) to application-specific integrated circuits (ASICs) are constantly reaching performance limits in latency, energy consumption, and other physical constraints. These limits give rise to next-generation analog neuromorhpic hardware platforms, such as optical neural networks (ONNs), for high parallel, low latency, and low energy computing to boost deep learning acceleration. △ Less

Submitted 9 November, 2023; originally announced November 2023.

Report number: LA-UR-23-32395

Bespoke Solvers for Generative Flow Models

Authors: Neta Shaul, Juan Perez, Ricky T. Q. Chen, Ali Thabet, Albert Pumarola, Yaron Lipman

Abstract: Diffusion or flow-based models are powerful generative paradigms that are notoriously hard to sample as samples are defined as solutions to high-dimensional Ordinary or Stochastic Differential Equations (ODEs/SDEs) which require a large Number of Function Evaluations (NFE) to approximate well. Existing methods to alleviate the costly sampling process include model distillation and designing dedica… ▽ More Diffusion or flow-based models are powerful generative paradigms that are notoriously hard to sample as samples are defined as solutions to high-dimensional Ordinary or Stochastic Differential Equations (ODEs/SDEs) which require a large Number of Function Evaluations (NFE) to approximate well. Existing methods to alleviate the costly sampling process include model distillation and designing dedicated ODE solvers. However, distillation is costly to train and sometimes can deteriorate quality, while dedicated solvers still require relatively large NFE to produce high quality samples. In this paper we introduce "Bespoke solvers", a novel framework for constructing custom ODE solvers tailored to the ODE of a given pre-trained flow model. Our approach optimizes an order consistent and parameter-efficient solver (e.g., with 80 learnable parameters), is trained for roughly 1% of the GPU time required for training the pre-trained model, and significantly improves approximation and generation quality compared to dedicated solvers. For example, a Bespoke solver for a CIFAR10 model produces samples with Fréchet Inception Distance (FID) of 2.73 with 10 NFE, and gets to 1% of the Ground Truth (GT) FID (2.59) for this model with only 20 NFE. On the more challenging ImageNet-64$\times$64, Bespoke samples at 2.2 FID with 10 NFE, and gets within 2% of GT FID (1.71) with 20 NFE. △ Less

Submitted 29 October, 2023; originally announced October 2023.

Training-free Linear Image Inverses via Flows

Authors: Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, Brian Karrer

Abstract: Solving inverse problems without any training involves using a pretrained generative model and making appropriate modifications to the generation process to avoid finetuning of the generative model. While recent methods have explored the use of diffusion models, they still require the manual tuning of many hyperparameters for different inverse problems. In this work, we propose a training-free met… ▽ More Solving inverse problems without any training involves using a pretrained generative model and making appropriate modifications to the generation process to avoid finetuning of the generative model. While recent methods have explored the use of diffusion models, they still require the manual tuning of many hyperparameters for different inverse problems. In this work, we propose a training-free method for solving linear inverse problems by using pretrained flow models, leveraging the simplicity and efficiency of Flow Matching models, using theoretically-justified weighting schemes, and thereby significantly reducing the amount of manual tuning. In particular, we draw inspiration from two main sources: adopting prior gradient correction methods to the flow regime, and a solver scheme based on conditional Optimal Transport paths. As pretrained diffusion models are widely accessible, we also show how to practically adapt diffusion models for our method. Empirically, our approach requires no problem-specific tuning across an extensive suite of noisy linear inverse problems on high-dimensional datasets, ImageNet-64/128 and AFHQ-256, and we observe that our flow-based method for solving inverse problems improves upon closely-related diffusion-based methods in most settings. △ Less

Submitted 10 March, 2024; v1 submitted 25 September, 2023; originally announced October 2023.

Comments: 40 pages, 30 figures. Added additional qualitative results in the appendix

Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization

Authors: Dinghuai Zhang, Ricky T. Q. Chen, Cheng-Hao Liu, Aaron Courville, Yoshua Bengio

Abstract: We tackle the problem of sampling from intractable high-dimensional density functions, a fundamental task that often appears in machine learning and statistics. We extend recent sampling-based approaches that leverage controlled stochastic processes to model approximate samples from these target densities. The main drawback of these approaches is that the training objective requires full trajector… ▽ More We tackle the problem of sampling from intractable high-dimensional density functions, a fundamental task that often appears in machine learning and statistics. We extend recent sampling-based approaches that leverage controlled stochastic processes to model approximate samples from these target densities. The main drawback of these approaches is that the training objective requires full trajectories to compute, resulting in sluggish credit assignment issues due to use of entire trajectories and a learning signal present only at the terminal time. In this work, we present Diffusion Generative Flow Samplers (DGFS), a sampling-based framework where the learning process can be tractably broken down into short partial trajectory segments, via parameterizing an additional "flow function". Our method takes inspiration from the theory developed for generative flow networks (GFlowNets), allowing us to make use of intermediate learning signals. Through various challenging experiments, we demonstrate that DGFS achieves more accurate estimates of the normalization constant than closely-related prior methods. △ Less

Submitted 9 March, 2024; v1 submitted 4 October, 2023; originally announced October 2023.

Comments: Accepted by ICLR 2024

Generalized Schrödinger Bridge Matching

Authors: Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos A. Theodorou, Ricky T. Q. Chen

Abstract: Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider a generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Generaliz… ▽ More Modern distribution matching algorithms for training diffusion or flow models directly prescribe the time evolution of the marginal distributions between two boundary distributions. In this work, we consider a generalized distribution matching setup, where these marginals are only implicitly described as a solution to some task-specific objective function. The problem setup, known as the Generalized Schrödinger Bridge (GSB), appears prevalently in many scientific areas both within and without machine learning. We propose Generalized Schrödinger Bridge Matching (GSBM), a new matching algorithm inspired by recent advances, generalizing them beyond kinetic energy minimization and to account for task-specific state costs. We show that such a generalization can be cast as solving conditional stochastic optimal control, for which efficient variational approximations can be used, and further debiased with the aid of path integral theory. Compared to prior methods for solving GSB problems, our GSBM algorithm better preserves a feasible transport map between the boundary distributions throughout training, thereby enabling stable convergence and significantly improved scalability. We empirically validate our claims on an extensive suite of experimental setups, including crowd navigation, opinion depolarization, LiDAR manifolds, and image domain transfer. Our work brings new algorithmic opportunities for training diffusion models enhanced with task-specific optimality structures. Code available at https://github.com/facebookresearch/generalized-schrodinger-bridge-matching △ Less

Submitted 18 April, 2024; v1 submitted 3 October, 2023; originally announced October 2023.

Comments: ICLR 2024 Camera Ready

On Kinetic Optimal Probability Paths for Generative Models

Authors: Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matt Le, Yaron Lipman

Abstract: Recent successful generative models are trained by fitting a neural network to an a-priori defined tractable probability density path taking noise to training examples. In this paper we investigate the space of Gaussian probability paths, which includes diffusion paths as an instance, and look for an optimal member in some useful sense. In particular, minimizing the Kinetic Energy (KE) of a path i… ▽ More Recent successful generative models are trained by fitting a neural network to an a-priori defined tractable probability density path taking noise to training examples. In this paper we investigate the space of Gaussian probability paths, which includes diffusion paths as an instance, and look for an optimal member in some useful sense. In particular, minimizing the Kinetic Energy (KE) of a path is known to make particles' trajectories simple, hence easier to sample, and empirically improve performance in terms of likelihood of unseen data and sample generation quality. We investigate Kinetic Optimal (KO) Gaussian paths and offer the following observations: (i) We show the KE takes a simplified form on the space of Gaussian paths, where the data is incorporated only through a single, one dimensional scalar function, called the \emph{data separation function}. (ii) We characterize the KO solutions with a one dimensional ODE. (iii) We approximate data-dependent KO paths by approximating the data separation function and minimizing the KE. (iv) We prove that the data separation function converges to $1$ in the general case of arbitrary normalized dataset consisting of $n$ samples in $d$ dimension as $n/\sqrt{d}\rightarrow 0$. A consequence of this result is that the Conditional Optimal Transport (Cond-OT) path becomes \emph{kinetic optimal} as $n/\sqrt{d}\rightarrow 0$. We further support this theory with empirical experiments on ImageNet. △ Less

Submitted 11 June, 2023; originally announced June 2023.

Integrated multi-operand optical neurons for scalable and hardware-efficient deep learning

Authors: Chenghao Feng, Jiaqi Gu, Hanqing Zhu, Rongxing Tang, Shupeng Ning, May Hlaing, Jason Midkiff, Sourabh Jain, David Z. Pan, Ray T. Chen

Abstract: The optical neural network (ONN) is a promising hardware platform for next-generation neuromorphic computing due to its high parallelism, low latency, and low energy consumption. However, previous integrated photonic tensor cores (PTCs) consume numerous single-operand optical modulators for signal and weight encoding, leading to large area costs and high propagation loss to implement large tensor… ▽ More The optical neural network (ONN) is a promising hardware platform for next-generation neuromorphic computing due to its high parallelism, low latency, and low energy consumption. However, previous integrated photonic tensor cores (PTCs) consume numerous single-operand optical modulators for signal and weight encoding, leading to large area costs and high propagation loss to implement large tensor operations. This work proposes a scalable and efficient optical dot-product engine based on customized multi-operand photonic devices, namely multi-operand optical neurons (MOON). We experimentally demonstrate the utility of a MOON using a multi-operand-Mach-Zehnder-interferometer (MOMZI) in image recognition tasks. Specifically, our MOMZI-based ONN achieves a measured accuracy of 85.89% in the street view house number (SVHN) recognition dataset with 4-bit voltage control precision. Furthermore, our performance analysis reveals that a 128x128 MOMZI-based PTCs outperform their counterparts based on single-operand MZIs by one to two order-of-magnitudes in propagation loss, optical delay, and total device footprint, with comparable matrix expressivity. △ Less

Submitted 31 May, 2023; originally announced May 2023.

Comments: 19 pages, 10 figures

Lightening-Transformer: A Dynamically-operated Optically-interconnected Photonic Transformer Accelerator

Authors: Hanqing Zhu, Jiaqi Gu, Hanrui Wang, Zixuan Jiang, Zhekai Zhang, Rongxing Tang, Chenghao Feng, Song Han, Ray T. Chen, David Z. Pan

Abstract: The wide adoption and significant computing resource of attention-based transformers, e.g., Vision Transformers and large language models (LLM), have driven the demand for efficient hardware accelerators. There is a growing interest in exploring photonics as an alternative technology to digital electronics due to its high energy efficiency and ultra-fast processing speed. Photonic accelerators hav… ▽ More The wide adoption and significant computing resource of attention-based transformers, e.g., Vision Transformers and large language models (LLM), have driven the demand for efficient hardware accelerators. There is a growing interest in exploring photonics as an alternative technology to digital electronics due to its high energy efficiency and ultra-fast processing speed. Photonic accelerators have shown promising results for CNNs, which mainly rely on weight-static linear operations. However, they encounter issues when efficiently supporting Transformer architectures, questioning the applicability of photonics to advanced ML tasks. The primary hurdle lies in their inefficiency in handling unique workloads in Transformers, i.e., dynamic and full-range tensor multiplication. In this work, we propose Lightening-Transformer, the first light-empowered, high-performance, and energy-efficient photonic Transformer accelerator. To overcome prior designs' fundamental limitations, we introduce a novel dynamically-operated photonic tensor core, DPTC, a crossbar array of interference-based optical vector dot-product engines supporting highly parallel, dynamic, and full-range matrix multiplication. Furthermore, we design a dedicated accelerator that integrates our novel photonic computing cores with photonic interconnects for inter-core data broadcast, fully unleashing the power of optics. Comprehensive evaluations show that ours achieves >2.6x energy and >12x latency reductions compared to prior photonic accelerators and delivers the lowest energy cost and 2 to 3 orders of magnitude lower energy-delay product compared to electronic Transformer accelerators, all while maintaining digital-comparable accuracy. Our work highlights the immense potential of photonics for advanced ML workloads, such as Transformer-backboned LLM. Our work is available at https://github.com/zhuhanqing/Lightening-Transformer. △ Less

Submitted 31 December, 2023; v1 submitted 30 May, 2023; originally announced May 2023.

Comments: Published as a conference paper in HPCA 2024. Recieved the Reproducibility Badges at IEEE. Our implementation is available at https://github.com/zhuhanqing/Lightening-Transformer

M3ICRO: Machine Learning-Enabled Compact Photonic Tensor Core based on PRogrammable Multi-Operand Multimode Interference

Authors: Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Zixuan Jiang, Ray T. Chen, David Z. Pan

Abstract: Photonic computing shows promise for transformative advancements in machine learning (ML) acceleration, offering ultra-fast speed, massive parallelism, and high energy efficiency. However, current photonic tensor core (PTC) designs based on standard optical components hinder scalability and compute density due to their large spatial footprint. To address this, we propose an ultra-compact PTC using… ▽ More Photonic computing shows promise for transformative advancements in machine learning (ML) acceleration, offering ultra-fast speed, massive parallelism, and high energy efficiency. However, current photonic tensor core (PTC) designs based on standard optical components hinder scalability and compute density due to their large spatial footprint. To address this, we propose an ultra-compact PTC using customized programmable multi-operand multimode interference (MOMMI) devices, named M3ICRO. The programmable MOMMI leverages the intrinsic light propagation principle, providing a single-device programmable matrix unit beyond the conventional computing paradigm of one multiply-accumulate (MAC) operation per device. To overcome the optimization difficulty of customized devices that often requires time-consuming simulation, we apply ML for optics to predict the device behavior and enable a differentiable optimization flow. We thoroughly investigate the reconfigurability and matrix expressivity of our customized PTC, and introduce a novel block unfolding method to fully exploit the computing capabilities of a complex-valued PTC for near-universal real-valued linear transformations. Extensive evaluations demonstrate that M3ICRO achieves a 3.4-9.6x smaller footprint, 1.6-4.4x higher speed, 10.6-42x higher compute density, 3.7-12x higher system throughput, and superior noise robustness compared to state-of-the-art coherent PTC designs, while maintaining close-to-digital task accuracy across various ML benchmarks. Our code is open-sourced at https://github.com/JeremieMelo/M3ICRO-MOMMI. △ Less

Submitted 28 December, 2023; v1 submitted 30 May, 2023; originally announced May 2023.

Comments: 12 pages. Accepted to APL Machine Learning 2023

Multisample Flow Matching: Straightening Flows with Minibatch Couplings

Authors: Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, Ricky T. Q. Chen

Abstract: Simulation-free methods for training continuous-time generative models construct probability paths that go between noise distributions and individual data samples. Recent works, such as Flow Matching, derived paths that are optimal for each data sample. However, these algorithms rely on independent data and noise samples, and do not exploit underlying structure in the data distribution for constru… ▽ More Simulation-free methods for training continuous-time generative models construct probability paths that go between noise distributions and individual data samples. Recent works, such as Flow Matching, derived paths that are optimal for each data sample. However, these algorithms rely on independent data and noise samples, and do not exploit underlying structure in the data distribution for constructing probability paths. We propose Multisample Flow Matching, a more general framework that uses non-trivial couplings between data and noise samples while satisfying the correct marginal constraints. At very small overhead costs, this generalization allows us to (i) reduce gradient variance during training, (ii) obtain straighter flows for the learned vector field, which allows us to generate high-quality samples using fewer function evaluations, and (iii) obtain transport maps with lower cost in high dimensions, which has applications beyond generative modeling. Importantly, we do so in a completely simulation-free manner with a simple minimization objective. We show that our proposed methods improve sample consistency on downsampled ImageNet data sets, and lead to better low-cost sample generation. △ Less

Submitted 24 May, 2023; v1 submitted 28 April, 2023; originally announced April 2023.

Distributional GFlowNets with Quantile Flows

Authors: Dinghuai Zhang, Ling Pan, Ricky T. Q. Chen, Aaron Courville, Yoshua Bengio

Abstract: Generative Flow Networks (GFlowNets) are a new family of probabilistic samplers where an agent learns a stochastic policy for generating complex combinatorial structure through a series of decision-making steps. Despite being inspired from reinforcement learning, the current GFlowNet framework is relatively limited in its applicability and cannot handle stochasticity in the reward function. In thi… ▽ More Generative Flow Networks (GFlowNets) are a new family of probabilistic samplers where an agent learns a stochastic policy for generating complex combinatorial structure through a series of decision-making steps. Despite being inspired from reinforcement learning, the current GFlowNet framework is relatively limited in its applicability and cannot handle stochasticity in the reward function. In this work, we adopt a distributional paradigm for GFlowNets, turning each flow function into a distribution, thus providing more informative learning signals during training. By parameterizing each edge flow through their quantile functions, our proposed \textit{quantile matching} GFlowNet learning algorithm is able to learn a risk-sensitive policy, an essential component for handling scenarios with risk uncertainty. Moreover, we find that the distributional approach can achieve substantial improvement on existing benchmarks compared to prior methods due to our enhanced training algorithm, even in settings with deterministic rewards. △ Less

Submitted 17 February, 2024; v1 submitted 11 February, 2023; originally announced February 2023.

Comments: Accepted by TMLR

Flow Matching on General Geometries

Authors: Ricky T. Q. Chen, Yaron Lipman

Abstract: We propose Riemannian Flow Matching (RFM), a simple yet powerful framework for training continuous normalizing flows on manifolds. Existing methods for generative modeling on manifolds either require expensive simulation, are inherently unable to scale to high dimensions, or use approximations for limiting quantities that result in biased training objectives. Riemannian Flow Matching bypasses thes… ▽ More We propose Riemannian Flow Matching (RFM), a simple yet powerful framework for training continuous normalizing flows on manifolds. Existing methods for generative modeling on manifolds either require expensive simulation, are inherently unable to scale to high dimensions, or use approximations for limiting quantities that result in biased training objectives. Riemannian Flow Matching bypasses these limitations and offers several advantages over previous approaches: it is simulation-free on simple geometries, does not require divergence computation, and computes its target vector field in closed-form. The key ingredient behind RFM is the construction of a relatively simple premetric for defining target vector fields, which encompasses the existing Euclidean case. To extend to general geometries, we rely on the use of spectral decompositions to efficiently compute premetrics on the fly. Our method achieves state-of-the-art performance on many real-world non-Euclidean datasets, and we demonstrate tractable training on general geometries, including triangular meshes with highly non-trivial curvature and boundaries. △ Less

Submitted 26 February, 2024; v1 submitted 7 February, 2023; originally announced February 2023.

Journal ref: ICLR 2024

A Point-of-Care Biosensor for Rapid Detection and Differentiation of COVID-19 Virus (SARS-CoV-2) and Influenza Virus Using Subwavelength Grating Micro-ring Resonator

Authors: Shupeng Ning, Hao-Chen Chang, Kang-Chieh Fan, Po-yu Hsiao, Chenghao Feng, Devan Shoemaker, Ray T. Chen

Abstract: In the context of continued spread of coronavirus disease 2019 (COVID-19) caused by SARS-CoV-2 and the emergence of new variants, the demand for rapid, accurate, and frequent detection is increasing. Besides, the new predominant strain, Omicron variant, manifests more similar clinical features to those of other common respiratory infections. The concurrent detection of multiple potential pathogens… ▽ More In the context of continued spread of coronavirus disease 2019 (COVID-19) caused by SARS-CoV-2 and the emergence of new variants, the demand for rapid, accurate, and frequent detection is increasing. Besides, the new predominant strain, Omicron variant, manifests more similar clinical features to those of other common respiratory infections. The concurrent detection of multiple potential pathogens helps distinguish SARS-CoV-2 infection from other diseases with overlapping symptoms, which is significant for patients to receive tailored treatment and containing the outbreak. Here, we report a lab-on-a-chip biosensing platform for SARS-CoV-2 detection based on subwavelength grating micro-ring resonator. The sensing surface is functionalized by specific antibody against SARS-CoV-2 spike protein, which could produce redshifts of resonant peaks by antigen-antibody combination, thus achieving quantitative detection. Additionally, the sensor chip is integrated with a microfluidic chip with an anti-backflow Y-shaped structure that enables the concurrent detection of two analytes. In this study, we realized the detection and differentiation of COVID-19 and influenza A H1N1. Experimental results show that the limit of detection of our device reaches 100 fg/mL (1.31 fM) within 15 min detecting time, and cross-reactivity tests manifest the specificity of the optical diagnostic assay. Further, the integrated packaging and streamlined workflow facilitate its use for clinical applications. Thus, the biosensing platform offers a promising solution to achieve ultrasensitive, selective, multiplexed, and quantitative point-of-care detection of COVID-19. △ Less

Submitted 11 January, 2023; originally announced January 2023.

Latent Discretization for Continuous-time Sequence Compression

Authors: Ricky T. Q. Chen, Matthew Le, Matthew Muckley, Maximilian Nickel, Karen Ullrich

Abstract: Neural compression offers a domain-agnostic approach to creating codecs for lossy or lossless compression via deep generative models. For sequence compression, however, most deep sequence models have costs that scale with the sequence length rather than the sequence complexity. In this work, we instead treat data sequences as observations from an underlying continuous-time process and learn how to… ▽ More Neural compression offers a domain-agnostic approach to creating codecs for lossy or lossless compression via deep generative models. For sequence compression, however, most deep sequence models have costs that scale with the sequence length rather than the sequence complexity. In this work, we instead treat data sequences as observations from an underlying continuous-time process and learn how to efficiently discretize while retaining information about the full sequence. As a consequence of decoupling sequential information from its temporal discretization, our approach allows for greater compression rates and smaller computational complexity. Moreover, the continuous-time approach naturally allows us to decode at different time intervals. We empirically verify our approach on multiple domains involving compression of video and motion capture sequences, showing that our approaches can automatically achieve reductions in bit rates by learning how to discretize. △ Less

Submitted 27 December, 2022; originally announced December 2022.

Flow Matching for Generative Modeling

Authors: Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le

Abstract: We introduce a new paradigm for generative modeling built on Continuous Normalizing Flows (CNFs), allowing us to train CNFs at unprecedented scale. Specifically, we present the notion of Flow Matching (FM), a simulation-free approach for training CNFs based on regressing vector fields of fixed conditional probability paths. Flow Matching is compatible with a general family of Gaussian probability… ▽ More We introduce a new paradigm for generative modeling built on Continuous Normalizing Flows (CNFs), allowing us to train CNFs at unprecedented scale. Specifically, we present the notion of Flow Matching (FM), a simulation-free approach for training CNFs based on regressing vector fields of fixed conditional probability paths. Flow Matching is compatible with a general family of Gaussian probability paths for transforming between noise and data samples -- which subsumes existing diffusion paths as specific instances. Interestingly, we find that employing FM with diffusion paths results in a more robust and stable alternative for training diffusion models. Furthermore, Flow Matching opens the door to training CNFs with other, non-diffusion probability paths. An instance of particular interest is using Optimal Transport (OT) displacement interpolation to define the conditional probability paths. These paths are more efficient than diffusion paths, provide faster training and sampling, and result in better generalization. Training CNFs using Flow Matching on ImageNet leads to consistently better performance than alternative diffusion-based methods in terms of both likelihood and sample quality, and allows fast and reliable sample generation using off-the-shelf numerical ODE solvers. △ Less

Submitted 8 February, 2023; v1 submitted 6 October, 2022; originally announced October 2022.

Neural Conservation Laws: A Divergence-Free Perspective

Authors: Jack Richter-Powell, Yaron Lipman, Ricky T. Q. Chen

Abstract: We investigate the parameterization of deep neural networks that by design satisfy the continuity equation, a fundamental conservation law. This is enabled by the observation that any solution of the continuity equation can be represented as a divergence-free vector field. We hence propose building divergence-free neural networks through the concept of differential forms, and with the aid of autom… ▽ More We investigate the parameterization of deep neural networks that by design satisfy the continuity equation, a fundamental conservation law. This is enabled by the observation that any solution of the continuity equation can be represented as a divergence-free vector field. We hence propose building divergence-free neural networks through the concept of differential forms, and with the aid of automatic differentiation, realize two practical constructions. As a result, we can parameterize pairs of densities and vector fields that always exactly satisfy the continuity equation, foregoing the need for extra penalty methods or expensive numerical simulation. Furthermore, we prove these models are universal and so can be used to represent any divergence-free vector field. Finally, we experimentally validate our approaches by computing neural network-based solutions to fluid equations, solving for the Hodge decomposition, and learning dynamical optimal transport maps. △ Less

Submitted 11 December, 2022; v1 submitted 4 October, 2022; originally announced October 2022.

Journal ref: NeurIPS 2022

Latent State Marginalization as a Low-cost Approach for Improving Exploration

Authors: Dinghuai Zhang, Aaron Courville, Yoshua Bengio, Qinqing Zheng, Amy Zhang, Ricky T. Q. Chen

Abstract: While the maximum entropy (MaxEnt) reinforcement learning (RL) framework -- often touted for its exploration and robustness capabilities -- is usually motivated from a probabilistic perspective, the use of deep probabilistic models has not gained much traction in practice due to their inherent complexity. In this work, we propose the adoption of latent variable policies within the MaxEnt framework… ▽ More While the maximum entropy (MaxEnt) reinforcement learning (RL) framework -- often touted for its exploration and robustness capabilities -- is usually motivated from a probabilistic perspective, the use of deep probabilistic models has not gained much traction in practice due to their inherent complexity. In this work, we propose the adoption of latent variable policies within the MaxEnt framework, which we show can provably approximate any policy distribution, and additionally, naturally emerges under the use of world models with a latent belief state. We discuss why latent variable policies are difficult to train, how naive approaches can fail, then subsequently introduce a series of improvements centered around low-cost marginalization of the latent state, allowing us to make full use of the latent state at minimal additional cost. We instantiate our method under the actor-critic framework, marginalizing both the actor and critic. The resulting algorithm, referred to as Stochastic Marginal Actor-Critic (SMAC), is simple yet effective. We experimentally validate our method on continuous control tasks, showing that effective marginalization can lead to better exploration and more robust training. Our implementation is open sourced at https://github.com/zdhNarsil/Stochastic-Marginal-Actor-Critic. △ Less

Submitted 10 February, 2023; v1 submitted 3 October, 2022; originally announced October 2022.

Comments: Accepted by ICLR 2023

NeurOLight: A Physics-Agnostic Neural Operator Enabling Parametric Photonic Device Simulation

Authors: Jiaqi Gu, Zhengqi Gao, Chenghao Feng, Hanqing Zhu, Ray T. Chen, Duane S. Boning, David Z. Pan

Abstract: Optical computing is an emerging technology for next-generation efficient artificial intelligence (AI) due to its ultra-high speed and efficiency. Electromagnetic field simulation is critical to the design, optimization, and validation of photonic devices and circuits. However, costly numerical simulation significantly hinders the scalability and turn-around time in the photonic circuit design loo… ▽ More Optical computing is an emerging technology for next-generation efficient artificial intelligence (AI) due to its ultra-high speed and efficiency. Electromagnetic field simulation is critical to the design, optimization, and validation of photonic devices and circuits. However, costly numerical simulation significantly hinders the scalability and turn-around time in the photonic circuit design loop. Recently, physics-informed neural networks have been proposed to predict the optical field solution of a single instance of a partial differential equation (PDE) with predefined parameters. Their complicated PDE formulation and lack of efficient parametrization mechanisms limit their flexibility and generalization in practical simulation scenarios. In this work, for the first time, a physics-agnostic neural operator-based framework, dubbed NeurOLight, is proposed to learn a family of frequency-domain Maxwell PDEs for ultra-fast parametric photonic device simulation. We balance the efficiency and generalization of NeurOLight via several novel techniques. Specifically, we discretize different devices into a unified domain, represent parametric PDEs with a compact wave prior, and encode the incident light via masked source modeling. We design our model with parameter-efficient cross-shaped NeurOLight blocks and adopt superposition-based augmentation for data-efficient learning. With these synergistic approaches, NeurOLight generalizes to a large space of unseen simulation settings, demonstrates 2-orders-of-magnitude faster simulation speed than numerical solvers, and outperforms prior neural network models by ~54% lower prediction error with ~44% fewer parameters. Our code is available at https://github.com/JeremieMelo/NeurOLight. △ Less

Submitted 19 September, 2022; originally announced September 2022.

Comments: 13 pages. Accepted to NeurIPS 2022

Unifying Generative Models with GFlowNets and Beyond

Authors: Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, Yoshua Bengio

Abstract: There are many frameworks for deep generative modeling, each often presented with their own specific training algorithms and inference methods. Here, we demonstrate the connections between existing deep generative models and the recently introduced GFlowNet framework, a probabilistic inference machine which treats sampling as a decision-making process. This analysis sheds light on their overlappin… ▽ More There are many frameworks for deep generative modeling, each often presented with their own specific training algorithms and inference methods. Here, we demonstrate the connections between existing deep generative models and the recently introduced GFlowNet framework, a probabilistic inference machine which treats sampling as a decision-making process. This analysis sheds light on their overlapping traits and provides a unifying viewpoint through the lens of learning with Markovian trajectories. Our framework provides a means for unifying training and inference algorithms, and provides a route to shine a unifying light over many generative models. Beyond this, we provide a practical and experimentally verified recipe for improving generative modeling with insights from the GFlowNet perspective. △ Less

Submitted 30 January, 2023; v1 submitted 6 September, 2022; originally announced September 2022.

Comments: expanded version of the ICML 2022 workshop paper

Theseus: A Library for Differentiable Nonlinear Optimization

Authors: Luis Pineda, Taosha Fan, Maurizio Monge, Shobha Venkataraman, Paloma Sodhi, Ricky T. Q. Chen, Joseph Ortiz, Daniel DeTone, Austin Wang, Stuart Anderson, Jing Dong, Brandon Amos, Mustafa Mukadam

Abstract: We present Theseus, an efficient application-agnostic open source library for differentiable nonlinear least squares (DNLS) optimization built on PyTorch, providing a common framework for end-to-end structured learning in robotics and vision. Existing DNLS implementations are application specific and do not always incorporate many ingredients important for efficiency. Theseus is application-agnost… ▽ More We present Theseus, an efficient application-agnostic open source library for differentiable nonlinear least squares (DNLS) optimization built on PyTorch, providing a common framework for end-to-end structured learning in robotics and vision. Existing DNLS implementations are application specific and do not always incorporate many ingredients important for efficiency. Theseus is application-agnostic, as we illustrate with several example applications that are built using the same underlying differentiable components, such as second-order optimizers, standard costs functions, and Lie groups. For efficiency, Theseus incorporates support for sparse solvers, automatic vectorization, batching, GPU acceleration, and gradient computation with implicit differentiation and direct loss minimization. We do extensive performance evaluation in a set of applications, demonstrating significant efficiency gains and better scalability when these features are incorporated. Project page: https://sites.google.com/view/theseus-ai △ Less

Submitted 18 January, 2023; v1 submitted 19 July, 2022; originally announced July 2022.

Comments: Advances in Neural Information Processing Systems (NeurIPS), 2022

Lab-on-a-Chip Optical Biosensor Platform: Micro Ring Resonator Integrated with Near-Infrared Fourier Transform Spectrometer

Authors: Kyoung Min Yoo, May Hlaing, Sourabh Jain, James Fan, Yue An, Ray T. Chen

Abstract: A micro-ring-resonator (MRR) optical biosensor based on the evanescent field sensing mechanism has been extensively studied due to its high sensitivity and compact device size. However, a suitable on-chip integrated spectrometer device has to be demonstrated for the lab-on-a-chip applications, which can read the resonance wavelength shift from MRR biosensors based on minuscule changes in refractiv… ▽ More A micro-ring-resonator (MRR) optical biosensor based on the evanescent field sensing mechanism has been extensively studied due to its high sensitivity and compact device size. However, a suitable on-chip integrated spectrometer device has to be demonstrated for the lab-on-a-chip applications, which can read the resonance wavelength shift from MRR biosensors based on minuscule changes in refractive index. In this paper, we demonstrated the design and experimental results of the near-infrared lab-on-a-chip optical biosensor platform that monolithically integrates the MRR and the on-chip spectrometer on the silicon-on-insulator (SOI) wafer, which can eliminate the external optical spectrum analyzer for scanning the wavelength spectrum. The symmetric add-drop MRR biosensor is designed to have a free spectral range (FSR) of ~19 nm, and a bulk sensitivity of ~73 nm/RIU; then the drop-port output resonance peaks are reconstructed from the integrated spatial-heterodyne Fourier transform spectrometer (SHFTS) with the spectral resolution of ~3.1 nm and bandwidth of ~50 nm, which results in the limit of detection of 0.042 RIU. The MRR output spectrum with air- and water-claddings are measured and reconstructed from the MRR-SHFTS integrated device experimentally to validate the wavelength shifting measurement. △ Less

Submitted 15 July, 2022; originally announced July 2022.

Comments: 23 pages, 9 figures including supplementary

Matching Normalizing Flows and Probability Paths on Manifolds

Authors: Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Aditya Grover, Maximilian Nickel, Ricky T. Q. Chen, Yaron Lipman

Abstract: Continuous Normalizing Flows (CNFs) are a class of generative models that transform a prior distribution to a model distribution by solving an ordinary differential equation (ODE). We propose to train CNFs on manifolds by minimizing probability path divergence (PPD), a novel family of divergences between the probability density path generated by the CNF and a target probability density path. PPD i… ▽ More Continuous Normalizing Flows (CNFs) are a class of generative models that transform a prior distribution to a model distribution by solving an ordinary differential equation (ODE). We propose to train CNFs on manifolds by minimizing probability path divergence (PPD), a novel family of divergences between the probability density path generated by the CNF and a target probability density path. PPD is formulated using a logarithmic mass conservation formula which is a linear first order partial differential equation relating the log target probabilities and the CNF's defining vector field. PPD has several key benefits over existing methods: it sidesteps the need to solve an ODE per iteration, readily applies to manifold data, scales to high dimensions, and is compatible with a large family of target paths interpolating pure noise and data in finite time. Theoretically, PPD is shown to bound classical probability divergences. Empirically, we show that CNFs learned by minimizing PPD achieve state-of-the-art results in likelihoods and sample quality on existing low-dimensional manifold benchmarks, and is the first example of a generative model to scale to moderately high dimensional manifolds. △ Less

Submitted 11 July, 2022; originally announced July 2022.

Comments: ICML 2022

Semi-Discrete Normalizing Flows through Differentiable Tessellation

Authors: Ricky T. Q. Chen, Brandon Amos, Maximilian Nickel

Abstract: Mapping between discrete and continuous distributions is a difficult task and many have had to resort to heuristical approaches. We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, complete with exact likelihood evaluations. This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism wit… ▽ More Mapping between discrete and continuous distributions is a difficult task and many have had to resort to heuristical approaches. We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, complete with exact likelihood evaluations. This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism with an efficient log determinant Jacobian. We explore this approach in two application settings, mapping from discrete to continuous and vice versa. Firstly, a Voronoi dequantization allows automatically learning quantization boundaries in a multidimensional space. The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values. Secondly, a Voronoi mixture model has near-constant computation cost for likelihood evaluation regardless of the number of mixture components. Empirically, we show improvements over existing methods across a range of structured data modalities. △ Less

Submitted 11 December, 2022; v1 submitted 13 March, 2022; originally announced March 2022.

Journal ref: NeurIPS 2022

Packaging-enhanced optical fiber-chip interconnect with enlarged grating coupler and multimode fiber

Authors: Chao Wang, Chingwen Chang, Jason Midkiff, Aref Asghari, James Fan, Jianying Zhou, Xiaochuan Xu, Huiping Tian, Ray T. Chen

Abstract: Optical I/O plays a crucial role in the lifespan of lab-on-a-chip systems, from preliminary testing to operation in the target environment. However, due to the precise alignments required, efficient and reliable fiber-to-chip connections remain challenging, yielding inconsistent test results and unstable packaged performance. To overcome this issue, for use in single mode on-chip systems, we propo… ▽ More Optical I/O plays a crucial role in the lifespan of lab-on-a-chip systems, from preliminary testing to operation in the target environment. However, due to the precise alignments required, efficient and reliable fiber-to-chip connections remain challenging, yielding inconsistent test results and unstable packaged performance. To overcome this issue, for use in single mode on-chip systems, we propose the incorporation of area-enlarged grating couplers working in conjunction with multimode fibers. This combination enables simpler, faster, and more reliable connections than the traditional small area grating coupler with single-mode fiber. In this work, we experimentally demonstrate a 3dB in-plane (X, Y) spatial tolerance of (10.2 μm, 17.3 μm) for the large area configuration, being at least (2.49, 3.33) times that of the small area one, and agreeing well with theoretical calculations. The simple concept is readily applicable to a range of photonic systems where cheaper more robust optical I/O is desired. △ Less

Submitted 18 January, 2022; originally announced January 2022.

Comments: 11 pages, 5 figures

ADEPT: Automatic Differentiable DEsign of Photonic Tensor Cores

Authors: Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Zixuan Jiang, Mingjie Liu, Shuhan Zhang, Ray T. Chen, David Z. Pan

Abstract: Photonic tensor cores (PTCs) are essential building blocks for optical artificial intelligence (AI) accelerators based on programmable photonic integrated circuits. PTCs can achieve ultra-fast and efficient tensor operations for neural network (NN) acceleration. Current PTC designs are either manually constructed or based on matrix decomposition theory, which lacks the adaptability to meet various… ▽ More Photonic tensor cores (PTCs) are essential building blocks for optical artificial intelligence (AI) accelerators based on programmable photonic integrated circuits. PTCs can achieve ultra-fast and efficient tensor operations for neural network (NN) acceleration. Current PTC designs are either manually constructed or based on matrix decomposition theory, which lacks the adaptability to meet various hardware constraints and device specifications. To our best knowledge, automatic PTC design methodology is still unexplored. It will be promising to move beyond the manual design paradigm and "nurture" photonic neurocomputing with AI and design automation. Therefore, in this work, for the first time, we propose a fully differentiable framework, dubbed ADEPT, that can efficiently search PTC designs adaptive to various circuit footprint constraints and foundry PDKs. Extensive experiments show superior flexibility and effectiveness of the proposed ADEPT framework to explore a large PTC design space. On various NN models and benchmarks, our searched PTC topology outperforms prior manually-designed structures with competitive matrix representability, 2-30x higher footprint compactness, and better noise robustness, demonstrating a new paradigm in photonic neural chip design. The code of ADEPT is available at https://github.com/JeremieMelo/ADEPT using the https://github.com/JeremieMelo/pytorch-onn (TorchONN) library. △ Less

Submitted 3 May, 2022; v1 submitted 16 December, 2021; originally announced December 2021.

Comments: Accepted to ACM/IEEE Design Automation Conference (DAC), 2022

ELight: Enabling Efficient Photonic In-Memory Neurocomputing with Life Enhancement

Authors: Hanqing Zhu, Jiaqi Gu, Chenghao Feng, Mingjie Liu, Zixuan Jiang, Ray T. Chen, David Z. Pan

Abstract: With the recent advances in optical phase change material (PCM), photonic in-memory neurocomputing has demonstrated its superiority in optical neural network (ONN) designs with near-zero static power consumption, time-of-light latency, and compact footprint. However, photonic tensor cores require massive hardware reuse to implement large matrix multiplication due to the limited single-core scale.… ▽ More With the recent advances in optical phase change material (PCM), photonic in-memory neurocomputing has demonstrated its superiority in optical neural network (ONN) designs with near-zero static power consumption, time-of-light latency, and compact footprint. However, photonic tensor cores require massive hardware reuse to implement large matrix multiplication due to the limited single-core scale. The resultant large number of PCM writes leads to serious dynamic power and overwhelms the fragile PCM with limited write endurance. In this work, we propose a synergistic optimization framework, ELight, to minimize the overall write efforts for efficient and reliable optical in-memory neurocomputing. We first propose write-aware training to encourage the similarity among weight blocks, and combine it with a post-training optimization method to reduce programming efforts by eliminating redundant writes. Experiments show that ELight can achieve over 20X reduction in the total number of writes and dynamic power with comparable accuracy. With our ELight, photonic in-memory neurocomputing will step forward towards viable applications in machine learning with preserved accuracy, order-of-magnitude longer lifetime, and lower programming energy. △ Less

Submitted 15 December, 2021; originally announced December 2021.

Comments: 7 pages, 8 figures, accepted by ASPDAC 2022

Dual-Polarization Bandwidth-Bridged On-Chip Bandpass Sampling Fourier Transform Spectrometer from Visible to Near-Infrared

Authors: Kyoung Min Yoo, Ray T. Chen

Abstract: The on-chip broadband optical spectrometers which cover the entire tissue transparency window (λ=650-1050 nm) with high resolution are highly demanded for the miniaturized bio-sensing and bio-imaging applications. Here, we propose a novel type of spatial heterodyne Fourier transform spectrometer (SHFTS) integrated with a sub-wavelength grating coupler (SWGC) for the dual-polarization bandpass samp… ▽ More The on-chip broadband optical spectrometers which cover the entire tissue transparency window (λ=650-1050 nm) with high resolution are highly demanded for the miniaturized bio-sensing and bio-imaging applications. Here, we propose a novel type of spatial heterodyne Fourier transform spectrometer (SHFTS) integrated with a sub-wavelength grating coupler (SWGC) for the dual-polarization bandpass sampling on the Si3N4 platform. Through tuning the coupling angles with different polarization, we experimentally demonstrated the unprecedented broadband spectrum retrieval results with the overall bandwidth coverage of 400 nm, bridging the wavelengths from 650 nm to 1050 nm, with the resolution of 2-5 nm. By applying the bandpass sampling theorem, we circumvented the intrinsic trade-off limitation between the bandwidth and resolution of SHFTS without having an outrageous number of Mach-Zehnder interferometer (MZI) arrays or adding additional active components. The bandpass sampling SHFTS is designed to have linearly unbalanced 32 MZIs with the maximum optical path length difference of 93 μm within an overall footprint size of 4.7 mm x 0.65 mm, and the coupling angles of SWGC are varied from 0° to 32° to cover the entire tissue transparency window. △ Less

Submitted 13 December, 2021; originally announced December 2021.

Comments: 48 Pages, 6 figures, 14 supportive figures

A compact butterfly-style silicon photonic-electronic neural chip for hardware-efficient deep learning

Authors: Chenghao Feng, Jiaqi Gu, Hanqing Zhu, Zhoufeng Ying, Zheng Zhao, David Z. Pan, Ray T. Chen

Abstract: The optical neural network (ONN) is a promising hardware platform for next-generation neurocomputing due to its high parallelism, low latency, and low energy consumption. Previous ONN architectures are mainly designed for general matrix multiplication (GEMM), leading to unnecessarily large area cost and high control complexity. Here, we move beyond classical GEMM-based ONNs and propose an optical… ▽ More The optical neural network (ONN) is a promising hardware platform for next-generation neurocomputing due to its high parallelism, low latency, and low energy consumption. Previous ONN architectures are mainly designed for general matrix multiplication (GEMM), leading to unnecessarily large area cost and high control complexity. Here, we move beyond classical GEMM-based ONNs and propose an optical subspace neural network (OSNN) architecture, which trades the universality of weight representation for lower optical component usage, area cost, and energy consumption. We devise a butterfly-style photonic-electronic neural chip to implement our OSNN with up to 7x fewer trainable optical components compared to GEMM-based ONNs. Additionally, a hardware-aware training framework is provided to minimize the required device programming precision, lessen the chip area, and boost the noise robustness. We experimentally demonstrate the utility of our neural chip in practical image recognition tasks, showing that a measured accuracy of 94.16% can be achieved in hand-written digit recognition tasks with 3-bit weight programming precision. △ Less

Submitted 17 July, 2022; v1 submitted 11 November, 2021; originally announced November 2021.

Comments: 17 pages,5 figures

L2ight: Enabling On-Chip Learning for Optical Neural Networks via Efficient in-situ Subspace Optimization

Authors: Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Zixuan Jiang, Ray T. Chen, David Z. Pan

Abstract: Silicon-photonics-based optical neural network (ONN) is a promising hardware platform that could represent a paradigm shift in efficient AI with its CMOS-compatibility, flexibility, ultra-low execution latency, and high energy efficiency. In-situ training on the online programmable photonic chips is appealing but still encounters challenging issues in on-chip implementability, scalability, and eff… ▽ More Silicon-photonics-based optical neural network (ONN) is a promising hardware platform that could represent a paradigm shift in efficient AI with its CMOS-compatibility, flexibility, ultra-low execution latency, and high energy efficiency. In-situ training on the online programmable photonic chips is appealing but still encounters challenging issues in on-chip implementability, scalability, and efficiency. In this work, we propose a closed-loop ONN on-chip learning framework L2ight to enable scalable ONN mapping and efficient in-situ learning. L2ight adopts a three-stage learning flow that first calibrates the complicated photonic circuit states under challenging physical constraints, then performs photonic core mapping via combined analytical solving and zeroth-order optimization. A subspace learning procedure with multi-level sparsity is integrated into L2ight to enable in-situ gradient evaluation and fast adaptation, unleashing the power of optics for real on-chip intelligence. Extensive experiments demonstrate our proposed L2ight outperforms prior ONN training protocols with 3-order-of-magnitude higher scalability and over 30X better efficiency, when benchmarked on various models and learning tasks. This synergistic framework is the first scalable on-chip learning solution that pushes this emerging field from intractable to scalable and further to efficient for next-generation self-learnable photonic neural chips. From a co-design perspective, L2ight also provides essential insights for hardware-restricted unitary subspace optimization and efficient sparse training. We open-source our framework at https://github.com/JeremieMelo/L2ight. △ Less

Submitted 27 October, 2021; originally announced October 2021.

Comments: 10 pages. Accepted to NeurIPS 2021

Towards Memory-Efficient Neural Networks via Multi-Level in situ Generation

Authors: Jiaqi Gu, Hanqing Zhu, Chenghao Feng, Mingjie Liu, Zixuan Jiang, Ray T. Chen, David Z. Pan

Abstract: Deep neural networks (DNN) have shown superior performance in a variety of tasks. As they rapidly evolve, their escalating computation and memory demands make it challenging to deploy them on resource-constrained edge devices. Though extensive efficient accelerator designs, from traditional electronics to emerging photonics, have been successfully demonstrated, they are still bottlenecked by expen… ▽ More Deep neural networks (DNN) have shown superior performance in a variety of tasks. As they rapidly evolve, their escalating computation and memory demands make it challenging to deploy them on resource-constrained edge devices. Though extensive efficient accelerator designs, from traditional electronics to emerging photonics, have been successfully demonstrated, they are still bottlenecked by expensive memory accesses due to tremendous gaps between the bandwidth/power/latency of electrical memory and computing cores. Previous solutions fail to fully-leverage the ultra-fast computational speed of emerging DNN accelerators to break through the critical memory bound. In this work, we propose a general and unified framework to trade expensive memory transactions with ultra-fast on-chip computations, directly translating to performance improvement. We are the first to jointly explore the intrinsic correlations and bit-level redundancy within DNN kernels and propose a multi-level in situ generation mechanism with mixed-precision bases to achieve on-the-fly recovery of high-resolution parameters with minimum hardware overhead. Extensive experiments demonstrate that our proposed joint method can boost the memory efficiency by 10-20x with comparable accuracy over four state-of-the-art designs, when benchmarked on ResNet-18/DenseNet-121/MobileNetV2/V3 with various tasks. △ Less

Submitted 5 September, 2021; v1 submitted 25 August, 2021; originally announced August 2021.

Comments: Accepted by International Conference on Computer Vision (ICCV) 2021

Fully differentiable optimization protocols for non-equilibrium steady states

Authors: Rodrigo A. Vargas-Hernández, Ricky T. Q. Chen, Kenneth A. Jung, Paul Brumer

Abstract: In the case of quantum systems interacting with multiple environments, the time-evolution of the reduced density matrix is described by the Liouvillian. For a variety of physical observables, the long-time limit or steady state solution is needed for the computation of desired physical observables. For inverse design or optimal control of such systems, the common approaches are based on brute-forc… ▽ More In the case of quantum systems interacting with multiple environments, the time-evolution of the reduced density matrix is described by the Liouvillian. For a variety of physical observables, the long-time limit or steady state solution is needed for the computation of desired physical observables. For inverse design or optimal control of such systems, the common approaches are based on brute-force search strategies. Here, we present a novel methodology, based on automatic differentiation, capable of differentiating the steady state solution with respect to any parameter of the Liouvillian. Our approach has a low memory cost, and is agnostic to the exact algorithm for computing the steady state. We illustrate the advantage of this method by inverse designing the parameters of a quantum heat transfer device that maximizes the heat current and the rectification coefficient. Additionally, we optimize the parameters of various Lindblad operators used in the simulation of energy transfer under natural incoherent light. We also present a sensitivity analysis of the steady state for energy transfer under natural incoherent light as a function of the incoherent-light pumping rate. △ Less

Submitted 23 November, 2021; v1 submitted 23 March, 2021; originally announced March 2021.

Comments: Main work 10 pages and 5 Figures. Supplemental Material 12 pages, 1 Figure, 3 Tables

Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations

Authors: Winnie Xu, Ricky T. Q. Chen, Xuechen Li, David Duvenaud

Abstract: We perform scalable approximate inference in continuous-depth Bayesian neural networks. In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation. We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors. We also derive a nove… ▽ More We perform scalable approximate inference in continuous-depth Bayesian neural networks. In this model class, uncertainty about separate weights in each layer gives hidden units that follow a stochastic differential equation. We demonstrate gradient-based stochastic variational inference in this infinite-parameter setting, producing arbitrarily-flexible approximate posteriors. We also derive a novel gradient estimator that approaches zero variance as the approximate posterior over weights approaches the true posterior. This approach brings continuous-depth Bayesian neural nets to a competitive comparison against discrete-depth alternatives, while inheriting the memory-efficient training and tunable precision of Neural ODEs. △ Less

Submitted 30 January, 2022; v1 submitted 12 February, 2021; originally announced February 2021.

Efficient On-Chip Learning for Optical Neural Networks Through Power-Aware Sparse Zeroth-Order Optimization

Authors: Jiaqi Gu, Chenghao Feng, Zheng Zhao, Zhoufeng Ying, Ray T. Chen, David Z. Pan

Abstract: Optical neural networks (ONNs) have demonstrated record-breaking potential in high-performance neuromorphic computing due to their ultra-high execution speed and low energy consumption. However, current learning protocols fail to provide scalable and efficient solutions to photonic circuit optimization in practical applications. In this work, we propose a novel on-chip learning framework to releas… ▽ More Optical neural networks (ONNs) have demonstrated record-breaking potential in high-performance neuromorphic computing due to their ultra-high execution speed and low energy consumption. However, current learning protocols fail to provide scalable and efficient solutions to photonic circuit optimization in practical applications. In this work, we propose a novel on-chip learning framework to release the full potential of ONNs for power-efficient in situ training. Instead of deploying implementation-costly back-propagation, we directly optimize the device configurations with computation budgets and power constraints. We are the first to model the ONN on-chip learning as a resource-constrained stochastic noisy zeroth-order optimization problem, and propose a novel mixed-training strategy with two-level sparsity and power-aware dynamic pruning to offer a scalable on-chip training solution in practical ONN deployment. Compared with previous methods, we are the first to optimize over 2,500 optical components on chip. We can achieve much better optimization stability, 3.7x-7.6x higher efficiency, and save >90% power under practical device variations and thermal crosstalk. △ Less

Submitted 5 September, 2021; v1 submitted 21 December, 2020; originally announced December 2020.

Comments: 7 pages content, 2 pages of reference, 6 figures, 4 tables, accepted to Association for the Advancement of Artificial Intelligence (AAAI) 2021

Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization

Authors: Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, Aaron Courville

Abstract: Flow-based models are powerful tools for designing probabilistic models with tractable density. This paper introduces Convex Potential Flows (CP-Flow), a natural and efficient parameterization of invertible models inspired by the optimal transport (OT) theory. CP-Flows are the gradient map of a strongly convex neural potential function. The convexity implies invertibility and allows us to resort t… ▽ More Flow-based models are powerful tools for designing probabilistic models with tractable density. This paper introduces Convex Potential Flows (CP-Flow), a natural and efficient parameterization of invertible models inspired by the optimal transport (OT) theory. CP-Flows are the gradient map of a strongly convex neural potential function. The convexity implies invertibility and allows us to resort to convex optimization to solve the convex conjugate for efficient inversion. To enable maximum likelihood training, we derive a new gradient estimator of the log-determinant of the Jacobian, which involves solving an inverse-Hessian vector product using the conjugate gradient method. The gradient estimator has constant-memory cost, and can be made effectively unbiased by reducing the error tolerance level of the convex optimization routine. Theoretically, we prove that CP-Flows are universal density approximators and are optimal in the OT sense. Our empirical results show that CP-Flow performs competitively on standard benchmarks of density estimation and variational inference. △ Less

Submitted 23 February, 2021; v1 submitted 10 December, 2020; originally announced December 2020.

Inverse design of dissipative quantum steady-states with implicit differentiation

Authors: Rodrigo A. Vargas-Hernández, Ricky T. Q. Chen, Kenneth A. Jung, Paul Brumer

Abstract: Inverse design of a property that depends on the steady-state of an open quantum system is commonly done by grid-search type of methods. In this paper we present a new methodology that allows us to compute the gradient of the steady-state of an open quantum system with respect to any parameter of the Hamiltonian using the implicit differentiation theorem. As an example, we present a simulation of… ▽ More Inverse design of a property that depends on the steady-state of an open quantum system is commonly done by grid-search type of methods. In this paper we present a new methodology that allows us to compute the gradient of the steady-state of an open quantum system with respect to any parameter of the Hamiltonian using the implicit differentiation theorem. As an example, we present a simulation of a spin-boson model where the steady-state solution is obtained using Redfield theory. △ Less

Submitted 25 November, 2020; originally announced November 2020.

Comments: 6 pages, 2 figures, accepted for publication in Third Workshop on Machine Learning and the Physical Sciences (NeurIPS 2020), Vancouver, Canada

Self-Tuning Stochastic Optimization with Curvature-Aware Gradient Filtering

Authors: Ricky T. Q. Chen, Dami Choi, Lukas Balles, David Duvenaud, Philipp Hennig

Abstract: Standard first-order stochastic optimization algorithms base their updates solely on the average mini-batch gradient, and it has been shown that tracking additional quantities such as the curvature can help de-sensitize common hyperparameters. Based on this intuition, we explore the use of exact per-sample Hessian-vector products and gradients to construct optimizers that are self-tuning and hyper… ▽ More Standard first-order stochastic optimization algorithms base their updates solely on the average mini-batch gradient, and it has been shown that tracking additional quantities such as the curvature can help de-sensitize common hyperparameters. Based on this intuition, we explore the use of exact per-sample Hessian-vector products and gradients to construct optimizers that are self-tuning and hyperparameter-free. Based on a dynamics model of the gradient, we derive a process which leads to a curvature-corrected, noise-adaptive online gradient estimate. The smoothness of our updates makes it more amenable to simple step size selection schemes, which we also base off of our estimates quantities. We prove that our model-based procedure converges in the noisy quadratic setting. Though we do not see similar gains in deep learning tasks, we can match the performance of well-tuned optimizers and ultimately, this is an interesting step for constructing self-tuning optimizers. △ Less

Submitted 9 November, 2020; originally announced November 2020.

Neural Spatio-Temporal Point Processes

Authors: Ricky T. Q. Chen, Brandon Amos, Maximilian Nickel

Abstract: We propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and enable flexible, high-fidelity models of discrete events that are localized in continuous time and space. Central to our approach is a combination of continuous-time neural networks with two novel neural architectures, i.e., Jump and Attentive Continuous-time Nor… ▽ More We propose a new class of parameterizations for spatio-temporal point processes which leverage Neural ODEs as a computational method and enable flexible, high-fidelity models of discrete events that are localized in continuous time and space. Central to our approach is a combination of continuous-time neural networks with two novel neural architectures, i.e., Jump and Attentive Continuous-time Normalizing Flows. This approach allows us to learn complex distributions for both the spatial and temporal domain and to condition non-trivially on the observed event history. We validate our models on data sets from a wide variety of contexts such as seismology, epidemiology, urban mobility, and neuroscience. △ Less

Submitted 17 March, 2021; v1 submitted 9 November, 2020; originally announced November 2020.

Journal ref: ICLR 2021