(arXiv:1902.08153) Learned Step Size Quantization 리뷰

Esser, Steven K., et al. “Learned step size quantization.” arXiv preprint arXiv:1902.08153 (2019).

도입

저비트 네트워크는 고비트 네트워크의 가중치로 학습할 수 있음이 계속해서 확인되고 있다. 이 과정에서의 양자화는 주로 실수를 이산 값으로 매핑하는 균일 양자화(uniform quantization)을 사용한다.

균일 양자화는 일반적으로 전체 실수 구간을 동일한 구간 크기(step size)로 나누어 각 구간별 대표값을 설정하는 방식을 사용한다.

이후에 (Polino et al., 2018)이 비균일한 구간을 양자화에 사용하거나, (Lioo & Liu, 2016; Zhou et al., 2016; Cai et al. 2018; McKinstry et al., 2018)이 데이터 분포에 따라 양자화 구간을 조정하거나, (Choi et al., 2018; Zhang et al. 2018)이 학습을 이용하여 양자화 오류를 최소화하려는 시도가 있기도 했다. 하지만 항상 네트워크의 성능을 최적화함 혹은 양자화 오류를 최소화함을 보장하지는 못했다.

그래서 저자들은 Learned Step Size Quantization(LSQ) 를 제안하게 되었다.

Learned Step Size Quantization (LSQ)

LSQ는 양자화 대상이 되는 각 레이어에 양자화 step size $s$를 학습 가능한 파라미터로 두고, 손실을 직접 최소화하는 방향으로 $s$를 학습하도록 했다.

LSQ는 두 가지 면에서 향상이 있었다.

1. round 연산에 대해 Straight-Through Estimator(STE)를 사용해 단순한 방법으로 step size $s$로의 그래디언트를 양자화 상태 전이점(transition point)에 민감하도록 근사한다.
2. $s$ 업데이트 크기가 가중치 업데이트와 비슷하게 만들도록 하여 가중치의 수렴 성능을 개선한다.

이러한 방법은 활성 함수, 가중치 모두 양자화하는데 사용할 수 있었고, 역전파나 스토캐스틱 경사 하강법과 같은 기존의 학습 방법과 함께 사용할 수 있었다.

방법

양자화기는 원본 데이터 $v$, 스텝 사이즈 $s$, 상/하한의 경계 $Q_P$, $Q_N$를 이용해 양자화 결과 $\bar{v}$를 계산한다.

\[\bar{v}=\lfloor \text{clip}(v/s, -Q_N, Q_P)\rceil\] \[\hat{v} = \bar{v} \times s\]

구체적으로는 먼저 $v$를 $s$로 나눈 후 $-Q_N$, $Q_P$로 클리핑하고, 반올림한다. $\bar{v}$를 얻는 과정에서 $s$에 대해 나눈 만큼 값이 축소되었으므로 $s$를 다시 곱한다.

$s$를 나누고 다시 곱하는 과정에서 클리핑과 반올림이 발생하였으므로, $\hat{v}$는 특정한 유형의 값만을 가지게 된다.

부호 없는 값은 $Q_N = 0$; $Q_P=2^b-1$, 부호 있는 값은 $Q_N = -2^{b-1}$; $Q_P=2^{b-1}-1$로 설정한다.

스텝 사이즈 그래디언트

양자화기의 양자화 스텝 사이즈 $s$는 학습을 통해 조정된다. $s$의 학습 loss는 아래와 같이 정의하여 사용한다. 이 정의는 (Bengio et al., 2013)의 STE에서 유도되었다.

\[\frac{\partial \hat{v}}{\partial s} = \begin{cases} - v / s + \lfloor v / s \rceil & \text{if}\ -Q_N < v/s < Q_P \\ -Q_N & \text{if}\ v/s \le -Q_N \\ Q_P & \text{if}\ v/s \ge Q_P \\ \end{cases}\]

입력이 클립 범위 내 $(Q_N < v / s < Q_P)$라면 $\delta \hat{v} / \delta {s} = -v / s + \lfloor v / s \rceil$, 좌측으로 포화라면 $\delta \hat{v} / \delta {s} = -Q_N$, 우측으로 포화라면 $\delta \hat{v} / \delta {s} = Q_P$이다.

이 식은 입력이 전이 지점(transition point)에 가까울수록, $s$가 조금만 변해도 양자화 구간(quantization bin) $\bar{v}$가 바뀌어 $\hat{v}$이 크게 변한다. 이 상황에서 $\delta \hat{v} / \delta {s}$는 더 커지고, LSQ 그래디언트는 더 민감해진다.

연구에서는 LSQ의 비교군으로서 (Jung et al., 2018)의 QIL, (Choi et al., 2018)의 PACT가, $v$가 전이 지점에 얼마나 가까운지가 그래디언트에 영향을 주지 않는다고 언급하였다.

스텝 사이즈 그래디언트의 크기

(You et al., 2017)에 의해 네트워크의 평균적인 가중치 업데이트 규모가 평균적인 매개변수의 규모와 비슷해야 수렴이 잘 되었음이 밝혀졌다. 다시 말해 전이점에 민감한 $\delta \hat{v} / \delta {s}$는 값이 지나치게 확대되어 수렴을 해칠 수 있다.

\[\begin{align*} R &= \frac{\text{update size for step-size }s\text{ per current step-size }s}{\text{update size for weight }w\text{ per current weight }w} \\ &= \frac{\nabla_s L}{s} / \frac{\lvert\lvert \nabla_w L \rvert\rvert}{\lvert\lvert w \rvert\rvert} \end{align*}\]

그래서 LSQ에서도 “평균 업데이트 크기”와 “평균 파라미터 크기”의 비가 대략적으로 비슷하게 되도록, $R = (\nabla_s L / s) / (\lvert\lvert \nabla_w L \rvert\rvert / \lvert\lvert w \rvert\rvert)$가 평균적으로 1에 가깝도록 크기를 조정한다.

구체적으로는 레이어의 가중치 수 $N_W$, 활성 함수의 feature 수 $N_F$에 대하여, 가중치 step size에는 $g = 1/\sqrt{N_W Q_P}$를, 활성 함수 step size에는 $g = 1/\sqrt{N_F Q_P}$를 사용한다.

학습

모델을 학습하는데는 일반적으로 널리 사용되는 (Courbariaux et al., 2015)의 방법을 따랐다.

• FP32 마스터 가중치 저장/업데이트
• 순/역전파에는 양자화된 가중치, 활성 함수 사용
• round 미분은 STE로 처리
• SGD로 학습 가능 파라미터 업데이트

step size $s$의 초기값은 $v$, $Q_P$에 대해서, $2 \langle \lvert v \rvert \rangle / \sqrt{Q_P}$을 초기 가중치 분포, 첫 배치의 활성 함수 양자화에서 사용했다.

마무리

LSQ는 이전의 네트워크 양자화 방법들을 상회하는 성능을 보였다. 레이어의 크기와 정밀도에 따라서 양자화 스텝 사이즈의 손실 그래디언트를 조정하는 것이 성능에 효과적인 개선을 주었다.

LSQ는 MSE/MAE 등의 단순 양자화 오류를 직접 최소화하지는 않았다. 모델의 과업(task)에 대한 손실을 사용해 $s$를 학습하여서, 단순 오류 지표와 정확도 사이의 높은 상관관계는 관찰되지 않았다.

단순 오류가 더 크더라도 모델은 덜 민감한 부분으로 노이즈를 배분해 정확도를 더 잘 유지하거나 향상시킬 수 있었다.