이전까지의 생성모델은 모두 배우려는 이미지들의 분포를 단순한 분포로 변환시키는 함수를 찾는 형태로 발전되어왔습니다. 이번 장은 에너지 기반 모델로, 물리 시스템 모델링(물리학에서 분자 등의 운동을 연구하기 위해 연구한 함수, 식 등)을 이미지 생성 분포에 활용한 내용을 다룹니다.
따라서 이전까지의 방법과 다른 방식을 적용한 것이 EBM의 차이입니다.
1. 사전 개념
볼츠만 분포
- 시스템이 해당 상태의 에너지와 온도의 함수로 특정 상태 에 있을 확률을 제공 하는 확률 분포 또는 확률 척도로 이 분포를 이미지 생성의 분포에 적용합니다.
- E(x)는 샘플 x의 에너지 함수(또는 점수)로 신경망이 학습시킬 대상입니다. 적분으로 표현되어서 복잡해보이지만, 마치 SoftMax 함수와 매우 유사해보입니다.
- 직관적으로 해석해보면, 가능성이 높은 샘플에는 낮은 점수를 출력하도록 하는 함수를 만든다고 보면 됩니다.(그러면 p(x)가 1에 가까워지고 반대는 0에 가까워짐)
스위시 활성화 함수
- reLU함수의 대안으로 위의 식을 따르고, 아래 그래프 모양을 따른다. ReLU함수와 모양이 비슷하지만 함수가 연속적이고 음수일 때 gradient descent로 0을 출력하지 않는 다는 점이 차이가 있다. 급격하게 x가 변한다고 해서 기울기가 변하지 않는다는 특징이 에너지 특성을 잘 뽑아내는데 도움이 되어 이 활성화 함수를 사용한다.
랑주뱅 동역학
분자 시스템에 대한 움직임을 수학적으로 모델링하기 위한 공식으로, 어떠한 시스템 안에서 시스템을 구성하는 요소가 움직이는 방향을 예측하는 모델을 연구하는 학문입니다.
주로 분자는 에너지가 높은 곳에서 낮은 곳으로 움직이기 때문에 이 공식을 활용하면 에너지가 낮은 쪽으로 향하도록 만들 수 있다. 앞서 볼츠만 분포의 공식에서 E(x)가 작을수록 p(x) == 1에 가까워진다.
우리가 이미지를 생성하려면 E(x)가 낮은 이미지(에너지가 낮은 이미지)를 만들 수 있어야 해서 이 랑주뱅 동역학이 사용된다.
랑주뱅 동역학에서는 t시점에서의 x라는 에너지 시스템이 변화하는 움직임을 위와 같이 정의했다. 이 때 gradient는 에너지가 감소하는 방향을 가리킨다.
t시점에서 t+1시점의 변화는 아래와 같이 정리해볼 수 있다.
여기서 gradient앞에 있는 계수는 learning_rate라고 보면 되고, w는 noise이다.
2. 해결하려는 문제
a. 이미지 생성의 어려움
새로운 샘플을 생성하려면 모델을 어떻게 사용해야하는지 명확하지가 않습니다. 어떤 샘플을 만들면 점수가 나오긴 하지만, 점수가 낮은 이미지(p(x) == 1)를 출력하려면 어떤 것을 넣어야할지 모르기 때문입니다.
b. 분모의 적분이 어려운 문제
컴퓨터로 이 모든 것을 연산하기 위해서는 현재의 볼츠만 분포 함수에 대해 약간의 근사를 가합니다.
3. 적용방법
전체 구조
전체 개념은 볼츠만 분포의 적분 계산을 쉽게 하기 위해서 어떠한 근사를 적용했냐입니다.
1. Energy Function 신경망 구조
이미지 데이터이기 때문에 신경망 구조는 activation function으로 swish 사용한 것을 제외하고는 다른 모델과 비슷하게 큰 차이는 없습니다.
ebm_input = layers.Input(shape = (32,32,1))
x = layers.Conv2D(16, kernel_size = 5, strides = 2,
padding = "same", activation = activations.swish)(x)
x = layers.Conv2D(32, kernel_size = 3, strides = 2,
padding = "same", activation = activations.swish)(x)
x = layers.Conv2D(64, kernel_size = 3, strides = 2,
padding = "same", activation = activations.swish)(x)
x = layers.Conv2D(64, kernel_size = 3, strides = 2,
padding = "same", activation = activations.swish)(x)
x = layers.Flatten()(x)
x = laers.Dense(64, activation = activations.swish)(x)
ebm_output = layers.Dense(1)(x) #마지막 층은 선형 활성화 함수를 가진 하나의 완전 연결 유닛입니다.
model = models.Model(ebm_input, ebm_output)
2. 랑주뱅 동역학 사용하기
에너지 함수는 위에서 말한 것처럼 주어진 입력(이미지)에 대해 하나의 점수를 출력합니다. 하지만 우리가 궁금한건 에너지 점수(output)가 낮은 새로운 샘플(input)을 생성하고 싶다는 점입니다.
이 문제에서 랑주뱅 동역학이 사용됩니다. 먼저 샘플 공간에서 임의의 지점을 뽑습니다. 그 다음부터 랑주뱅 동역학 공식을 사용한다면, 이 지점은 에너지가 낮은 방향(진짜 이미지가 생성되는 방향)으로 이동할 것입니다.
실제 코드는 아래와 같이 짜면 된다.
def generate_samples(model, inp_imgs, steps, step_size, noise):
imgs_per_step = []
for _ in range(steps):
#이미지에 작은 양의 잡음을 추가
inp_imgs += tf.random.normal(inp_imgs.shape, mean = 0, stddev = noise)
inp_imgs = tf.clip_by_value(inp_imgs, -1.0, 1.0) #이미지의 픽셀 크기를 제한
with tf.GradientTape() as tape:
tape.watch(inp_imgs) #inp_imgs를 지속적으로 미분 계산하기 위한 용도
out_score = -model(inp_imgs) #이미지를 모델에 통과시켜 에너지 점수 얻기
grads = tape.gradient(out_score, inp_imgs) #입력에 대한 출력의 그래디언트 계산
grads = tf.clip_by_value(grads, -0.03, 0.03)
inp_imgs += -step_size * grads #작은 양의 그레디언트를 입력 이미지에 더함
inp_imgs = tf.clip_by_value(inp_imgs, -1.0,1.0)
return inp_imgs
3. Loss Function
이제 학습시킬 데이터 샘플과 신경망도 만들었으니 Loss Function을 정의해야한다.
우리는 위와 같은 p(x)라는 확률분포를 최적화시켜야한다. 보통 확률분포는 negative log likelihood를 쓰기 때문에 아래와 같은 식을 최적화시켜야 한다.
하지만, p(x)는 분모가 적분으로 구성되어 있기 때문에 일반적으로 n이 커질수록 알고리즘의 시간복잡도가 매우 높다. 따라서 최적해가 아닌 근사해를 구하는 목적함수로 수정하는데 이를 Constrative Divergence 라는 알고리즘이 등장한다.
따라서 위 objective를 아래와 같이 수정한다.
- 모델에 의해 생성된 데이터 분포에 대한 기대값을 나타냅니다.
- 는 모델의 파라미터입니다.
직관적으로는 우리가 만들려는 이미지의 E(x;θ)와 실제 데이터의 E(x;θ)를 가깝게 만들면 0에 가까워지므로 우리가 원하는 모델을 만들 수 있게 된다.
헷갈렸던 점
loss function을 낮추는 방법으로는 둘의 차이를 줄이는 방법도 있지만 극단적으로 data가 만드는 에너지의 평균이 의도와 다르게 엄청 커지고, model이 만드는 에너지의 평균을 엄청 낮추는 방향으로 학습할 수도 있다.
그래서 저렇게 loss function을 세우는게 맞나? 이 생각이 들어서 한참 고민했는데 그래서 정규화 loss가 들어간 거였다. 항상 모델 설계할 때 의도한대로 작동하는지를 잘 살펴보아야할 것 같다.
class Buffer:
def __init__(self, model):
super().__init__()
self.model = model
#샘플 초기화
self.examples = [
tf.random.uniform(shape = (1,32,32,1)) * 2 -1
for _ in range(128)
]
def sample_new_exmps(self, steps, step_size, noise):
#매번 평균적으로 샘플의 5%가 처음부터 랜덤한 잡음으로 생성됨
n_new = np.random.binomial(128, 0.05)
rand_imgs = (
tf.random.uniform((n_new, 32, 32, 1)) * 2 - 1
)
#나머지는 기존 버퍼로부터 랜덤하게 추출됨
old_imgs = tf.concat(
random.choices(self.examples, k = 128 - n_new), axis = 0
)
inp_imgs = tf.concat([rand_imgs, old_imgs], axis = 0)
#샘플을 합치고, 랑주뱅 샘플링을 실행
inp_imgs = generate_samples(
self.model, inp_imgs, steps = steps, step_size = step_size, noise = noise)
#만들어진 샘플을 examples에 추가하고 최대 샘플 갯수를 8,192개로 제한한다.
self.examples = tf.split(inp_imgs, 128, axis = 0) + self.examples
self.examples = self.examples[:8192]
return inp_imgs
class EBM(models.Model):
def __init__(self):
super(EBM, self).__init__()
self.model = model
self.buffer = Buffer(self.model)
self.alpha = 0.1
self.loss_metric = metrics.Mean(name = "loss")
self.reg_loss_metric = metrics.Mean(name = "reg")
self.cdiv_loss_metric = metrics.Mean(name = "cdiv")
self.real_out_metric = metrics.Mean(name = "real")
self.fake_out_metric = metrics.Mean(name = "fake")
@property
def metrics(self):
return [
self.loss_metric,
self.reg_loss_metric,
self.cdiv_loss_metric,
self.real_out_metric,
self.fake_out_metric
]
def train_step(self, real_imgs):
real_imgs += tf.random.normal(
shape = tf.shape(real_imgs), mean = 0, stddev = 0.005) 2
real_imgs = tf.clip_by_value(real_imgs, -1.0,1.0)
fake_imgs = self.buffer.sample_new_exmps(
steps = 60, step_size = 10, noise = 0.005) 2
inp_imgs = tf.concat([real_imgs, fake_imgs], axis = 0)
with tf.GradientTape() as training_tape:
real_out, fake_out = tf.split(self.model(inp_imgs),2, axis = 0)
cdiv_loss = tf.reduce_mean(fake_out, axis = 0) - tf.reduce_mean(real_out, axis = 0)
reg_loss = self.alpha * tf.reduce_mean(real_out ** 2 + fake_out **2, axis = 0)
loss = reg_loss + cdiv_loss
grads = training_tape.gradient(loss, self.model.trainable_variables)
self.optimizer.apply_gradients(
zip(grads, self.model.trainable_variables)
)
self.loss_metric.update_state(loss)
self.reg_loss_metric.update_state(reg_loss)
self.cdiv_loss_metric.update_state(cdiv_loss)
self.real_out_loss_metric.update_state(tf.reduce_mean(real_out), axis = 0)
self.fake_out_loss_metric.update_state(tf.reduce_mean(fake_out), axis = 0)
return {m.name : m.result() for m in self.metrics}
def test_step(self, real_imgs):
batch_size = real_imgs.shape[0]
fake_imgs = tf.random.uniform((batch_size, 32, 32, 1)) * 2 -1
inp_imgs = tf.concat([real_imgs, fake_imgs), axis = 0)
real_out, fake_out = tf.split(self.model(inp_imgs), 2, axis = 0)
cdiv = tf.reduce_mean(fake_out, axis=0) - tf.reduce_mean(real_out, axis = 0)
self.cdiv_loss_metric.update_state(cdiv)
self.real_out_metric.update_state(tf.reduce_mean(real_out, axis = 0))
self.fake_out_metric.update_state(tf.reduce_mean(fake_out, axis = 0))
return {m.name : m.result() for m in self.metrics[2:]}
ebm = EBM()
ebm.compile(optimizer = optimizers.Adam(learning_rate = 0.0001), run_eagerly = True)
ebm.fit(x_train, epochs = 60, validation_data = x_test, )
4. 훈련 과정
- 위의 내가 헷갈렸던 점처럼 epoch이 진행됨에 따라 손실이 감소되는 폭이 비교적 변화가 작다. 그 이유는 모델도 성능이 향상되어서 실제 생성된 이미지의 퀄리티도 개선되기 때문이다.
- 그래서 buffer에서 샘플을 추출하지 않고 랜덤한 잡음으로 만든 샘플의 점수를 활용했더니 차이가 잘 벌어지는 걸 볼 수 있다.
5. 새로운 샘플 만들기
start_imgs = np.random.uniform(size = (10,32,32,1)) * 2 - 1
gen_img = generate_samples(
ebm.model,
start_imgs,
steps = 1000,
step_size = 10,
noise = 0.005,
return_img_per_step = True)
참고 자료
1. https://biomadscientist.tistory.com/62 (대조 발산 알고리즘이 필요한 이유)
2. 만들면서 배우는 생성AI 2판. 7장
'Machine Learning > 유튜브, 책, 아티클 정리' 카테고리의 다른 글
| 만들면서 배우는 생성AI 9장 : 트랜스포머 모델 (0) | 2023.12.12 |
|---|---|
| 만들면서 배우는 생성 AI 8장 : Diffusion 모델 (1) | 2023.12.11 |
| 만들면서 배우는 생성 AI 정리 6장 - Normalizing Flow 모델 (1) | 2023.12.08 |
| 만들면서 배우는 생성 AI 정리 5장 - Autoregressive 모델 (0) | 2023.12.04 |
| 만들면서 배우는 생성 AI 정리 4장 - GAN 모델 (1) | 2023.11.21 |
