Online Reducible Norm - Welford 기반 Streaming Statistics

Method Overview

많은 통계 연산은 평균과 분산을 계산하기 위해 두 번 이상의 데이터 순회를 요구한다.

예를 들어 일반적인 variance 계산은 다음과 같이 수행된다.

mean = sum(x) / N
var  = sum((x-mean)^2) / N

이 방식은 평균을 먼저 계산한 뒤 다시 데이터를 순회해야 하기 때문에 multi-pass reduction 이 된다.

하지만 Welford 알고리즘을 사용하면 평균과 분산을 single-pass streaming 방식으로 계산할 수 있다.

핵심 아이디어는 다음 세 가지 상태를 유지하는 것

count
running_mean
running_M2

이 상태는 새로운 데이터가 들어올 때마다 업데이트되며, 전체 데이터를 다시 보지 않아도 통계값을 계산할 수 있다.

결과적으로

• multi-pass reductino -> single-pass streaming reduction
• intermediate tensor 저장 불필요
• large batch statistics 에서도 HBM traffic 감소

 

Mathematical Structure

Welford 업데이트는 다음과 같다.

새로운 값 x 가 들어올 때

count += 1
delta = x - mean
mean += delta / count
delta2 = x - mean
M2 += delta * delta2

마지막에

variance = M2 / count

중요한 수학적 성질은 mergeability 이다

두 partial 통계 (mean_a, M2_a, count_a) 와 (mean_b, M2_b, count_b) 는 다음처럼 함칠 수 있다.

delta = mean_b - mean_a
count = count_a + count_b
mean  = mean_a + delta * (count_b / count)
M2    = M2_a + M2_b + delta^2 * count_a * count_b / count

이 성질 때문에

• streaming
• parallel reduction
• disstributed merge

가 모두 가능하다.

 

Hardware Implication

이 구조는 하드웨어 관점에서 중요한 효과를 가진다.

기존 방식

pass 1 : sum(x)
pass 2 : sum((x-mean)^2)

• 최소 2 번 global memory read

Welford 방식

Single streaming pass

• global memory read 1회

또한 다음이 가능하다

• warp-local reduction
• block-level partial stats
• hierarchical merge

이는 GPU reduction kernel 에 매우 적합한 구조이다.

대표 적용 사례

• LayerNorm
• BatchNorm statistics
• dataset statistic
• streaming metrics

 

Compiler Recognition (MCIR)

컴파일러는 다음 패턴을 인식해야 한다.

mean + variance calculation

그리고 이를 다음 property 로 추상화할 수 있다.

Property :  online_reducible_norm

Legality 조건

- statistic state 가 mergeable
- update rule 이 associative merge 지원

Rewrite

two-pass statistics
-> streaming statistics kernel

Kernel mapping

global load -> local accumulator
warp / block reduction
-> final merge

이렇게 되면 컴파일러는 LayerNorm 이나 dataset statistics 계산을 single-pass streaming kernel 로 자동 변환할 수 있다.