안녕하세요, donyy입니다!
올해 3월부터 시작한 졸업 연구도 어느덧 끝을 향해 달려가고 있습니다.
지난 스타트 학기 동안 Diffusion Model과 Quantization에 대한 기본적인 개념을 이해하기 위해 다양한 논문들을 리뷰했는데요!
이번 그로쓰 학기에는 그동안 학습한 내용을 바탕으로 구체적인 연구 주제를 설정하고, 실험을 통해 의미 있는 결론을 도출하는 과정에 집중하고 있습니다.
따라서 이번 포스팅에서는 이론적인 부분보다는 실제 구현에 초점을 맞추어, Simple Diffusion Model에 양자화를 적용하는 과정을 코드 중심으로 설명하려고 합니다.
우선 양자화 적용 방법에 대해 설명하기 전에 저희 팀의 연구 주제와 내용을 간략히 소개하겠습니다.
📃 연구 주제 및 개요
연구 주제
Layer-Adaptive Quantization on Diffusion Model using Fisher
피셔 정보를 활용한 확산 모델에의 레이어 적응형 양자화
저희 팀의 졸업 연구는 Diffusion Model이 timestep마다 서로 다른 layer를 활성화한다는 가설에서 출발했습니다.
이를 확인하기 위해 MNIST 데이터셋을 이용해 DDPM 모델을 학습시키고, Fisher Information 값을 추출해 각 timestep에서 활성화된 layer를 분석했습니다.
Diffusion Model은 timestep마다 노이즈를 제거하는 과정을 반복하여 이미지를 생성하는 방식을 채택하고 있기 때문에 고품질의 이미지를 생성한다는 장점을 갖지만, 동시에 대규모 파라미터로 인한 연산 복잡성, 느린 추론 속도, 높은 메모리 요구량이라는 기술적 한계를 가지고 있습니다.
저희 팀은 이중 높은 메모리 요구량 에 초점을 맞추어,
Diffusion Model의 메모리 요구량을 줄이되 이미지 품질 저하를 최소화하는 것을 연구 목표로 설정하였습니다!
그리고 이러한 연구 목표를 달성하기 위한 하나의 방안으로 Layer 별 양자화 기법이라는 아이디어를 떠올리게 되었습니다.
앞서 언급했듯이 저희의 연구는 Diffusion Model의 이미지 생성 과정에서 timestep에 따라 활성화되는 layer 가 다르다는 가설로부터 시작하게 되었는데요!
'해당 가설이 사실이라면 layer의 활성화 정도에 따라 layer 별로 다른 정도의 양자화를 적용하면 어떨까?' 하는 생각을 하게 된 것입니다!
다음의 예시를 통해 조금 더 쉽게 설명해볼게요.
하나의 모델이 총 두 개의 레이어(Layer 1, 2)로 구성되어있고, 하나의 레이어는 3개의 가중치 값으로 구성되어 있다고 해봅시다.
(사각형 안에 적힌 숫자가 가중치를 의미합니다.)
기존의 양자화 방식은 별도의 임곗값을 도입하지 않아, 모델을 구성하는 모든 가중치에 양자화를 적용합니다.
이미지에서 좌측을 보면, Layer 1과 2를 구성하는 모든 가중치 값이 양자화(회색)된 것을 볼 수 있죠.
반면에 레이어별로 임곗값을 설정하고, 임곗값을 넘지 못하는 가중치 값에만 양자화를 적용하면 높은 가중치 값, 즉, 중요도가 높은 가중치 값(파란색)은 유지되고, 그 외의 값(회색)에 대해서만 양자화가 적용되는 것이죠!
결론적으로 저희의 연구 과정은 다음과 같습니다.
- layer 별 fisher information 값을 바탕으로 layer 별 기여도(중요도) 분석
- 분석 결과를 바탕으로 (1) 단일 / (2) 레이어 그룹별 / (3) 레이어별 임곗값(threshold)을 설정 및 양자화 적용
세 가지의 실험 방법 중 저는 (3) 레이어별 임곗값 설정 실험을 담당하여 해당 실험을 기준으로 작성되었으며, 이번 포스팅에서는 개념적/이론적 내용이 아닌 코드 기반의 내용을 다루고 있으니 참고 부탁드립니다!
그럼 본격적으로 Simple Diffusion Model에 Quantization을 적용하는 방법을 소개하도록 하겠습니다!
🚀 Simple Diffusion Model에 Quantization 적용하기!
1. Simple Diffusion Model 구현하기
저희 팀이 연구에서 사용한 모델은 Simple Diffusion Model로 U-Net 구조를 기반으로 구현되었습니다.
이번 포스팅의 핵심은 양자화 적용이니 Simple Diffusion Model 구현에 대한 자세한 설명은 생략하도록 하겠습니다. 참고로, 모델은 다음의 글을 참고하여 구현하였습니다. Simple Diffusion Model 구현이 궁금하신 분들은 아래의 링크를 참고해주세요! :)
https://towardsdatascience.com/diffusion-model-from-scratch-in-pytorch-ddpm-9d9760528946
Diffusion Model from Scratch in Pytorch
Implementation of Denoising Diffusion Probabilistic Models (DDPM)
towardsdatascience.com
2. Fisher Information 기반 Layer 활성화도 분석하기
앞서 말씀드린 바와 같이 저희 팀의 연구 목표는 Layer별 활성화 분석을 기반으로 양자화 기법을 적용하는 것인데요.
이에 따라 양자화 적용에 앞서 Fisher information 계산이 필요합니다.
Fisher Information은 각 가중치가 모델 출력에 얼마나 중요한 영향을 미치는지를 나타냅니다. 따라서 이를 통해 Stable Diffusion Model이 이미지를 생성하는 과정에 있어서 중요한 가중치와 덜 중요한 가중치를 구별할 수 있습니다.
다음은 timestep별로 Fisher information 을 계산하고, 계산된 fisher information 값을 바탕으로 그래프를 그리는 코드입니다.
1) calculate_fisher_information_per_timestep: timestep별 Fisher information 계산
def calculate_fisher_information_per_timestep(model, z, scheduler, num_time_steps, per=50):
"""
각 time step에 대해 레이어 내의 weight별로 Fisher Information을 계산합니다.
Args:
model (torch.nn.Module): 모델
z (torch.Tensor): 초기 노이즈 텐서
scheduler (DDPM_Scheduler): 스케줄러
num_time_steps (int): time step의 수
per (int): Fisher 정보를 저장할 time step 간격
Returns:
dict: 각 time step별로 레이어와 weight별 Fisher 정보 값을 저장한 딕셔너리
"""
fisher_information_per_timestep = {t: {} for t in range(0, num_time_steps, per)}
criterion = nn.MSELoss()
model.eval()
for t in tqdm(reversed(range(1, num_time_steps)), desc='Calculating Fisher Information per Timestep'):
t_tensor = [t]
# `per` 간격마다만 Fisher 정보를 저장하므로 이때만 `z`를 새로 복사하여 `requires_grad=True` 설정
if t % per == 0:
z_temp = z.detach().clone().requires_grad_(True)
model.zero_grad()
else:
z_temp = z.detach() # Fisher 정보 계산 필요 없을 때는 `requires_grad=False`
output = model(z_temp, t_tensor)
# `per` 간격마다 손실 및 Fisher 정보 계산
if t % per == 0:
# 손실 계산 및 그래디언트 계산
loss = criterion(output, z_temp)
loss.backward()
# 각 레이어와 weight별로 Fisher 정보 저장
for name, param in model.named_parameters():
if 'Layer' in name and 'weight' in name and param.requires_grad:
layer_name = name.split('.')[0]
weight_name = name
# Fisher 정보 딕셔너리에 레이어와 weight 정보를 저장
if layer_name not in fisher_information_per_timestep[t]:
fisher_information_per_timestep[t][layer_name] = {}
# Fisher 정보 계산 및 저장
fisher_information_per_timestep[t][layer_name][weight_name] = param.grad.data.pow(2).mean().item()
# Fisher 정보 계산 후 그래디언트 해제
z_temp.grad = None
model.zero_grad()
# 노이즈 업데이트
temp = (scheduler.beta[t_tensor].cuda() / ((torch.sqrt(1 - scheduler.alpha[t_tensor].cuda())) *
(torch.sqrt(1 - scheduler.beta[t_tensor].cuda()))))
z = (1 / (torch.sqrt(1 - scheduler.beta[t_tensor].cuda()))) * z - (temp * output)
return fisher_information_per_timestep
2) plot_fisher_information_map: Fisher information 값 시각화
def plot_fisher_information_map(fisher_information_per_timestep, layer_group=None):
"""
Fisher Information 값을 각 레이어 또는 특정 레이어의 가중치별로 time step에 따라 시각화합니다.
Args:
fisher_information_per_timestep (dict): 각 time step별 레이어와 가중치의 Fisher 정보 딕셔너리
layer_group (str, optional): 특정 레이어 이름을 지정하면 해당 레이어의 가중치별로 시각화합니다.
None일 경우 전체 레이어를 시각화합니다.
"""
# file_name = f'fisher_information/fisher_information_map_{layer_group}.png'
file_name = f'fisher_information/fisher_information_map_{layer_group or "All_layers"}.png'
fisher_values = {}
for t, layers_info in fisher_information_per_timestep.items():
if layer_group is None:
# 전체 레이어에 대한 Fisher Information 평균 수집
for layer_name, weights_info in layers_info.items():
if weights_info: # 빈 딕셔너리 확인
layer_fisher_info = sum(weights_info.values()) / len(weights_info)
if layer_name not in fisher_values:
fisher_values[layer_name] = {'timesteps': [], 'fisher_info': []}
fisher_values[layer_name]['timesteps'].append(t)
fisher_values[layer_name]['fisher_info'].append(layer_fisher_info)
else:
# 특정 레이어의 가중치별 Fisher Information 수집
if layer_group in layers_info and layers_info[layer_group]:
for weight_name, fisher_value in layers_info[layer_group].items():
if weight_name not in fisher_values:
fisher_values[weight_name] = {'timesteps': [], 'fisher_info': []}
fisher_values[weight_name]['timesteps'].append(t)
fisher_values[weight_name]['fisher_info'].append(fisher_value)
# 시각화 설정
plt.figure(figsize=(12, 8))
plt.title(f"Fisher Information Over Timesteps - {layer_group if layer_group else 'All Layers'}")
plt.xlabel("Timestep")
plt.ylabel("Fisher Information")
plt.yscale('log')
# 각 레이어 또는 가중치별로 plot 추가
for name, data in fisher_values.items():
plt.plot(data['timesteps'], data['fisher_info'], label=name)
plt.legend(loc="upper right") # 레이어 또는 가중치 이름 표시
plt.tight_layout()
plt.show()
plt.savefig(file_name)
plt.close()
print(f"Fisher Information map plot saved to '{file_name}'")
위의 코드를 실행하여 Layer별 Fisher information 값을 시각화한 결과는 다음과 같았습니다.
저희가 구현한 Simple Diffusion Model은 총 여섯 개의 layer로 구성되어 있는데요.
초반에는 Down Layer 에 해당하는 Layer 1과 2가 비교적 높은 기여도를 보이는 것을 확인할 수 있습니다.
초기 단계에서 Layer 1과 Layer 2의 Fisher information 값이 높다는 것은, 두 레이어가 초기 타임스텝에서 노이즈 제거 및 이미지 생성에 중요한 역할을 한다는 것을 의미합니다.
참고로 Diffusion 모델의 reverse process에서는 초기 타임스텝에서 노이즈가 대부분의 이미지를 차지하고 있기 때문에, 초기 layer 에 해당하는 Layer 1, 2는 저수준 특징(low-level feature)을 복구하는 데 집중합니다.
예를 들면, 선이나 텍스처, 윤곽선 과 같은 것이 이러한 저수준 특징에 해당합니다.
다음은 Layer 별로 세부 모듈 가중치를 계산하고 시각화한 결과입니다.
여섯 개의 Layer 모두 전반적으로 비슷한 경향성을 보이는 것을 확인할 수 있습니다!
다만, y축의 값을 자세히 보면 Layer1 과 Layer 2의 값이 월등히 크다는 것을 알 수 있어요.
Layer 1의 서브 모듈들이 갖는 최대값은 $10^7$, Layer 2의 서브 모듈들이 갖는 최대값은 $10^4$ 인 것에 반해,
Layer 3~6의 최댓값은 약 $10^{-2}$ ~ $10$ 사이인 것을 확인할 수 있습니다.
이에 따라 Layer 1과 2의 중요도가 높음을 다시 한번 확인할 수 있었습니다.
3. Quantization 적용하기 ⭐️
자, 그럼 이제 모델에 양자화 적용을 해봅시다!
앞서 저희의 실험은 (1) 단일 / (2) 레이어 그룹별 / (3) 레이어별 임곗값(threshold)을 설정 및 양자화 적용이라는 총 세 가지 방법으로 진행되었다는 말씀과 함께, 저는 (3) 레이어별 임곗값 설정 실험을 담당하였다는 말씀드렸습니다.
이에 따라 해당 방법 및 코드 구현 과정을 설명해보도록 하겠습니다!
1) Layer 별 임계값 설정하기
우선 양자화를 적용시키는 기준이 될 임계값 설정이 필요합니다.
앞서 시각화한 도표를 통해 각각의 layer별 임계값을 수동으로 지정하는 방법도 있겠으나, 전반적인 경향성 확인을 위하여 임계값 비율을 설정하여 FP32를 FP16으로 양자화를 진행하였습니다.
여기서 잠깐, 임곗값 비율(threshold_ratio)이 무엇인지 잘 모르는 분들이 있으실텐데요!
임곗값 비율 p 가 주어졌을 때, 해당 레이어의 피셔 값 분포에서 (1-p)×100 백분위수에 해당하는 값 을 임곗값으로 선택하며, 이 값은 0.0에서 1.0 사이의 값을 갖습니다.
예를 들어 임곗값 비율이 0.3 인 경우, 피셔 값 분포의 상위 30%에 해당하는 값이 해당 레이어의 임곗값으로 채택되는 것 입니다!
# Fisher 값 기반으로 threshold 계산
thresholds = {}
for layer_name, weights_info in timestep_1_info.items():
fisher_values = list(weights_info.values()) # 모든 weight의 Fisher 값 가져오기
fisher_values.sort(reverse=True) # 내림차순 정렬
cutoff_index = int(len(fisher_values) * threshold_ratio) # 비율에 따라 cutoff 계산
threshold_value = fisher_values[cutoff_index] if cutoff_index < len(fisher_values) else fisher_values[-1]
thresholds[layer_name] = threshold_value # 각 레이어의 threshold 저장
- Fisher Information으로부터 각 레이어별 임계값을 리스트에 저장하고, 이를 내림차순으로 정렬합니다.
- 정렬된 fisher_values 리스트의 길이에 threshold_ratio를 곱하여 임계값의 인덱스를 구하고, 해당 인덱스를 통해 임계값을 구하여 threshold_value 리스트에 저장합니다.
- 해당 작업을 반복하여 마치고 나면 threshold_value에는 각 Layer 별 임계값이 순서대로 (Layer 1의 임계값, Layer 2의 임계값, ... , Layer 6) 저장되게 됩니다.
2) 양자화 적용하기
# 양자화 적용
with torch.no_grad():
quantized_count = {} # 각 레이어의 양자화된 가중치 개수를 저장
for layer_name, threshold_value in thresholds.items():
quantized_count[layer_name] = 0 # 초기화
for weight_name, fisher_value in timestep_1_info[layer_name].items():
print(f' layer_name = {layer_name}, weight_name={weight_name} fisher_value = {fisher_value} / threshold_value = {threshold_value}')
if fisher_value < threshold_value: # threshold보다 작은 weight만 양자화
for name, param in model.named_parameters():
if name == weight_name: # weight 이름이 일치하면
print(f"Applying half to {name} with Fisher value {fisher_value} (Threshold: {threshold_value})")
param.data = param.data.half() # FP16으로 변환
quantized_count[layer_name] += 1 # 양자화된 weight 개수 증가
break- 타임스텝별로 Layer 정보와 threshold를 순회하여 현재 Layer에 해당하는 임계값을 가져옵니다.
- Layer 내 가중치별 Fisher 값을 순회하여 Fisher information value가 임계값보다 낮은 경우, 해당 가중치를 양자화합니다.
- model.named_parameters()를 통해 모델의 가중치 이름과 Fisher Information의 weight_name을 매칭하고, 일치하는 가중치 데이터(param.data)를 FP16(half precision)으로 변환합니다.
- 양자화된 가중치의 개수를 기록하고, 이를 레이어별로 정리합니다.
양자화 적용 전체 코드는 다음과 같습니다!
def apply_fp16_quantization(model, fisher_information_per_timestep, threshold_ratio):
"""
Fisher 정보에 기반하여 timestep 1의 레이어별로 주어진 비율의 threshold 값을 적용하여
weight를 절반으로 양자화한 모델을 반환합니다.
Args:
model (torch.nn.Module): 모델.
fisher_information_per_timestep (dict): 각 time step과 레이어의 weight별 Fisher 정보 값을 포함한 딕셔너리.
threshold_ratio (float): 임계값 비율 (0.0 ~ 1.0)로, 하위 비율의 Fisher 정보를 걸러내기 위한 값.
Returns:
torch.nn.Module: 양자화된 모델.
"""
# Fisher 정보에서 timestep 1만 사용
timestep_1_info = fisher_information_per_timestep.get(1, {})
# Fisher 값 기반으로 threshold 계산
thresholds = {}
for layer_name, weights_info in timestep_1_info.items():
fisher_values = list(weights_info.values()) # 모든 weight의 Fisher 값 가져오기
fisher_values.sort(reverse=True) # 내림차순 정렬
cutoff_index = int(len(fisher_values) * threshold_ratio) # 비율에 따라 cutoff 계산
threshold_value = fisher_values[cutoff_index] if cutoff_index < len(fisher_values) else fisher_values[-1]
thresholds[layer_name] = threshold_value # 각 레이어의 threshold 저장
# 양자화 적용
with torch.no_grad():
quantized_count = {} # 각 레이어의 양자화된 가중치 개수를 저장
for layer_name, threshold_value in thresholds.items():
quantized_count[layer_name] = 0 # 초기화
for weight_name, fisher_value in timestep_1_info[layer_name].items():
print(f' layer_name = {layer_name}, weight_name={weight_name} fisher_value = {fisher_value} / threshold_value = {threshold_value}')
if fisher_value < threshold_value: # threshold보다 작은 weight만 양자화
for name, param in model.named_parameters():
if name == weight_name: # weight 이름이 일치하면
print(f"Applying half to {name} with Fisher value {fisher_value} (Threshold: {threshold_value})")
param.data = param.data.half() # FP16으로 변환
quantized_count[layer_name] += 1 # 양자화된 weight 개수 증가
break
# 각 레이어의 양자화된 weight 개수 출력
print("\nQuantization Summary:")
for layer_name, count in quantized_count.items():
print(f" Layer {layer_name}: Quantized {count} weights")
return model
4. Quantization 적용 결과
실험은 timestep 1000, 생성된 이미지(generated_images) 100개를 기준으로 진행되었으며,
성능 평가는 FID score와 메모리 용량(MEM)을 지표로 하여 진행하였습니다.
실험 결과 비교를 위하여 비교군을 두 가지 설정하였는데요!
- 양자화 적용 전(Simple Diffusion Model) 모델 → 메모리 사용량의 감소 정도를 효과적으로 비교 가능
- 기존의 양자화 방식(전체 양자화 적용)을 적용한 모델 → FID 점수의 개선 정도를 효과적으로 비교 가능
두 비교군의 결과는 다음과 같았습니다.
1) 양자화 적용 전(Simple Diffusion Model) 결과
- FID Score: 22.7512
- Total Parameter Count: 35178689
- Total Size in MB: 134.20 MB
2) 기존의 양자화 방식(전체 양자화 적용)을 적용한 모델
- FID Score: 24.1324
- Total Parameter Count: 35178689
- Total Size in MB: 67.70 MB
그럼 이제 대망의 layer별 임곗값 설정 실험 결과!
(두구두구두구... 🥁🥁)
3) layer별 임곗값 설정 및 양자화 적용 결과
| threshold_ratio | FID (↓) | MEM (↓) |
| 양자화 전 (비교군) |
22.7512 | 134.20 MB |
| 전체 양자화 (비교군) |
23.3597 | 67.70 MB |
| 0.1 | 23.2198 | 81.31 MB |
| 0.15 | 23.2370 | 91.50 MB |
| 0.2 | 23.8857 | 92.06 MB |
| 0.25 | 22.4168 | 102.26 MB |
| 0.3 | 23.6778 | 105.35 MB |
| 0.35 | 23.1871 | 115.27 MB |
| 0.4 | 22.5009 | 118.36 MB |
| 0.45 | 22.9136 | 128.56 MB |
| 0.5 | 23.1005 | 131.09 MB |
FID score는 실제 이미지와 생성 이미지를 비교하여 생성 이미지 품질을 측정하는 지표로, 값이 낮을 수록 품질이 우수함을 나타내는데요!
임곗값 비율이 0.25일 때 FID score = 22.4168 으로 가장 낮은 FID 값을 갖는 것을 확인하였습니다.
이는 전체 양자화를 적용하는 방식(FID = 23.3597) 대비 약 7% 개선된 결과입니다.
또한 임곗값 비율이 0.25일 때 양자화가 적용되지 않은 모델에 비해 메모리 사용량도 134.20MB → 102.26MB로, 약 23.8% 가량 줄어든 것을 확인할 수 있습니다.
여기서 잠깐, '뭐야, 양자화 전보다는 FID score가 높은데?' 라고 생각하실 수 있는데요,
이건 사실 당연한 결과랍니다!
양자화가 적용되지 않은 '양자화 전' 모델은 양자화로 인해 소실되는 정보가 없으니 가장 좋은 이미지 품질을 갖는 게 당연한 거죠!
따라서 우리가 주목해야할 것은 '양자화 전' 모델의 FID score 보다 낮은 점수를 갖는 경우가 있는지가 아니라, 다양한 임곗값 비율을 기준으로 양자화한 모델들 중 가장 낮은 FID score를 갖는 경우를 찾는 것입니다!
위의 결과를 바탕으로 그래프를 그려보면 다음과 같습니다.
실험 결과를 다시 한 번 정리하자면, 임곗값 비율을 0.25로 설정하여 양자화한 모델이 모델 크기는 양자화 이전 대비 23.8% 감소하고, FID가 전체 양자화 방식 대비 약 7% 개선되어 Diffusion Model의 경량화와 이미지 품질 간의 균형을 어느 정도 이룬 것을 확인할 수 있었습니다!
여기까지 하면 Simple Diffusion Model에 양자화 적용하기 완료!
이번 포스팅에서는 양자화 적용 및 코드 구현에 초점을 맞추어 글을 작성해보았는데요.
더욱 다양한 임계값에 따른 양자화 적용 결과와 저희 팀의 졸업 연구 결과가 궁금하시다면 다음 포스팅을 기다려주세요!
'AI > ML & DL' 카테고리의 다른 글
| [DL] 딥 러닝 기본 이론 3: 경사 하강법 (0) | 2024.09.22 |
|---|---|
| [DL] 딥 러닝 기본 이론 2: 활성화 함수 (0) | 2024.09.22 |
