[논문 리뷰] Differentiating a Tensor Language

Gilbert Bernstein, Michael W. Mara외 3명|arXiv (Cornell University)|2020. 08. 25.

Parallel Computing and Optimization Techniques참고 문헌 29인용 수 5

한 줄 요약

이 논문은 아이버슨의 지표 함수를 사용하여 희소성을 명시적으로 모델링함으로써, 텐서 프로그램의 역방향 미분을 순수 기능적으로, 증명 가능하게 효율적으로 계산하는 방법을 제시한다. 새로운 '텐서 SSA' 정규형과 내적 기반 자동 미분의 유도를 활용하여, 변형 없이도 저비용 기울기 원칙을 달성함으로써, 도함수 코드가 원래 함수만큼 점근적으로 효율적임을 보장한다.

ABSTRACT

How does one compile derivatives of tensor programs, such that the resulting code is purely functional (hence easier to optimize and parallelize) and provably efficient relative to the original program? We show that naively differentiating tensor code---as done in popular systems like Tensorflow and PyTorch---can cause asymptotic slowdowns in pathological cases, violating the Cheap Gradients Principle. However, all existing automatic differentiation methods that guarantee this principle (for variable size data) do so by relying on += mutation through aliases/pointers---which complicates downstream optimization. We provide the first purely functional, provably efficient, adjoint/reverse-mode derivatives of array/tensor code by explicitly accounting for sparsity. We do this by focusing on the indicator function from Iverson's APL. We also introduce a new "Tensor SSA" normal form and a new derivation of reverse-mode automatic differentiation based on the universal property of inner-products.

연구 동기 및 목표

• 순수 자동 미분의 비효율성으로 인해 기울기 계산에서 점근적 속도 저하가 발생할 수 있는 텐서 프로그램의 문제를 해결하기 위해.
• 덧셈 할당(+=)과 같은 명령형 측면 효과를 피하여 데이터 병렬성 유지 및 후속 최적화 가능하게 하기 위해.
• 밀도 및 희소 텐서 모두에 대해 순수 기능적 환경에서 저비용 기울기 원칙을 수학적으로 보장하기 위해.
• 효율적이고 정확한 도함수 컴파일을 가능하게 하는 새로운 정규형인 '텐서 SSA'를 도입하기 위해.
• 내적의 보편 성질을 기반으로 한 역방향 미분을 위한 수학적으로 엄밀한 기초를 제공하기 위해.

제안 방법

• 논문은 아이버슨의 지표 함수를 사용하여 불리언 조건을 표현하고, 텐서 연산에서의 희소성을 명시적으로 모델링한다.
• 정적 분석 및 텐서 프로그램 최적화를 가능하게 하는 새로운 '텐서 SSA' 정규형을 도입한다.
• 도함수 계산은 내적의 보편 성질을 기반으로 유도되며, 정확성과 효율성을 보장한다.
• 도함수 계산에서 희소성을 고려하여 불필요한 메모리 왤쓰를 방지한다.
• 점근적 효율성을 유지하는 대수적 단순화를 적용하며, 극단적인 경우에도 성능 저하가 발생하지 않는다.
• 희소성을 일급 시민으로 간주함으로써 += 할당을 피하고 순수 기능적 컴파일을 가능하게 한다.

실험 결과

연구 질문

• RQ1 명령형 측면 효과에 의존하지 않고 순수 기능적 텐서 언어에서 저비용 기울기 원칙을 달성할 수 있는가?
• RQ2 텐서 연산의 희소성은 어떻게 수학적으로 모델링할 수 있으며, 이로 인해 도함수 계산의 효율성이 보장될 수 있는가?
• RQ3 어떤 정규형이 데이터 병렬 텐서 프로그램에서 효율적이고 정확한 역방향 자동 미분을 가능하게 하는가?
• RQ4 내적의 보편 성질을 사용하여 증명 가능한 효율성의 자동 미분 시스템을 도출할 수 있는가?
• RQ5 기존 프레임워크(예: PyTorch, TensorFlow)에서 발생하는 점근적 비효율성은 어떻게 체계적으로 피할 수 있는가?

주요 결과

• 제안된 방법은 순수 기능적 환경에서 저비용 기울기 원칙을 달성하여, 도함수 계산이 원래 프로그램보다 최대 상수 배수 이내로 느려지지 않음을 보장한다.
• 아이버슨의 지표 함수를 통한 희소성 모델링으로 인해, 예를 들어 대각선 및 추적 연산에서 불필요한 메모리 왤쓰를 피할 수 있다.
• 내적 및 합산 조정을 고려하더라도, 정방향 및 역방향 계산 간 비용 동등성을 증명할 수 있다.
• 텐서 SSA 정규형은 정적 분석 및 최적화를 가능하게 하며, 데이터 병렬성을 유지하고 효율적인 컴파일을 가능하게 한다.
• 극단적인 경우에서도 기존의 간단한 미분 방식보다 성능이 뛰어나다. 예를 들어, trace(A) + ... + trace(A)의 기울기는 최적화 시 O(kN²)에서 O(kN)으로 감소한다.
• += 할당에 의존하지 않기 때문에 기능적 순수성을 유지하고 더 강력한 컴파일러 최적화를 가능하게 한다.