https://arxiv.org/pdf/1810.07951.pdf
Don't Unroll Adjoint: Differentiating SSA-form Programs
Michael J Innes, 2019

https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf
Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients
William S. Moses and Valentin Churavy, NeurIPS 2020

Обе статьи посвящены (обратному aka reverse-mode) автоматическому (или алгоритмическому) дифференцированию функций, представленных в форме Single Static Assignment aka SSA. И тем не менее, они описывают существенно различные подходы.

Первая статья даёт краткое введение в обратное дифференцирование и распространённый подход на основе Wengert Lists. Чтобы перейти к SSA форме, к Wengert Lists необходимо добавить метки, условные и безусловные переходы и фи-узлы (φ nodes). Соответственно, статья вводит правила дифференцирования этих управляющих конструкций (control flow constructs). Дополнительно вводятся правила дифференцирования для чтения и записи в ячейки памяти, поскольку основной прицел статьи — императивные языки (и Julia в особенности). Забавно, что на практике (на текущий момент) основанная на описанном подходе библиотека Zygote не поддерживает деструктивную модификацию массивов, несмотря на её (библиотеки) широкое использование, в особенности во фреймворке для машинного обучения Flux. 😊

Несмотря на использование SSA-формы, первая статья подразумевает сравнительно высокоуровневое представление, близкое к исходному языку, до проведения оптимизаций. Вторая же статья рассматривает внедрение автоматического дифференцирования непосредственно в фреймворк LLVM в виде одного из проходов компиляции, выполняемого над низкоуровневым SSA-представлением, не зависящим от исходного языка и прошедшего ряд оптимизаций. Поэтому основное внимание она уделяет низкоуровневым аспектам: теневой памяти (shadow memory), кешам, обработке указателей, в том числе — вызовам функций по указателю, и переиспользованию информации с других проходов, таких как type-based alias analysis.

Стремление проводить автоматическое дифференцирование настолько низкоуровневого представления продиктовано двумя соображениями. Во-первых, немедленная применимость к большому количеству промышленных языков — C, C++, Rust, Julia — без каких-либо изменений в самом языке. Во-вторых, оптимизация исходного кода может сильно упростить и ускорить порождаемый код расчёта градиента функции, в некоторых случаях — понизить сложность с квадратичной до линейной после применения loop-invariant code motion к исходному коду.

Для подтверждения ускорения, авторы провели замеры производительности и сравнения с традиционными подходами на задачах ADBench от Microsoft и нескольких сторонних реализациях численного решателя дифференциальных уравнений. Результаты и графики приведены в статье. 😊

В любом случае, обе работы полагаются на "классические компиляторные техники", такие как dataflow analysis, alias analysis, abstract interpretation, и оптимизации. И потому представляют собой интереснейшее расширение "поля деятельности компиляторщиков" в сравнительно новую, но стремительно набирающую популярность, область.