확률적 프로그래밍 컴파일러 임펄스: 베이즈 추론을 최대 190배 빠르게 만든 방법

핵심 메시지

• 베이즈 추론에서 과학자가 자연스럽게 쓰고 싶은 코드와 샘플러가 효율적으로 실행해야 할 코드는 서로 다르다는 구조적 불일치가 성능을 가로막는다.
• 기존 확률적 프로그래밍 프레임워크는 확률적 의미론을 컴파일러에 넘기기 전에 해소해버려, 컴파일러가 반복 계산의 낭비를 알아볼 수 없다.
• 임펄스는 MLIR 위에서 확률적 구조를 1급 컴파일러 프리미티브로 보존해, 기존 최적화와 확률 전용 최적화를 모두 자연스러운 코드에 적용한다.
• 핵심 기법인 '샘플 불변 코드 모션'은 샘플링 파라미터에 의존하지 않는 결정적 계산을 정적으로 찾아 루프 밖으로 빼낸다.
• 표준 컴파일러 패스만으로 8~60배, 확률 전용 최적화까지 더하면 190배 가속을 얻었고, 블랙홀 영상 복원 워크로드에서는 3.7배 빨라졌다.

쉽게 이해하기

베이즈 추론은 잡음 섞인 관측에서 숨겨진 값을 추정하는 도구다. 샘플러가 가능한 값들을 반복해서 탐색하며 답을 좁혀 가는 방식이라, 데이터가 적거나 간접적이거나 잡음이 많을수록 진가를 발휘한다. 확률적 프로그래밍 프레임워크는 과학자가 이 통계 모형을 코드로 표현하면 추론 알고리즘을 자동으로 돌려주는 역할을 맡는다. 일리노이대 어배너-섐페인의 브랜트 첸이 PLDI 2026 학생 연구 부문에서 발표한 임펄스(Impulse)는 바로 이 지점의 성능 문제를 겨냥한다.

문제는 과학자가 쓰고 싶은 코드와 샘플러가 실행해야 할 코드가 다르다는 데 있다. 과학 공식을 그대로 옮겨 적으면 프레임워크가 그 모형을 반복 호출하는 루프를 만들어주는데, 이때 값비싼 촐레스키 분해 같은 연산이 매 반복마다 다시 계산된다. 사람이 자세히 들여다보면 이 식을 '루프 불변 인자의 제곱근'과 '루프 불변 행렬의 촐레스키'의 곱으로 분해해 루프 밖으로 통째로 빼낼 수 있다. 그러나 오늘날의 프레임워크는 그 재작성을 사용자에게 떠넘긴다. 어디에 낭비가 있는지 찾아내는 전문가 수준의 분석은 물론, 내부 헬퍼 함수가 무엇을 하는지까지 이해해야 하기 때문에 실무에서는 매우 어렵다.

기존 프레임워크가 이 낭비를 스스로 고치지 못하는 이유는 추상화 계층에 있다. 확률적 의미론이 컴파일러에 도달하기 전에 이미 해소되어 사라지기 때문이다. 낮은 수준으로 내려온 코드에서 컴파일러는 어떤 연산이 촐레스키이고 어떤 연산이 그 기울기인지부터 알아내야 하는데, 발표자는 이를 감당할 수 없는 분석량이라고 표현했다. 그래서 연구진은 두 가지 질문을 던진다. 확률적 구조를 컴파일러 중간 표현에 보존할 수 있는가, 그리고 컴파일러가 자연스러운 코드에서 빠른 추론을 자동으로 끌어낼 수 있는가.

임펄스의 답은 확률적 구조를 1급 컴파일러 프리미티브로 만드는 것이다. 고수준 중간 표현에 확률적 의미론이 남아 있으면 두 가지가 가능해진다. 하나는 기존의 일반 최적화를 과학자가 자연스럽게 쓴 확률 코드에 그대로 적용하는 것이고, 다른 하나는 기존 프레임워크가 할 수 없던 확률 전용 최적화다. 대표적인 것이 샘플 불변 코드 모션이다. 컴파일러가 샘플 의존성 분석으로 어떤 계산이 샘플링된 파라미터에 의존하고 어떤 계산이 그렇지 않은지를 정적으로 가려낸 뒤, 반복될 필요가 없는 결정적 계산을 밖으로 꺼낸다. 흥미로운 점은 이 재작성이 제곱근 연산을 새로 추가한다는 것이다. 범용 텐서 컴파일러라면 연산이 늘어나므로 절대 하지 않을 선택이지만, 이 코드가 수백만 번 반복된다는 사실을 아는 컴파일러에게는 명백한 이득이다.

성능 평가에서는 고수준 중간 표현에 표준 컴파일러 패스를 돌리는 것만으로 8~60배 가속이 나왔고, 샘플 불변 코드 모션 같은 확률 전용 최적화까지 더하자 190배까지 올라갔다. PPL 벤치 같은 표준 벤치마크에서는 커뮤니티에서 널리 쓰이는 스탠(Stan)을 포함한 주요 프레임워크들보다 빨랐다. 실제 응용으로는 망원경 측정값에서 블랙홀 영상을 복원하는 천문학 워크로드에 적용했다. 임펄스는 여러 겹의 라이브러리 호출 뒤에 숨어 있던 푸리에 변환 연산을 하나의 미리 계산된 행렬 곱으로 접어 넣어 3.7배 가속을 얻었는데, 발표자는 이런 중복을 사람이 손으로 찾아내려면 천체 모형 라이브러리가 내부에서 어떤 변환 라이브러리를 부르는지까지 파악해야 해 사실상 불가능에 가깝다고 설명했다.

주요 인사이트

• 이 연구의 핵심은 최적화 기법 자체보다 책임의 이동에 있다. 성능 튜닝의 부담을 사용자에게 떠넘기던 구조를 컴파일러 쪽으로 옮기자는 주장이다.
• 범용 컴파일러가 절대 하지 않는 '연산을 늘리는 재작성'이 여기서는 이득이 된다. 코드가 수백만 번 반복된다는 도메인 지식이 최적화의 판단 기준 자체를 바꾼다.
• 소스 코드 수준의 변환으로는 이 문제를 풀기 어렵다. 최적화 대상이 라이브러리 호출 여러 겹 아래 숨어 있어, 라이브러리를 인라이닝하거나 소스를 직접 고쳐야 하기 때문이다.
• 질의응답에서 나온 지적처럼 '변하지 않는 부분을 루프 밖으로 빼는' 발상 자체는 일반적이지만, 이번 기법의 근거는 샘플링 파라미터가 반복마다 흔들린다는 추론 알고리즘의 의미론이라 결정적 신경망에 그대로 옮기기는 어렵다.
• 천문학 사례가 보여주듯 이런 최적화의 가치는 벤치마크 점수보다 GPU가 필요할 만큼 무거운 실제 과학 워크로드에서 크게 드러난다.

자주 묻는 질문

관련 AI 소식