Eric Jang의 튜토리얼을 번역한 게시물입니다.
이전 블로그 포스트에서는, 어떻게 정규화 플로우를 통해 가우시안 같은 간단한 분포를 “변형”하여 복잡한 데이터 분포에 피팅시킬 수 있는지 설명했습니다. 2차원 아핀 Bijector와 비선형 PReLU를 연쇄해서 간단한 플로우를 구현했고, 이로 자그마한 가역 신경망을 만들었습니다.
그렇지만 이 다층 퍼셉트론 플로우는 매우 빈약합니다: 은닉층마다 유닛 두 개 뿐이니까요. 게다가, 비선형 부분은 단조이며, 조각별로 선형적이라서, 원점에 대해서 데이터 매니폴드를 약간 왜곡해주는 것이 전부입니다. 이 플로우로는 등방성 가우시안을 두 개의 덩어리로 나누어 밑의 그림과 같이 “두 개의 달” 데이터셋을 나타내는 것과 같은 복잡한 변환은 절대로 할 수가 없습니다:
다행히 요새 머신러닝 연구를 통해 보다 강력한 정규화 플로우 기법들이 등장했습니다. 이 튜토리얼에서는 이러한 테크닉들을 살펴보겠습니다.
오토리그레시브 모델은 정규화 플로우
WaveNet과 PixelRNN과 같은 오토리그레시브 밀도 추정 테크닉은 복잡한 결합 밀도 p(x1:D)를 각 xi가 이전의 i-1값들에만 의존하는 일차원 조건부 밀도의 곱으로 나누어서 학습합니다:
p(x) = ∏ip(xi|x1:i-1)
조건부 밀도는 보통 학습 가능한 파라미터를 갖습니다. 예를 들어, 보통 오토리그레시브의 조건부 확률은 단변량 가우시안이죠; 그 평균과 표준 편차는 그 전까지의 x1:i-1를 가지고 신경망으로 계산합니다.
p(xi|x1:i-1) = N(xi|μi, (exp αi)2)
μi = fμi(x1:i-1)
αi = fαi(x1:i-1)
데이터를 오토리그레시브 밀도 추정으로 학습한다는 것은, 앞 쪽의 변수는 뒤 쪽의 변수에 의존하지 않는다는 상당히 대담한 귀납적 편향을 받아들이겠다는 뜻입니다. 직관적으로 생각해봐도 현실 세계의 데이터에서 이는 전혀 사실이 아니겠죠 (이미지의 위 행의 픽셀은 아래의 이미지에 대해서 인과적, 조건부 의존을 지닙니다). 하지만 이런 방식으로도 그럴듯한 이미지를 만들어 낼 수 있습니다(연구자들도 무척 놀랐습니다)!
이 분포에서 샘플하려면, D개의 “소음 변량” u1:D를 표준 정규 분포 N(0, 1)에서 계산하고, 다음 재귀를 적용해서 x1:D를 구합니다.
xi = ui exp αi + μi
ui ~ N(0, 1)
오토리그레시프 샘플링은 (N(0, I)에서 샘플된) 기저 소음 변량을 새로운 분포로 바꾸는 결정론적 변환 절차이므로, 오토리그레시브 샘플은 사실 표준 분포의 변환된 분포라고 해석할수가 있죠!
이걸 이해했으면, 다수의 오토리그레시브 변환을 쌓아서 정규화 플로우를 만들 수 있습니다. 이렇게 하는 이점은 플로우의 각 Bijector에 대해서 변수 x1, …, xD의 순서를 바꿀 수 있다는 것입니다. 그래서 한 오토리그래시브 분해가 (변수 순서가 효과적이지 않아서) 분포를 잘 모델링하지 못하더라도, 다음 레이어는 잘 할 수도 있는 것이죠.
Masked Autoregressive Flow(MAF) Bijector는 그런 조건부 가우시안 오토리그레시브 모델을 구현합니다. 다음은 변환된 분포의 샘플 하나인 xi에 대한 정방향 패스를 그린 도식입니다:
회색 유닛 xi는 계산하고자 하는 유닛이고, 파란 유닛들은 그 계산에 필요한 값들입니다. αi와 μi는 x1:i-1를 신경망에 넣어서 나온 스칼라 값입니다. 변환은 단순히 스케일과 시프트 이지만, 그 스케일과 시프트는 이전 변수에 대해 복잡한 의존성을 지닐 수 있습니다. 첫번째 유닛인 x1에 대해서 μ와 α는 보통 어떤 x나 u에도 의존하지 않는 학습가능한 스칼라 변수로 설정됩니다.
변환을 이런식으로 디자인하는 보다 중요한 이유는, 역함수 u = f-1(x)를 계산 할 때 fα 나 fμ의 역을 구할 필요가 없기 때문입니다. 변환이 스케일과 시프트로 파라미터라이즈되어서, 그 시프트와 스케일만 반대로 해주면 원래의 소음 변량을 회복할 수 있습니다: u = (x – fμ(x)) / exp(fα(x)). 이 Bijector의 정방향과 역방향 패스는 fμ(x)와 fα(x)의 정방향 평가에만 의존하기 때문에, 역을 취하는 것이 불가능한 ReLU같은 함수를 사용하거나, fα와 fμ 신경망에서 정사각형이 아닌 행렬곱을 사용할 수 있는 것이죠.
MAF 모델의 역방향 패스를 이용해서 밀도를 평가할 수 있습니다:
distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x))
런타임 복잡도와 MADE
오토리그레시브 모델과 MAF는, 현대 GPU의 배치 병렬 계산을 이용해서 모든 조건부 우도 p(x1), p(x2|x1), …, p(xD|x1:D-1)를 D 쓰레드의 단일 패스로 동시에 평가하여 “빠르게” 학습시킬 수 있습니다. CPU나 GPU에서의 SIMD 벡터라이제이션과 같은 병렬계산이 런타임 오버헤드를 갖지 않는다고 가정하겠습니다.
그와 반대로, 오토리그레시브 모델에서 샘플링 하는 것은 느린데, 이는 xi를 계산하기전 모든 이전 x1:i-1이 계산되기를 기다려야하기 때문입니다. 하나의 샘플을 생성하는 런타임 복잡도는 단일 쓰레드를 D번 시퀀셜 패스하는 것과 같죠, 병렬 프로세싱을 활용할 수 없습니다.
다른 이슈 하나 더: 병렬가능한 역방향 패스에서 각 αi와 μi를 구하기 위해서 (형태가 다른 인풋을 받는) 별도의 신경망을 사용해야 할까요? 이는 비효율적입니다. D개의 네트워크간 학습된 표상이 (오토리그레시브 의존성이 지켜지는 한) 공유되어야 한다는 점을 생각하면 특히나 더 그렇죠. Masked Autoencoder for Distribution Estimation (MADE) 논문의 저자들은 매우 멋진 해결책을 제시합니다: 하나의 신경망을 사용해서 동시에 α와 μ의 모든 값을 출력하되, 오토리그레시브한 성격이 지켜지도록 웨이트에 마스킹을 하는 것이죠.
이 기법은 신경망 하나에 대한 단 한 번의 패스를 통해 x의 모든 값에서 u의 모든 값을 수복할 수 있도록 해줍니다 (D 입력, D 출력). 이는 D개의 신경망을 동시에 처리하는 것 보다 (D(D+1) / 2 입력, D 출력) 훨씬 효율적이죠.
정리하자면 MAF는 MADE 구조를 사용해서 오토리그레시브 변환의 시프트-스케일 비선형 파라미터를 효율적으로 구하고, 이러한 효율적 오토리그레시브 모델을 정규화 플로우 프레임워크로 바꾸어줍니다.
역방향 오토리그레시브 플로우 (IAF)
역방향 오토리그레시브 플로에서는 이전 데이터 샘플 대신 이전 소음 변량 u1:i-1에서 비선형적 시프트/스케일 통계값을 계산합니다:
xi = ui exp αi + μi
μi = fμi(u1:i-1)
αi = fαi(u1:i-1)
IAF의 정방향 패스(샘플링)는 빠릅니다: D 쓰레드를 한 번에 병렬로 패스해서 모든 xi를 계산할 수 있으니까요. IAF도 마찬가지로 MADE 네트워크를 사용해서 이러한 병렬계산을 효율적으로 구현합니다.
하지만 새로운 데이터 포인트가 주어지고 그에 대한 밀도를 평가하려면 u를 수복해야 하는데, 이 프로세스가 느리죠: 우선 u1 = (x – μ1) * exp(-α1)을 회복하고, 그리고 차후 ui = (x – μi (u1:i-1)) * exp(-αi(u1:i-1))를 순차적으로 회복해야 합니다. 반면에, IAF로 생성한 샘플의 (로그) 확률을 추적하는 것은 매우 쉬운데, 이는 x의 역을 취할 필요가 없이 이미 모든 u 값을 알기 때문입니다.
날카로운 독자라면 아래 행을 x1, … xD 그리고 위 행을 u1, …, uD로 바꾸어주면, MAF Bijector 역방향 패스와 정확히 같다는 것을 눈치챘을 겁니다! 마찬가지로, IAF의 역은 (x와 u를 바꾼) MAF의 정방향 패스와 동일하죠. 그러므로, TensoFlow Distributions에서는, 사실 완전히 동일한 Bijector 클래스를 사용해 MAF와 IAF를 구현합니다, 그리고 편리한 “Invert” 기능을 이용해서 Bijector의 역을 취해 역방향 정방향 패스를 바꿀 수 있죠.
iaf_bijector = tfb.Invert(maf_bijector)
IAF와 MAF는 상반되는 계산 트레이드오프가 있습니다 – MAF는 학습은 빨리 되지만 샘플링은 느리고, IAF는 학습은 느리지만 샘플링이 빠릅니다. 신경망을 학습시키기 위해서는 보통 밀도 평가가 샘플링 보다는 중요하므로 MAF가 분포를 배우기에 더 적절한 선택입니다.
Parallel WaveNet
당연히 그럼 이런 의문이 생길 것입니다, 이 두 접근법을 합쳐서 장점만 가져올 수 없는지, 다시 말해 빠른 학습과 빠른 샘플링을 가져올 수 없는지.
그 답은 예스입니다! DeepMind에서 공개해 큰 화제가 된 Parallel WaveNet이 바로 이걸 하는거죠: 오토리그레시브 모델(MAF)로 생성모델을 효과적으로 학습한 후, IAF 모델을 사용해서 이 선생 모델에 대한 자신의 샘플의 우도를 최대화 하도록 학습합니다. IAF의 경우 (학습 셋의 데이터 포인트와 같은) 외부 데이터 포인트의 밀도를 계산하는 것이 매우 값비쌉니다, 하지만 자신의 샘플의 밀도를 계산하는 것은 매우 간단하죠. 노이즈 변량 u1:D을 캐싱해서 역방향 패스를 부르는 것을 피하면 됩니다. 그러므로 학생과 선생의 분포 사이의 발산을 최소화해서 “학생” IAF 모델을 학습할 수 있습니다.
이는 정규화 플로우 리서치를 통해 나온 믿을 수 없을 만큼 영향력이 큰 어플리케이션으로 – 그 결과 20배 더 빠르게 샘플할 수 있는 실시간 오디오 합성 모델이 탄생했는데요, 구글 어시스턴트와 같은 실제 서비스에서 이미 사용되고 있습니다.
NICE와 Real-NVP
끝으로 IAF Bijector의 특별 케이스라고 볼 수 있는 Real-NVP를 살펴봅시다.
NVP “커플링 레이어”에서는, 정수를 0 < d < D로 고정합니다. IAF처럼, xd+1은 이전 ud 값들에 의존하는 시프트와 스케일 변환 결과입니다. 차이점은 xd+2, xd+3, …, xD 또한 이 ud 값들에만 의존하도록 강제하여 한 번의 네트워크 패스로 αd+1:D와 μd+1:D를 만들어냅니다.
x1:d의 경우 u1:d와 동일하게 설정되는 “패스-쓰루” 유닛입니다. 그러므로 Real-NVP 또한 MAF Bijector의 특이 케이스이죠 (α(u1:d)=α(x1:d)이기 때문입니다).
전체 레이어에 대한 시프트와 스케일 통계값을 x1:d이나 u1:d에서 한 번의 패스를 통해 계산할 수 있기 때문에, NVP는 정방향과 역방향 계산을 한 번의 병렬 패스에서 수행할 수 있습니다 (샘플링과 추정이 둘 모두 빠릅니다). MADE가 필요치도 않죠.
하지만 경험적 연구에 비춰보면 Real-NVP는 MAF나 IAF에 비해서 성능이 떨어지는 경향이 있습니다. 제 경험으로는 같은 수의 레이어를 이용했을 때, Real-NVP가 이차원 토이 데이터셋(예. SIGGRAPH 데이터셋)에 더 피팅을 못하는 경향이 있었습니다. 이차원의 경우 Real-NVP와 IAF가 거의 동일합니다만, 한 가지 다른 점은 IAF는 스케일과 시프트를 통해 변환되는데 반해 Real-NVP는 첫번째 유닛을 바꾸지 않고 그대로 둔다는 점입니다..
NICE bijector는 α = 0을 가정하여 시프트만 사용하는데요, Real-NVP는 그 후속 연구였습니다. NICE는 분포를 스케일 하지 않기 때문에 ILDJ가 상수이죠!
배치 노말라이제이션 Bijector
Real-NVP 논문은 몇 가지 새로운 제시를 하는데, 그 중 하나가 배치 정규화 Bijector를 사용해서 학습을 안정화시키는 것입니다. 보통은 배치 정규화를 신경망 학습에 적용합니다. 정방향 통계값이 평균으로 중심이 옮겨지고 대각 단위 공분산에 따라 스케일 되며, 배치 정규화 통계값(러닝 평균, 러닝 분산)은 지수적 이동 평균을 통해 축적되는 것이죠. 테스트 시에는 축적된 통계값이 데이터를 정규화하는데 사용됩니다.
그런데 정규화 플로우의 경우, 학습 동안 배치 정규화가 bijector.inverse에서 사용되고, 축적된 통계는 “테스트 시”에 데이터를 역-정규화 하기 위해서 사용됩니다 (bijector.forward). 구체적으로 배치 정규화 Bijector는 보통 다음처럼 구현됩니다.
역방향 패스:
- 데이터 분포 x의 현재 평균과 표준 편차를 계산합니다.
- 러닝 평균과 표준 편차를 업데이트 합니다.
- 현재 평균/표준 편차를 사용해서 배치 노말라이제이션을 합니다.
정방향 패스:
- 러닝 평균과 표준 편차를 사용해서 데이터 분포를 비정규화합니다.
TF Bijectors를 사용하면 몇 줄 되지않는 코드로 이를 구현할 수 있습니다.
class BatchNorm(tfb.Bijector):
def __init__(self, eps=1e-5, decay=0.95, validate_args=False, name="batch_norm"):
super(BatchNorm, self).__init__(
event_ndims=1, validate_args=validate_args, name=name)
self._vars_created = False
self.eps = eps
self.decay = decay
def _create_vars(self, x):
n = x.get_shape().as_list()[1]
with tf.variable_scope(self.name):
self.beta = tf.get_variable('beta', [1, n], dtype=DTYPE)
self.gamma = tf.get_variable('gamma', [1, n], dtype=DTYPE)
self.train_m = tf.get_variable(
'mean', [1, n], dtype=DTYPE, trainable=False)
self.train_v = tf.get_variable(
'var', [1, n], dtype=DTYPE, initializer=tf.ones_initializer, trainable=False)
self._vars_created = True
def _forward(self, u):
if not self._vars_created:
self._create_vars(u)
return (u - self.beta) * tf.exp(-self.gamma) * tf.sqrt(self.train_v + self.eps) + self.train_m
def _inverse(self, x):
# Eq 22. Called during training of a normalizing flow.
if not self._vars_created:
self._create_vars(x)
# statistics of current minibatch
m, v = tf.nn.moments(x, axes=[0], keep_dims=True)
# update train statistics via exponential moving average
update_train_m = tf.assign_sub(
self.train_m, self.decay * (self.train_m - m))
update_train_v = tf.assign_sub(
self.train_v, self.decay * (self.train_v - v))
# normalize using current minibatch statistics, followed by BN scale and shift
with tf.control_dependencies([update_train_m, update_train_v]):
return (x - m) * 1. / tf.sqrt(v + self.eps) * tf.exp(self.gamma) + self.beta
def _inverse_log_det_jacobian(self, x):
# at training time, the log_det_jacobian is computed from statistics of the
# current minibatch.
if not self._vars_created:
self._create_vars(x)
_, v = tf.nn.moments(x, axes=[0], keep_dims=True)
abs_log_det_J_inv = tf.reduce_sum(
self.gamma - .5 * tf.log(v + self.eps))
return abs_log_det_J_inv
간단하게 역함수의 로그 미분을 취해서 ILDJ를 유도할 수 있습니다 (단변 케이스를 생각해보세요).
코드 예시
Google Bayesflow팀과 Josh Dillon 덕분에, 이미 MADE 신경망을 사용해서 효율적으로 u를 수복하는 MaskedAutoregressive Flow Bijector가 구현되어 있죠.
우선 이 블렌더 스크립트를 사용해서 점을 “SIGGRAPH”라는 글자 모양으로 흩뿌린 복잡한 2차원 분포를 만들었습니다. 이전 튜토리얼과 매우 비슷하게 데이터셋, Bijector, 변환분포를 생성할테니 이에 대한 코드를 여기에 또 올리진 않을게요 – Jupyter Notebook은 여기서 참고하시면 됩니다. 이 Notebook으로 “두 개의 달”과 “SIGGRAPH” 데이터 분포에 대해서 MAF, IAF, Real-NVP를 사용하여 정규화 플로우를 학습할 수 있습니다.
잊기 쉬운 중요한 디테일 하나는 각 플로우에서 변수의 순서를 치환해주지 않으면 모델이 전혀 동작을 하지 않는다는 것입니다. 그렇지 않으면 그 어떤 레이어의 오토리그레시브 분해도 p(x1|x2) 구조를 배울 수가 없죠. 다행히 TensorFlow의 치환 Bijector가 바로 이러한 기능을 수행합니다.
for i in range(num_bijectors):
bijectors.append(tfb.MaskedAutoregressiveFlow(
shift_and_log_scale_fn=tfb.masked_autoregressive_default_template(
hidden_layers=[512, 512])))
bijectors.append(tfb.Permute(permutation=[1, 0]))
flow_bijector = tfb.Chain(list(reversed(bijectors[:-1])))
최종 경과물과 학습된 플로우입니다. 태피 반죽기가 생각나네요.
토의
TensorFlow distribution을 통해서 가독성 높고 깔끔한 코드로 연쇄 속 자코비안의 행렬식을 자동적으로 축적하여 정규화 플로우를 쉽게 구현할 수 있습니다. 어떤 정규화 플로우를 사용할지 결정할 때는 빠른 정방향 패스과 빠른 역방향 패스 사이, 그리고 플로우의 표현력과 ILJD 계산의 속도 사이에 트레이드오프가 있다는 점을 고려하세요.
튜토리얼 파트 1에서 정규화 플로우가 강화학습과 생성 모델에 사용할 수 있는 보다 강력한 분포를 만들어낼 수 있는 한 방법이라고 소개드렸습니다. 하지만 더 큰 그림을 보자면 부피 추적 정규화 플로우가 로보틱, 구조적 예측과 같은 AI 어플리케이션에 있어서 최고의 도구인지는 아직 확실하지 않습니다. 변분 추론이나 암시적 밀도 모델은 벌써 엄청난 성능을 보여주는데 말이죠. 그럼에도 정규화 플로우는 알아둘만한 유용한 방법론이고 구글 어시스턴트에 보급된 실시간 오디오 생성 모델과 같이 실제 어플리케이션에 사용될만큼의 뛰어난 성능도 증명했습니다.
정규화 플로우와 같은 명시적 밀도 모델은 최대 우도를 통해서 학습하기에 적합하지만, VAE나 GAN의 일부로 사용할 수도 있습니다. VAE 사전확률이나 GAN의 잠재 코드처럼 가우시안 분포를 사용하는 어느 곳에도 정규화 플로우를 대체재로 활용할 수 있는 것이죠. 예를 들어서, 이 논문은 정규화 플로우를 유연한 변분 사전확률로 사용하고, Tensorflow distributions 논문은 정규화 플로우를 사전확률로 삼아 PixelCNN 디코더와 함께 사용하는 VAE를 소개합니다. Parallel WaveNet은 KL 발산을 통해서 IAF “학생” 모델을 학습하고요.
정규화 플로우의 가장 흥미로운 특성 중 하나는 가역적인 계산을 구현한다는 것입니다 (다시 말해, 강한 표현력의 함수가 정의된 역함수를 가진다는 것이죠). 이는 정방향 패스 때 액티베이션 값을 메모리에 저장할 필요없이, 역방향 패스 때 다시 계산해서 (큰 그래프에 대해서는 계산이 값비쌀 가능성도 있겠죠) 역전파 할 수 있다는 뜻입니다. 신용 배분이 상당히 긴 시간동안 일어나는 상황에서 가역 계산을 이용해서 메모리 사용량을 제한하면서도 과거 결정 상태를 “수복”할 수 있죠. 사실 이 아이디어는 NICE bijector의 가역성으로부터 기원했고, RevNets 논문에서 활용되었습니다. 메멘토라는 영화의 주인공이 생각나네요. 기억을 저장할 수 없어 가역 계산을 사용하여 중요한 정보를 잊지 않도록 하죠.
읽어주신 독자분들 감사합니다.
감사의 말
이 포스트를 검수해준 Dustin Tran, Luke Metz, Jonathan Shen, Katherine Lee, Samy Bengio에게 감사를 전합니다.
참조와 추가 문헌
- 이 블로그 포스트의 구조와 내용은 Masked Autoregressive Flow for Density Estimation 논문의 영향을 많이 받았습니다. 무척이나 잘 쓴 논문이어서 제가 이 토픽을 이해하는데 큰 도움이 됐죠. 한 번 읽어보세요!
- NF에 대한 이전 연구들입니다: https://math.nyu.edu/faculty/tabak/publications/Tabak-Turner.pdf 와 https://arxiv.org/pdf/1302.5125.pdf 와 https://arxiv.org/abs/1505.05770
- Laurent Dinh의 & Twitter Cortex 리서처들의 토의 내용입니다. 재밌는 아이디어가 많죠.
- PyMC를 사용하는 정규화 플로우 튜토리얼
- 정규화 플로우를 랑게빈 플로우와 해밀토니안 플로우로 확장하는 연구는 아직 제가 완벽히 이해하지는 못했습니다. 정규화 플로우의 Bijector의 수가 무한으로 늘어나면, 훨씬 풍부한 변환을 표현할 수 있는 연속-시간 플로우가 됩니다.
Studies 