Feedforward computation, such as evaluating a neural network or sampling from an autoregressive model, is ubiquitous in machine learning. The sequential nature of feedforward computation, however, requires a strict order of execution and cannot be easily accelerated with parallel computing. To enable parallelization, we frame the task of feedforward computation as solving a system of nonlinear equations. We then propose to find the solution using a Jacobi or Gauss-Seidel fixed-point iteration method, as well as hybrid methods of both. Crucially, Jacobi updates operate independently on each equation and can be executed in parallel. Our method is guaranteed to give exactly the same values as the original feedforward computation with a reduced (or equal) number of parallelizable iterations, and hence reduced time given sufficient parallel computing power. Experimentally, we demonstrate the effectiveness of our approach in accelerating (i) backpropagation of RNNs, (ii) evaluation of DenseNets, and (iii) autoregressive sampling of MADE and PixelCNN++, with speedup factors between 2.1 and 26 under various settings.
JFB: Jacobian-Free Backpropagation for Implicit Networks
Colorado School of Mines(科罗拉多矿业学院);Typal Research, Typal LLC;Alibaba Group (US), Damo Academy(阿里巴巴集团(美国),达摩院);University of California, Los Angeles(加利福尼亚大学洛杉矶分校)
Accelerating Parallel Sampling of Diffusion Models
School of Science and Engineering, The Chinese University of Hong Kong, Shenzhen(香港中文大学(深圳)理工学院);DAMO Academy, Alibaba Group(阿里巴巴集团达摩院);Shenzhen Research Institute of Big Data(深圳大数据研究院)
Future-proofing geotechnics workflows: accelerating problem-solving with large language models
Research Organization of Information and Systems(信息与系统研究机构);The Institute of Statistical Mathematics(统计数学研究所);The Graduate University for Advanced Studies(综合研究大学院大学);Tohoku University(东北大学);Kyoto University(京都大学);Coastal Development Institute of Technology(沿海开发技术研究所);NIKKEN SEKKEI LTD(日建设计有限公司)
DOF: Accelerating High-order Differential Operators with Forward Propagation
National Key Laboratory of General Artificial Intelligence, Peking University(北京大学通用人工智能国家重点实验室);California Institute of Technology(加州理工学院);Stanford University(斯坦福大学);Center for Machine Learning Research, Peking University(北京大学机器学习研究中心)
Accelerating Proximal Policy Optimization Learning Using Task Prediction for Solving Environments with Delayed Rewards
Boston University(波士顿大学);Worcester Polytechnic Institute(伍斯特理工学院);MIT Lincoln Laboratory(麻省理工学院林肯实验室);Anduril Industries(安杜里尔工业公司);Lehigh University(利哈伊大学);University of Maryland(马里兰大学)
Chebyshev Spectral Neural Networks for Solving Partial Differential Equations
School of Mathematical Sciences, Shanghai Jiao Tong University(上海交通大学数学科学学院);Institute of Natural Sciences, Shanghai Jiao Tong University(上海交通大学自然科学研究院)