Rose: Composable Autodiff for the Interactive Web

Authors Sam Estep , Wode Ni , Raven Rothkopf , Joshua Sunshine



PDF
Thumbnail PDF

File

LIPIcs.ECOOP.2024.15.pdf
  • Filesize: 1.98 MB
  • 27 pages

Document Identifiers

Author Details

Sam Estep
  • Software and Societal Systems Department, Carnegie Mellon University, Pittsburgh, PA, USA
Wode Ni
  • Software and Societal Systems Department, Carnegie Mellon University, Pittsburgh, PA, USA
Raven Rothkopf
  • Barnard College, Columbia University, New York, NY, USA
Joshua Sunshine
  • Software and Societal Systems Department, Carnegie Mellon University, Pittsburgh, PA, USA

Acknowledgements

Thanks to Adam Paszke for corresponding about JAX and Dex. The Rose icons were created by Aaron Weiss; we use them via the CC BY 4.0 license.

Cite AsGet BibTex

Sam Estep, Wode Ni, Raven Rothkopf, and Joshua Sunshine. Rose: Composable Autodiff for the Interactive Web. In 38th European Conference on Object-Oriented Programming (ECOOP 2024). Leibniz International Proceedings in Informatics (LIPIcs), Volume 313, pp. 15:1-15:27, Schloss Dagstuhl – Leibniz-Zentrum für Informatik (2024)
https://doi.org/10.4230/LIPIcs.ECOOP.2024.15

Abstract

Reverse-mode automatic differentiation (autodiff) has been popularized by deep learning, but its ability to compute gradients is also valuable for interactive use cases such as bidirectional computer-aided design, embedded physics simulations, visualizing causal inference, and more. Unfortunately, the web is ill-served by existing autodiff frameworks, which use autodiff strategies that perform poorly on dynamic scalar programs, and pull in heavy dependencies that would result in unacceptable webpage sizes. This work introduces Rose, a lightweight autodiff framework for the web using a new hybrid approach to reverse-mode autodiff, blending conventional tracing and transformation techniques in a way that uses the host language for metaprogramming while also allowing the programmer to explicitly define reusable functions that comprise a larger differentiable computation. We demonstrate the value of the Rose design by porting two differentiable physics simulations, and evaluate its performance on an optimization-based diagramming application, showing Rose outperforming the state-of-the-art in web-based autodiff by multiple orders of magnitude.

Subject Classification

ACM Subject Classification
  • Software and its engineering → Compilers
  • Information systems → Web applications
  • Software and its engineering → Domain specific languages
  • Computing methodologies → Symbolic and algebraic manipulation
  • Software and its engineering → Formal language definitions
  • General and reference → Performance
  • Computing methodologies → Neural networks
  • General and reference → General conference proceedings
Keywords
  • Automatic differentiation
  • differentiable programming
  • compilers
  • web

Metrics

  • Access Statistics
  • Total Accesses (updated on a weekly basis)
    0
    PDF Downloads

References

  1. Martín Abadi and Gordon D. Plotkin. A simple differentiable programming language. Proc. ACM Program. Lang., 4(POPL), December 2019. URL: https://doi.org/10.1145/3371106.
  2. Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mane, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viegas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous distributed systems, 2016. URL: https://arxiv.org/abs/1603.04467.
  3. Anthony Anjorin, Li-yao Xia, and Vadim Zaytsev. Bidirectional transformations wiki, 2011. URL: http://bx-community.wikidot.com/.
  4. Francis J. Anscombe. Graphs in statistical analysis. The American Statistician, 27(1):17-21, 1973. Google Scholar
  5. Hudson Ayers, Evan Laufer, Paul Mure, Jaehyeon Park, Eduardo Rodelo, Thea Rossman, Andrey Pronin, Philip Levis, and Johnathan Van Why. Tighten Rust’s belt: Shrinking embedded Rust binaries. In Proceedings of the 23rd ACM SIGPLAN/SIGBED International Conference on Languages, Compilers, and Tools for Embedded Systems, LCTES 2022, pages 121-132, New York, NY, USA, 2022. Association for Computing Machinery. URL: https://doi.org/10.1145/3519941.3535075.
  6. Gilbert Bernstein, Michael Mara, Tzu-Mao Li, Dougal Maclaurin, and Jonathan Ragan-Kelley. Differentiating a tensor language, 2020. URL: https://arxiv.org/abs/2008.11256.
  7. Jeff Bezanson, Stefan Karpinski, Viral B. Shah, and Alan Edelman. Julia: A fast dynamic language for technical computing, 2012. URL: https://arxiv.org/abs/1209.5145.
  8. Dan Cascaval, Mira Shalah, Phillip Quinn, Rastislav Bodik, Maneesh Agrawala, and Adriana Schulz. Differentiable 3D CAD programs for bidirectional editing. Computer Graphics Forum, 41(2):309-323, 2022. URL: https://doi.org/10.1111/cgf.14476.
  9. Bartosz Ciechanowski, 2014. URL: https://ciechanow.ski/.
  10. Sam Estep, Wode Ni, Raven Rothkopf, and Joshua Sunshine. Rose: Composable autodiff for the interactive web, 2024. URL: https://arxiv.org/abs/2402.17743.
  11. Figma, Inc. Figma, 2016. URL: https://figma.com/.
  12. Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018. Google Scholar
  13. Google LLC. Google Slides, 2006. URL: https://google.com/slides.
  14. Andreas Griewank and Andrea Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation. SIAM, 2008. Google Scholar
  15. Andreas Haas, Andreas Rossberg, Derek L. Schuff, Ben L. Titzer, Michael Holman, Dan Gohman, Luke Wagner, Alon Zakai, and JF Bastien. Bringing the Web up to speed with WebAssembly. In Proceedings of the 38th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2017, pages 185-200, New York, NY, USA, 2017. Association for Computing Machinery. URL: https://doi.org/10.1145/3062341.3062363.
  16. Laurent Hascoet and Valérie Pascual. The Tapenade automatic differentiation tool: Principles, model, and specification. ACM Trans. Math. Softw., 39(3), May 2013. URL: https://doi.org/10.1145/2450153.2450158.
  17. Brian Hempel, Justin Lubin, and Ravi Chugh. Sketch-n-Sketch: Output-directed programming for SVG. In Proceedings of the 32nd Annual ACM Symposium on User Interface Software and Technology, UIST '19, pages 281-292, New York, NY, USA, 2019. Association for Computing Machinery. URL: https://doi.org/10.1145/3332165.3347925.
  18. Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Ragan-Kelley, and Frédo Durand. DiffTaichi: Differentiable programming for physical simulation, 2020. URL: https://arxiv.org/abs/1910.00935.
  19. Yuanming Hu, Tzu-Mao Li, Luke Anderson, Jonathan Ragan-Kelley, and Frédo Durand. Taichi: A language for high-performance computation on spatially sparse data structures. ACM Trans. Graph., 38(6), November 2019. URL: https://doi.org/10.1145/3355089.3356506.
  20. Michael Innes. Don't unroll adjoint: Differentiating SSA-form programs, 2019. URL: https://arxiv.org/abs/1810.07951.
  21. Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckas, Elliot Saba, Viral B. Shah, and Will Tebbutt. ∂ P: A differentiable programming system to bridge machine learning and scientific computing, 2019. URL: https://arxiv.org/abs/1907.07587.
  22. Wenzel Jakob, Sébastien Speierer, Nicolas Roussel, and Delio Vicini. Dr.Jit: A just-in-time compiler for differentiable rendering. ACM Trans. Graph., 41(4), July 2022. URL: https://doi.org/10.1145/3528223.3530099.
  23. Jerzy Karczmarczuk. Functional differentiation of computer programs. In Proceedings of the Third ACM SIGPLAN International Conference on Functional Programming, ICFP '98, pages 195-203, New York, NY, USA, 1998. Association for Computing Machinery. URL: https://doi.org/10.1145/289423.289442.
  24. Benjamin Kenwright. Introduction to the WebGPU API. In ACM SIGGRAPH 2022 Courses, SIGGRAPH '22, New York, NY, USA, 2022. Association for Computing Machinery. URL: https://doi.org/10.1145/3532720.3535625.
  25. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL: https://arxiv.org/abs/1412.6980.
  26. Chris Lattner and Vikram Adve. LLVM: A compilation framework for lifelong program analysis & transformation. In International Symposium on Code Generation and Optimization, 2004. CGO 2004., pages 75-86, 2004. URL: https://doi.org/10.1109/CGO.2004.1281665.
  27. Dong C. Liu and Jorge Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programming, 45(1):503-528, 1989. Google Scholar
  28. Oleksandr Manzyuk, Barak A. Pearlmutter, Alexey Andreyevich Radul, David R. Rush, and Jeffrey Mark Siskind. Perturbation confusion in forward automatic differentiation of higher-order functions. Journal of Functional Programming, 29:e12, 2019. URL: https://doi.org/10.1017/S095679681900008X.
  29. William Moses and Valentin Churavy. Instead of rewriting foreign code for machine learning, automatically synthesize fast gradients. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 12472-12485. Curran Associates, Inc., 2020. URL: https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf.
  30. William S. Moses, Valentin Churavy, Ludger Paehler, Jan Hückelheim, Sri Hari Krishna Narayanan, Michel Schanen, and Johannes Doerfert. Reverse-mode automatic differentiation and optimization of GPU kernels via Enzyme. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '21, New York, NY, USA, 2021. Association for Computing Machinery. URL: https://doi.org/10.1145/3458817.3476165.
  31. William S. Moses, Sri Hari Krishna Narayanan, Ludger Paehler, Valentin Churavy, Michel Schanen, Jan Hückelheim, Johannes Doerfert, and Paul Hovland. Scalable automatic differentiation of multiple parallel paradigms through compiler augmentation. In Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis, SC '22. IEEE Press, 2022. Google Scholar
  32. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. dquotesingle Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL: https://proceedings.neurips.cc/paper_files/paper/2019/file/bdbca288fee7f92f2bfa9f7012727740-Paper.pdf.
  33. Adam Paszke, Daniel D. Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew J. Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. Getting to the point: Index sets and parallelism-preserving autodiff for pointful array programming. Proc. ACM Program. Lang., 5(ICFP), August 2021. URL: https://doi.org/10.1145/3473593.
  34. Amit Patel. Red Blob Games, 2013. URL: https://www.redblobgames.com/.
  35. Barak A. Pearlmutter and Jeffrey Mark Siskind. Lazy multivariate higher-order forward-mode AD. In Proceedings of the 34th Annual ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages, POPL '07, pages 155-160, New York, NY, USA, 2007. Association for Computing Machinery. URL: https://doi.org/10.1145/1190216.1190242.
  36. Barak A. Pearlmutter and Jeffrey Mark Siskind. Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator. ACM Trans. Program. Lang. Syst., 30(2), March 2008. URL: https://doi.org/10.1145/1330017.1330018.
  37. Pyodide contributors and Mozilla. Pyodide, 2019. URL: https://pyodide.org/.
  38. Dan Quinlan and Chunhua Liao. The ROSE source-to-source compiler infrastructure. In Cetus Users and Compiler Infrastructure Workshop, in conjunction with PACT, volume 2011, page 1. Citeseer, 2011. Google Scholar
  39. Alexey Radul, Adam Paszke, Roy Frostig, Matthew J. Johnson, and Dougal Maclaurin. You only linearize once: Tangents transpose to gradients. Proc. ACM Program. Lang., 7(POPL), January 2023. URL: https://doi.org/10.1145/3571236.
  40. Amir Shaikhha, Andrew Fitzgibbon, Simon Peyton Jones, and Dimitrios Vytiniotis. Destination-passing style for efficient memory management. In Proceedings of the 6th ACM SIGPLAN International Workshop on Functional High-Performance Computing, FHPC 2017, pages 12-23, New York, NY, USA, 2017. Association for Computing Machinery. URL: https://doi.org/10.1145/3122948.3122949.
  41. Xipeng Shen, Guoqiang Zhang, Irene Dea, Samantha Andow, Emilio Arroyo-Fang, Neal Gafter, Johann George, Melissa Grueter, Erik Meijer, Olin Grigsby Shivers, Steffi Stumpos, Alanna Tempest, Christy Warden, and Shannon Yang. Coarsening optimization for differentiable programming. Proc. ACM Program. Lang., 5(OOPSLA), October 2021. URL: https://doi.org/10.1145/3485507.
  42. Jeffrey Mark Siskind and Barak A Pearlmutter. Nesting forward-mode AD in a functional framework. Higher-Order and Symbolic Computation, 21(4):361-376, 2008. Google Scholar
  43. Jeffrey Mark Siskind and Barak A. Pearlmutter. Divide-and-conquer checkpointing for arbitrary programs with no user annotation. Optimization Methods and Software, 33(4-6):1288-1330, 2018. URL: https://doi.org/10.1080/10556788.2018.1459621.
  44. Tom J. Smeding and Matthijs I. L. Vákár. Efficient dual-numbers reverse AD via well-known program transformations. Proc. ACM Program. Lang., 7(POPL), January 2023. URL: https://doi.org/10.1145/3571247.
  45. Tom J. Smeding and Matthijs I. L. Vákár. Efficient CHAD. Proc. ACM Program. Lang., 8(POPL), January 2024. URL: https://doi.org/10.1145/3632878.
  46. Daniel Smilkov, Nikhil Thorat, Yannick Assogba, Charles Nicholson, Nick Kreeger, Ping Yu, Shanqing Cai, Eric Nielsen, David Soegel, Stan Bileschi, Michael Terry, Ann Yuan, Kangyi Zhang, Sandeep Gupta, Sarah Sirajuddin, D Sculley, Rajat Monga, Greg Corrado, Fernanda Viegas, and Martin M Wattenberg. TensorFlow.js: Machine learning for the web and beyond. In A. Talwalkar, V. Smith, and M. Zaharia, editors, Proceedings of Machine Learning and Systems, volume 1, pages 309-321, 2019. URL: https://proceedings.mlsys.org/paper_files/paper/2019/file/acd593d2db87a799a8d3da5a860c028e-Paper.pdf.
  47. Bobbie Soedirgo. Compile and run LLVM IR in the browser, October 2023. original-date: 2021-02-24T14:29:16Z. URL: https://github.com/soedirgo/llvm-wasm.
  48. Bert Speelpenning. Compiling Fast Partial Derivatives of Functions Given by Algorithms. PhD thesis, University of Illinois at Urbana-Champaign, 1980. Copyright - Database copyright ProQuest LLC; ProQuest does not claim copyright in the individual underlying works; Last updated - 2023-02-19. URL: https://www.proquest.com/dissertations-theses/compiling-fast-partial-derivatives-functions/docview/302969224/se-2.
  49. Matthijs Vákár and Tom Smeding. CHAD: Combinatory homomorphic automatic differentiation. ACM Trans. Program. Lang. Syst., 44(3), August 2022. URL: https://doi.org/10.1145/3527634.
  50. Bret Victor. Explorable explanations, 2011. URL: https://worrydream.com/ExplorableExplanations/.
  51. Fei Wang, Daniel Zheng, James Decker, Xilun Wu, Grégory M. Essertel, and Tiark Rompf. Demystifying differentiable programming: Shift/reset the penultimate backpropagator. Proc. ACM Program. Lang., 3(ICFP), July 2019. URL: https://doi.org/10.1145/3341700.
  52. Guillermo Webster. g9: Automatically interactive graphics, 2016. URL: https://omrelli.ug/g9/.
  53. Yuting Yang, Connelly Barnes, Andrew Adams, and Adam Finkelstein. Aδ: Autodiff for discontinuous programs - Applied to shaders. ACM Trans. Graph., 41(4), July 2022. URL: https://doi.org/10.1145/3528223.3530125.
  54. Katherine Ye, Wode Ni, Max Krieger, Dor Ma'ayan, Jenna Wise, Jonathan Aldrich, Joshua Sunshine, and Keenan Crane. Penrose: From mathematical notation to beautiful diagrams. ACM Trans. Graph., 39(4), August 2020. URL: https://doi.org/10.1145/3386569.3392375.
Questions / Remarks / Feedback
X

Feedback for Dagstuhl Publishing


Thanks for your feedback!

Feedback submitted

Could not send message

Please try again later or send an E-mail