클래스 불균형 데이터셋의 효과적인 분류를 위한 조건부 확산 모델 기반 데이터 증강 기법

Ⅰ. 서 론

이미지 분류는 다양한 산업에서 다루는 핵심과제 중 하나이다. 대용량의 이미지 데이터가 끊임없이 생성되는 현대 산업에서, 수많은 이미지를 분류하는 것은 정보를 효율적으로 검색하고 활용하게 하는 데 필수적이다. 머신러닝 및 딥 러닝 기술의 발전은 이미지 분류 모델의 상당한 발전을 이루어 냈다. 예를 들어, 딥 러닝 모델을 통해 X-ray와 MRI 이미지를 분석하여 병변을 찾아내거나, 유방암 및 기타 종양을 빠르고 정확하게 진단할 수 있다[1][2]. 또한, 이미지 분류 모델은 제조 분야에서도 제품 불량을 판별하거나, 생산 공정에서의 이상 탐지에 사용되고 있다[3][4]. 이 외에도 해충, 패션 이미지 분류, 상품 카테고리 분류 등과 같이 다양한 분야에서 활용되고 있다[5]-[7].

그러나 이미지 분류 작업에서는 데이터셋 구성에 주의가 요구된다. 모델의 학습 과정에서 특정 클래스에 치우쳐 학습되어 학습될 수 있기 때문이다. 이는 분류 문제에서 흔히 발생하는 클래스 불균형 문제로 데이터의 구성이 몇 개의 클래스로 치우쳐 있을 때 오분류의 가능성을 높이고 모델의 편향을 증가시킨다. 이는 다양한 산업에서 활용되는 머신러닝 모델의 성능과 신뢰도를 저하시키는 주요 원인이 된다. 따라서 클래스 불균형을 해결하는 것은 모델의 예측 성능 개선과 모델의 신뢰성 및 신용성을 확보하는데 중요한 요구사항이다[8].

클래스 불균형 문제를 해결하기 위한 다양한 방법들이 제안되어왔다. 대표적인 방법으로 오버샘플링, 다운샘플링 등이 있다. 그러나 이러한 방법들은 기존 데이터 손실이나 과적합 문제를 야기할 수 있으므로 더 효과적인 클래스 불균형 처리를 위한 연구가 지속되고 있다. 최근에는 딥 러닝 기반의 VAE(Variational AutoEncoder)[9]와 GAN(Generative Adversarial Network)[10]과 같은 생성 모델을 활용한 접근법들이 주목받고 있다. 생성 모델의 발전에 따라 다양한 분야에서 데이터 불균형 문제를 해결하기 위해 사용되고 있다[11]-[14]. 그러나 VAE는 생성된 이미지가 흐릿하고 학습이 불안정하며 GAN은 데이터의 다양성을 제대로 반영하지 못하는 한계점이 있다.

한편, 최근에는 고품질의 이미지를 안정적으로 만들어 내며 VAE, GAN의 한계점을 극복한 Diffusion 모델이 등장하였다. Diffusion 모델은 데이터를 생성하는 모델 중 하나로, 원본 데이터의 분포를 학습하여 새로운 데이터 포인트를 생성하는 데 사용된다[15]. 특히, Diffusion 모델은 복잡하고 다양한 데이터 분포를 정확하게 표현할 수 있어 다양한 응용 분야에서 주목받고 있다.

본 연구에서는 Diffusion 모델의 조건부 생성을 활용하여 클래스 불균형을 해결하는 두 가지 방법을 제안한다. 조건부 생성 기법은 텍스트, 클래스와 같은 추가 정보에 따라 데이터를 생성하는 작업을 의미한다. 첫 번째 제안 방법은 Classifier-free guidance[16]에 클래스 불균형 데이터셋을 학습시켜 데이터의 개수가 부족한 클래스의 데이터를 생성하였다. 이미지를 생성하는 과정에서 생성하고자 하는 클래스 정보를 입력하여 생성한다. 두 번째 제안 방법은 사전 학습된 Stable Diffusion[17]의 텍스트 조건부 생성을 활용해 원하는 클래스의 데이터를 생성하였다. 생성하는 과정에서 텍스트 정보를 입력하여 생성한다. 데이터가 부족한 클래스의 이미지를 생성하기 위해 텍스트 정보는 해당 클래스의 이름(예: cat)을 사용하였다. 본 논문에서는 생성된 데이터와 기존의 클래스 불균형 데이터를 결합하여 균형 잡힌 데이터셋으로 재구성하였다. 이를 통해 이미지 분류 작업을 수행하였으며 사용한 모델은 ResNet18과 ResNet50[18] 이다.

실험은 제안된 두 가지 기법을 적용하여 생성된 이미지를 기존 이미지 데이터 증강연구에서 많이 활용되었던 방법인 Mixup[19]와 CutMix[20]를 비교 방법론으로 활용하였다. 기존 CIFAR10 데이터셋과 생성된 이미지의 유사성을 평가하기 위해 FID(Fréchet Inception Distance) 점수를 사용하여 비교하였으며, 각각 생성된 이미지로 균형 잡힌 데이터셋을 분류 모델에 학습시켜 Accuracy와 F1-score를 비교하였다. 실험 결과, 본 논문에서 제안된 방법이 기존의 증강기법에 비해 FID 성능을 향상하고 CIFAR10에 근사한 이미지를 생성하였으며 분류 성능 또한 유의미하게 향상되었다.

논문의 구성은 다음과 같다. 2장에서는 클래스 불균형 관련 연구들과 논문에서 활용한 Diffusion 모델을 서술한다. 3장과 4장에서는 제안한 방법론과 실험 설계 및 결과를 서술하며, 마지막으로 5장에서는 결론으로 연구의 요약과 추후 연구에 대해 제시하였다.

Ⅱ. 관련 연구

2.2 Diffusion model

Diffusion 모델은 이미지 생성에 사용되는 확률적 방법론 중 하나이다. Diffusion 모델은 원본 이미지에 점진적으로 노이즈를 추가하는 forward process와, 노이즈가 추가된 이미지를 원본 이미지로 복원하는 reverse process로 이루어져 있다. Reverse process를 통해서 forward process에 점진적으로 추가된 노이즈를 제거하여 원본 이미지를 추정하며 학습한다. 그림 1은 Diffusion 모델의 forward process와 reverse process를 표현한 것이다.

Fig. 1.

Diffusion process

이후, Diffusion 모델의 기본 구조를 발전시킨, DDPM(Denoising Diffusion Probabilistic Model)[23]이라는 새로운 모델이 제안되었다. DDPM은 Diffusion 모델의 수식과 모델 학습 절차를 개선하여 높은 품질의 이미지를 생성하였다. DDPM의 forward process 는 주어진 q(x)의 분포를 따르는 원본 이미지 x₀ ~ q(x)에 점진적으로 노이즈를 더해가는 과정이며, 이는 t - 1 시점에서 Markov Chain 방식에 따라 β_t의 분산을 가진 가우시안 노이즈가 추가되는 방식으로 진행된다. Forward process은 식 (3)과 같이 나타낼 수 있다.

$\[ \begin{array}{c}q\left(x_{1:T}\left|x_0\right.\right):=\prod _{t=1}^{T}q\left(x_t\left|x_{t-1}\right.\right)\hfill \\ q\left(x_t\left|x_{t-1}\right.\right):=N\left(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I\right)\hfill \end{array} \]$

(3)

Reverse process에서는 모델은 Denoising을 학습하여, 특정 시점 t의 이미지 x_t로부터 x_t-1의 평균 μ_θ와 분산 Σ_θ을 예측하며 학습한다. Reverse process는 아래 식 (4)로 표현된다.

$\[ \begin{array}{c}p\left(x_{0:T}\right)=p_{\theta }\left(x_T\right)\prod _{t=1}^T{p}_{\theta }\left(x_{t-1}\left|x_t\right.\right)\hfill \\ p_{\theta }\left(x_{t-1}\left|x_t\right.\right)=N\left(x_{t-1};{\mu }_{\theta }\left(x_t,t\right),\sum _{\theta }\left(x_t,t\right)\right.\hfill \end{array} \]$

(4)

Diffusion 모델의 목적은 원본 데이터의 분포 p_θ(x₀)를 학습하는 것이며, 이는 식 (5)의 loss를 최소화하면서 진행된다. 식 (5)의 loss를 최소화하는 reverse process를 학습하여 노이즈로부터 데이터를 만들어내는 생성 모델을 만들어낸다.

$\[ \begin{array}{c}E\left[-\text{log}{p}_{\theta }\left(x_0\right)\right]\le E_q\left[-\text{log}\frac{p_{\theta }\left(x_{0:T}\right)}{q\left(x_{1:T}\left|x_0\right.\right)}\right]\hfill \\ =E_q\left[-\text{log}p\left(x_T\right)-{\sum }_{t\ge 1}\text{log}\frac{p_{\theta }\left(x_{t-1}{x}_t\right)}{q\left(x_t\left|x_{t-1}\right.\right)}\right]\hfill \\ =:L\hfill \end{array} \]$

(5)

DDPM은 기존의 loss 함수가 복잡한 데이터 분포와 관련된 문제를 효과적으로 처리하기 어렵다는 한계점을 해결하기 위해 ELBO(Evidence Lower BOund)를 도입하였다. ELBO의 KL-Divergence를 사용하여 분포의 차이를 계산하고, 이를 통해 모델의 loss 함수를 최적화하였다. 아래 식 (6)는 ELBO를 적용한 loss 함수이다.

$\[ \begin{array}{c}\text{log}p\left(x\right)\ge \hfill \\ E_q\left(x_1\left|x_0\right.\right)\left[\text{log}{p}_{\theta }\left(x_0\left|x_1\right.\right)\right.\hfill \\ -D_{KL}\left(q\left(x_T\left|x_0\right.\right)\left|p\left(x_T\right)\right.\right)\hfill \\ \left.-\sum _{t=2}^T{E}_{q\left(x_t\left|x_0\right.\right)}\left[D_{KL}\left(q\left(x_{t-1}\left|x_t,x_0\right.\right)\left|p_{\theta }\right.\right)\right)\right]\hfill \\ =L_0-L_T-\sum _{t=2}^T{L}_{t-1}\hfill \end{array} \]$

(6)

또한, DDPM에서는 inductive bias를 늘려 모델의 학습과 최적화 과정을 개선하였다. 최종적인 DDPM 식 (7)로 표현할 수 있으며 은 추정되는 노이즈를 의미한다.

$\[ \begin{array}{c}{L}_t^{simple }=\hfill \\ E_{x_0,t,ϵ}\left[{∥ϵ-{ϵ}_{\theta }\left(\sqrt{{\stackrel{-}{\alpha }}_t}{x}_0+\sqrt{1-{\stackrel{-}{\alpha }}_t}ϵ,t\right)∥}^2\right]\hfill \end{array} \]$

(7)

따라서, 본 논문에서 Diffusion model을 활용한 이미지 생성을 통해 클래스 불균형을 해결하고자 한다.

Ⅲ. 방법론

본 논문에서는 클래스 불균형 데이터셋을 균형 잡힌 데이터셋으로 재구성하기 위해 이미지 생성 방법론 중 하나인 Diffusion 모델을 사용하였다. 본 논문에서는 Diffusion 모델의 조건부 생성 기법을 활용하여 소수 클래스에 대한 클래스 불균형의 문제를 해결하고자 한다. 이미지를 생성하여, 클래스 불균형 데이터셋을 균형 잡힌 데이터셋으로 재구성하였다. 첫 번째로 제안한 방법은 Classifier-free guidance[16] 모델을 사용하여 클래스 조건부 생성을 통해 새로운 이미지를 만들었으며, 두 번째 제안 방법에서는 LDM[17]을 활용해 텍스트 조건부 생성을 통해 이미지를 만들었다. 이를 통해, 데이터가 부족한 클래스의 이미지를 기존의 균형 잡힌 원본 CIFAR10 이미지와 유사하게 재구성할 수 있었다. 이후, 재구성된 데이터셋을 통해 이미지 분류 성능을 확인하였다.

3.1 Classifier-free guidance

DDPM[23]은 뛰어난 생성 능력을 통해 기존의 VAE[9], GAN[10]보다 다양성, 품질 측면에서 우수한 성능을 보여주었다. Diffusion 모델은 UNet[24] 기반의 구조로 구현되었다. 이 구조는 이미지를 생성하는 과정에서 중요한 역할을 하는데, 각 계층에서의 특징 추출을 통해 세밀하고 정확한 이미지 생성이 가능하게 만든다. DDPM 이후에는 Classifier를 추가로 학습해 원하는 클래스의 이미지를 생성할 수 있는 Classifier-Guidance(ADM)[25]와 Classifier를 따로 학습하지 않고, Diffusion 모델에서 학습하는 CFG가 등장하였다. 먼저, ADM은 생성된 이미지의 품질을 향상하기 위해 별도의 Classifier를 학습하는 방법을 제안하였다. 기존 Diffusion 모델의 학습 방법은 동일하나, Classifier로부터 나오는 gradient를 예측한 노이즈에 추가한다. Classifier의 기울기를 조절함으로써 해당 클래스의 이미지를 생성할 수 있었다. 그러나 이러한 모델 구조는 추가적인 Classifier를 학습시켜야 하므로 Diffusion 모델의 학습이 복잡해지고, Classifier는 노이즈가 있는 데이터에 대해 학습해야 하므로 사전학습된 Classifier를 활용할 수 없다는 한계점이 존재했다.

이러한 문제점을 해결한 방법이 CFG이다. CFG는 Classifier의 기울기를 이용한 샘플링 대신, 조건부 Diffusion 모델의 출력값과 비조건부 Diffusion 모델의 출력값 차이의 방향으로 샘플링 하도록 유도한다. 아래 식 (8)은 CFG의 샘플링 함수를 나타낸다.

$\[ {\widetilde{ϵ}}_{\theta }\left(z_{\lambda },c\right)=\left(1+w\right){ϵ}_{\theta }\left(z_{\lambda },c\right)-w{ϵ}_{\theta }\left(z_{\lambda }\right) \]$

(8)

w는 guidance scale로 w가 커질수록 특정 클래스의 이미지만 생성되므로 CFG를 활용하여 원하는 클래스의 데이터를 생성할 수 있다. 본 논문에서는 클래스가 불균형한 CIFAR10 데이터셋을 CFG 모델에 학습시켜 소수 클래스의 데이터를 생성하였다. 그림 2는 제안된 방법론의 전체 과정을 나타낸다. 불균형한 CIFAR10 데이터셋을 CFG 모델에 학습시켜 이미지 데이터를 생성하였다. 이후, 생성된 데이터는 원래의 불균형 데이터셋에 추가하여 최종적으로 균형 잡힌 클래스를 가진 데이터셋으로 재구성하였다.

Fig. 2.

Image classification framework utilizing classifier-free guidance

3.2 Stable Diffusion(LDM)

클래스 정보뿐만 아니라 텍스트 정보를 활용한 GLIDE[26], Imagen[27], DALL-E2[28] 같은 방법론들도 연구되고 있다. 그러나 [26]-[28]의 Diffusion 모델은 픽셀값을 직접 예측하는 Diffusion model의 특성상 고해상도 이미지를 생성하기 위해서는 많은 컴퓨팅 자원을 필요로 하는 한계가 있다. 고해상도 이미지는 이미지의 세부 정보와 패턴을 더 잘 보존하여 정교한 결과를 얻을 수 있다. 따라서 고해상도 이미지를 용얻기 위해서는 별도로 Super-Resolution 모델을 학습하여 적용한다[29].

고해상도 이미지 생성은 시각적 세부 사항을 더 정확하게 보존함으로써 더욱 풍부하고 실제적인 결과를 만들어낸다. 이에 따라 데이터 증강과 함께 고해상도 이미지 생성 기술은 다양한 응용 분야에서 효과적인 결과를 나타냈다[30].

Stable Diffusion(LDM)[17]은 이러한 기존 Diffusion 모델의 제약을 극복하기 위해 잠재공간을 활용해 고해상도의 이미지를 생성하고, 더 적은 컴퓨팅 자원으로 효율적인 Diffusion 모델을 개발하였다. 이를 통해, 기존의 잠재공간을 사용하지 않는 GLIDE, Imagen, DALL-E2 보다 적은 컴퓨팅 자원과 시간을 소모하여 뛰어난 품질의 이미지를 생성할 수 있게 되었다.

식 (9)은 LDM의 loss 함수를 나타내며, 이를 통해 픽셀값이 아닌 잠재공간 임베딩 z가 입력으로 사용된다는 것을 확인할 수 있다.

$\[ \begin{array}{c}{L}_{DM}=E_{x,ϵ\sim N\left(\mathrm{0,1}\right), t}\left[{∥ϵ-{ϵ}_{\theta }\left(x_t,t\right)∥}_2^2\right]\hfill \\ L_{LDM}=E_{ϵ\left(x\right),ϵ\sim N\left(\mathrm{0,1}\right),t}\left[{∥ϵ-{ϵ}_{\theta }\left(z_t,t\right)∥}_2^2\right]\hfill \end{array} \]$

(9)

본 연구에서는 사전 학습된 LDM[17]을 사용하여 클래스 불균형 데이터셋에서 부족한 클래스의 이미지를 생성하였다.

고해상도 이미지를 생성하여 원본 CIFAR10에 근사한 이미지를 생성하고자 하였다. LDM의 생성 과정에 사용되는 텍스트 정보는 클래스 불균형 데이터셋의 각 클래스 이름으로 설정하였다. 그림 3은 LDM을 활용한 이미지 생성 과정을 나타낸 것이다. LDM의 출력 이미지는 512x512 이미지이기 때문에 클래스 불균형 데이터셋의 이미지 사이즈로 축소하여 재구성하였다.

Fig. 3.

Framework of proposed method with generation

재구성된 이미지는 CFG 방법론과 동일한 방식으로 기존의 클래스 불균형 데이터셋과 결합하여 균형 잡힌 데이터셋을 구성하였다.

Ⅳ. 실험 및 결과

4.1 Dataset

본 논문에서는 CIFAR10 데이터셋을 사용하였다. CIFAR10 데이터셋은 10개의 클래스로 구성되어 있으며, 각 클래스별로 32x32 해상도의 5,000개의 학습 데이터와 1,000개의 평가 데이터로 구성되어 총 60,000개의 데이터셋으로 이루어져 있다.

본 논문에서는 CIFAR10 데이터셋을 클래스 불균형에 대한 실험에 적용하기 위해 클래스별로 데이터 개수가 같은 CIFAR10 데이터를 재구성하여 사용하였다[30]. 클래스 불균형한 CIFAR10 데이터셋을 구성하기 위해 식 (10)을 통해 적용하였다.

$\[ N_i^{train }=500\cdot \left(i+1\right) \]$

(10)

i는 각 클래스의 인덱스 번호이며, 클래스별로 점진적으로 개수가 늘어나도록 설정하였다. 즉, 0번 클래스의 데이터 개수는 500개, 9번 클래스의 데이터 개수는 5000개이다. 최종적으로 학습에 사용한 각 클래스별 데이터의 개수는 아래 표 1과 같다. 테스트 데이터셋의 경우 기존 테스트셋을 사용하여 평가하였다.

Table 1.

Reconstructed CIFAR10 data sets

4.3 실험결과

4.3.1 이미지 생성 성능 평가

표 2은 대표적인 데이터 증강기법인 Mixup[19]와 CutMix[20], 그리고 본 연구에서 제안하는 두 가지 기법을 사용하여 생성된 이미지들의 FID 결과를 보여준다. CutMix의 이미지를 섞는 비율을 나타내는 하이퍼 파라미터는 0.6으로 설정하였다. 표 2의 첫 번째 행과 두 번째 행은 각각 기존의 데이터 증강기법인 Mixup와 CutMix를 이용하여 생성된 이미지의 FID 결과를 나타낸다. 첫 번째로 제안된 기법인 CFG는 guidance scale(w)에 따라 다양한 이미지를 생성하는데, 본 논문에서는 guidance scale(w)에 따라 이미지 생성 품질을 비교하기 위해 두 가지 guidance scale(w)을 사용하여 이미지를 생성하였다. guidance scale(w)의 영향을 거의 받지 않는 w = 0.1와 LDM에서 사용한 guidance scale인 w = 7.5을 사용하여 비교하였다. 먼저, w = 0.1일 때의 FID는 28.028였으며 w = 7.5일 때는 FID가 17.678로 guidance scale(w)의 크기에 따라 생성된 이미지의 분포에 차이가 발생하며 guidance scale(w)이 클수록 기존 CIFAR10에 유사한 이미지를 생성함을 확인할 수 있었다.

Table 2.

Performance of FID of reconstructed CIFAR10 datasets with generated images added

또한, guidance scale(w)에 관계없이 CFG를 사용하여 재구성한 데이터셋 모두 기존의 Mixup, CutMix보다 기존 원본 CIFAR10 데이터셋의 분포와 유사한 것을 확인하였다.

두 번째 제안 방법론인 LDM은 텍스트 프롬프트를 이용하여 각 클래스에 맞는 이미지를 생성하였다. 이 방법으로 생성된 이미지의 FID 값은 6.654로 기존 증강기법뿐 아니라 CFG보다도 더 높은 성능을 보여주었다.

그림 4는 제안된 생성 기법을 통해 생성된 이미지를 정성적으로 비교한 것이다. CFG의 경우 크기가 커질수록 이미지의 객체가 선명하게 나타났으며, LDM의 경우 CFG로 생성된 이미지보다 고품질의 이미지를 생성하며 ‘비행기’에 가까운 이미지를 생성하는 것을 확인할 수 있다.

Fig. 4.

Image generation results with proposed methods

4.3.2 이미지 분류 성능 평가

본 논문에서는 이미지 분류 성능을 평가하기 위해 ResNet18, ResNet50[18] 모델을 사용하였다. ResNet은 이미지 분류에 광범위하게 활용되는 대표적인 모델 중 하나이다. ViT[31], ConvNext[32] 등과 같이 최근에 제안된 더 깊고 복잡한 모델들은 CIFAR10을 적용할 경우, 데이터의 양이 상대적으로 적기 때문에 성능이 떨어지는 경향이 있어 선택하지 않았다. 이에, ResNet[18]만을 활용하여 분류 성능을 평가하였다.

ResNet18을 사용하여 기존의 데이터 증강기법인 Mixup[19]과 CutMix[20]로 재구성된 데이터를 학습한 경우에 비해, CFG(w = 0.1)로 생성된 데이터를 적용하였을 때 정확도가 향상되었으며, 특히 CFG(w = 7.5)로 생성된 데이터를 적용한 것은 더 높은 정확도를 보였다.

Table 3.

Performance of accuracy of reconstructed CIFAR10 with generated images added

ResNet50에서도 기존 증강기법들보다 개선된 성능을 확인할 수 있었으나 ResNet18보다 성능이 개선되진 않았다. 이는 실험에 활용된 CIFAR10 데이터의 수가 많지 않으므로 네트워크의 깊이와 복잡성이 성능 향상을 보장하지 않는 것을 알 수 있다.

표 4는 각 클래스별 F1 score를 비교한 결과이다. 재구성된 실험 데이터셋에서 생성된 데이터의 비율이 높지 않은 클래스 8, 9번은 제안된 방법론과 기존의 증강기법과 큰 차이가 없는 것을 확인할 수 있다. 반면, 생성된 데이터의 비율이 높은 클래스 0, 1번에서는 제안된 방법론의 F1 score가 더 높은 것을 확인할 수 있었다. 이를 통해서 제안된 방법론을 통해 재구성된 데이터셋이 불균형 문제를 효과적으로 완화되고 분류 모델이 소수 클래스를 더 잘 인식하게 되었다는 것을 알 수 있다.

Table 4.

Performance of F1-score of reconstructed CIFAR10 with generated images added

또한, LDM의 분류 성능이 CFG를 포함한 기존 증강 기법보다 우수한 성능을 보이며 고해상도의 이미지 생성을 통한 데이터 증강이 분류 성능의 향상에도 유의미한 영향을 끼쳤음을 확인하였다.

Ⅴ. 결론 및 향후 과제

본 논문에서는 클래스 불균형 이미지 데이터에 분류 성능 개선을 위한 데이터 증강기법으로 Diffusion 모델 기반 2가지 방법론을 제시하였다. 첫 번째로 제안한 방법론은 클래스 불균형 데이터셋을 학습한 Classifier-free guidance를 활용해 소수 클래스의 이미지를 생성하였다. 두 번째로 제안한 방법론은 사전 학습된 LDM의 텍스트 입력을 클래스명으로 사용해 소수 클래스의 이미지를 생성하였다. 제안된 방법론을 통해 생성된 데이터셋은 클래스 불균형 데이터셋과 결합해 재구성하여 이미지 분류 모델인 ResNet으로 학습하였다.

제안된 두 방법론은 이미지 생성 및 분류 분야에서 널리 활용된 기존 데이터 증강 기법들에 비해 효과적인 성능을 보였으며, 특히 두 번째로 제안된 방법론인 LDM을 활용한 기법이 가장 우수한 성능을 보였다. 두 가지 제안된 방법론 모두 기존 증강 기법에 비해 원본 CIFAR10 데이터셋과 유사한 이미지를 생성하며 데이터를 증강하였고, 이는 이미지 분류 성능의 향상으로 이어졌다.

본 연구의 한계점 중 하나는 CFG가 소수 클래스의 이미지를 효과적으로 학습하고 생성하는 데 성공했지만 데이터셋의 크기가 더욱 줄어들 경우 일관된 성능을 유지하는 데 어려움이 있다는 것이다. LDM의 경우 사전 학습된 모델을 활용하여 적은 클래스 이미지를 고품질로 생성할 수 있다. 그러나 실험에 사용된 데이터는 비교적 간단한 이미지(예: 개, 고양이)로 구성되어 있어 사전 학습된 LDM으로 고품질의 이미지를 효과적으로 생성할 수 있었지만, 실제 산업에서의 복잡하고 희소한 데이터를 생성하는 것에는 한계점이 존재한다. 따라서, 실제 산업에서 클래스 불균형 문제를 해결하려면 기존 모델에 관련 이미지를 추가로 학습하는 등의 작업을 고려할 수 있다. 본 연구에서 제안하는 생성 기반 분류 모델은 클래스 불균형 문제를 가진 다양한 분야에서의 확장된 연구와 실제 활용 가능한 방법으로의 적용을 기대해 볼 수 있다.

Acknowledgments

본 연구는 4단계 BK21에 의해 지원되었습니다.

References

• Q. Li, W. Cai, X. Wang, Y. Zhou, D. D. Feng, and M. Chen, "Medical image classification with convolutional neural network", 2014 13th international conference on control automation robotics & vision (ICARCV), Singapore, Dec. 2014. [https://doi.org/10.1109/ICARCV.2014.7064414]
• H. Wang, Q. Zhang, H. Lu, D. Won, and S. W. Yoon, "3D medical image classification with depthwise separable networks", Procedia Manufacturing, Vol. 39, pp. 349-356, 2019. [https://doi.org/10.1016/j.promfg.2020.01.369]
• G.-E. Kim and C.-H. Son, "Multiscale Crosss Attention Vision Transformer for Pest Image Classification", The Journal of Korean Institute of Information Technology, Vol. 21, No. 7, pp. 77-84, Jul. 2023. [https://doi.org/10.14801/jkiit.2023.21.7.77]
• K. Imoto, T. Nakai, T. Ike, K. Haruki, and Y. Sato, "A CNN-Based transfer learning method for defect classification in semiconductor manufacturing", IEEE Transactions on Semiconductor Manufacturing, Vol. 32, No. 4, pp. 455-459, Nov. 2019. [https://doi.org/10.1109/TSM.2019.2941752]
• J. Myeong, "Classify products in shopping malls based on deep learning models using image and text information", Korean Institute of Information Scientists and Engineers, pp. 681-683, Dec. 2020.
• E. Choi and K. Yoon, "Fashion Image Classification System Considering Object Aspect Ratio", Journal of Broadcast Engineering, Vol. 29, No. 1, pp. 1-9, Jan. 2024. [https://doi.org/10.5909/JBE.2024.29.1.1]
• Y. Sho, M. Son, H. Park, S. Yoon, H. Yi, and I. Yoon, "Image Classification of Insects of Rice Using Machine Learnisng", Proc. of KIIT Conference, Jeju, Korea, pp. 466-467, Jun. 2022.
• L. Perez and J. Wang, "The Effectiveness of Data Augmentation in Image Classification using Deep Learning", arXiv preprint arXiv:1712.04621, , pp. 1-8, Dec. 2017. [https://doi.org/10.48550/arXiv.1712.04621]
• D. P. Kingma and M. Welling, "Auto-Encoding Variational Bayes", arXiv preprint arXiv:1312.6114, , Dec. 2013. [https://doi.org/10.48550/arXiv.1312.6114]
• A. Creswell, T. White, V. Dumoulin, K. Arulkumaran, B. Sengupta, and A. A. Bharath, "Generative adversarial networks: An overview", IEEE signal processing magazine, Vol. 35, No. 1, pp. 53-65, Jan. 2018. [https://doi.org/10.1109/MSP.2017.2765202]
• D.-G. Kang and J. Choi, "Validation of Generative Adversarial Network Based Generated images in Class Imbalance Situation", Journal of the Korea Institute of Information and Communication Engineering, Vol. 28, No. 2, pp. 139-145, Feb. 2024. [https://doi.org/10.6109/jkiice.2024.28.2.139]
• W. Liu and W. Hu, "Oversampling of Tabular Data for Imbalanced Learning Via Denoising Diffusion Probabilistic Models", Available at SSRN 4673719, Dec. 2023. [https://doi.org/10.2139/ssrn.4673719]
• Q. Ai, P. Wang, L. He, L. Wen, L. Pan, and Z. Xu, "Generative Oversampling for Imbalanced Data via Majority-Guided VAE", arXiv preprint arXiv:2302.10910, . Feb. 2023. [https://doi.org/10.48550/arXiv.2302.10910]
• A. Majeed and S. O. Hwang, "CTGAN-MOS: Conditional Generative Adversarial Netwrok Based Minority-Class-Augmented Oversampling Scheme for Imbalanced Problems", IEEE Access, Vol. 11, pp. 85878-85899, Jun. 2023. [https://doi.org/10.1109/ACCESS.2023.3303509]
• J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli, "Deep Unsupervised Learning using Nonequilibrium Thermodynamics", In International conference on machine learning, pp. 2256-2265, PMLR, Nov. 2015. [https://doi.org/10.48550/arXiv.1503.03585]
• J. Ho and T. Salimans, "Classifier-Free Diffusion Guidance", arXiv preprint arXiv:2207.12598, , Jul. 2022. [https://doi.org/10.48550/arXiv.2207.12598]
• R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, "High-Resolution Image Synthesis with Latent Diffusion Models", In Proc. of the IEEE.CVF conference on computer vision and pattern recognition, pp. 10684-10695, Dec. 2021. [https://doi.org/10.48550/arXiv.2112.10752]
• K. He, X. Zhang, S. Ren, and J. Sun, "Deep Residual Learning for Image Recognition", In Proc. of the IEEE conference on computer vision and pattern recognition, pp. 770-778, Dec. 2015. [https://doi.org/10.48550/arXiv.1512.03385]
• H. Zhang, M. Cisse, Y. N. Dauphin, and D. Lopez-Paz, "mixup: Beyond Empirical Risk Minimization", arXiv preprint arXiv:1710.09412, , Oct. 2017. [https://doi.org/10.48550/arXiv.1710.09412]
• S. Yun, D. Han, S. J. Oh, S. Chun, J. Choe, and Y. Yoo, "CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features", In Proc. of the IEEE/CVF international conference on computer vision, pp. 6023-6032, May 2019. [https://doi.org/10.48550/arXiv.1905.04899]
• Q. Su, H. N. A. Hamed, M. A. Isa, X. Hao, and X. Dai, "A GAN-based Data Augmentation Mthod for Imbalanced Multi-class Skin Lesion Classification", IEEE Access, Vol. 12, pp. 16398-16513, Jan. 2024. [https://doi.org/10.1109/ACCESS.2024.3360215]
• S. J. Bang, M. J. Kang, M. G. Lee, and S. M. Lee, "STO-CVAE: state transition-oriented conditional variational autoencoder for data augmentation in disability classification", Complex Intelligent System, Mar. 2024. [https://doi.org/10.1007/s40747-024-01370-x]
• J. Ho, A. Jain, and P. Abbeel, "Denoising Diffusion Probabilistic Models, Advances in neural information processing systems, Vol. 33, pp. 6840-6851, Jun. 2020. [https://doi.org/10.48550/arXiv.2006.11239]
• O. Ronneberger, P. Fischer, and T. Brox, "U-Net: Convolutional Networks for Biomedical Image Segmentation", Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, Vol. 9351, pp. 234-241, May 2015. [https://doi.org/10.1007/978-3-319-24574-4_28]
• P. Dhariwal and A. Nichol, "Diffusion Models Beat GANs on Image Synthesis", Advances in neural information processing systems, Vol. 34, 8780-8794, May 2021. [https://doi.org/10.48550/arXiv.2105.05233]
• A. Nichol, P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, and M. Chen, "GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models", arXiv preprint arXiv:2112.10741, , Dec. 2021. [https://doi.org/10.48550/arXiv.2112.10741]
• C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, and M. Norouzi, "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding", Advances in Neural Information Processing Systems, Vol. 35, pp. 36479-36494, May 2022. [https://doi.org/10.48550/arXiv.2205.11487]
• A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen, "Hierarchical Text-Conditional Image Generation with CLIP Latents", arXiv preprint arXiv:2204.06125, , Apr. 2022. [https://doi.org/10.48550/arXiv.2204.06125]
• M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, "GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium", Advances in neural information processing systems, Vol 30, Jun. 2017. [https://doi.org/10.48550/arXiv.1706.08500]
• J.-C. Park, J.-U. Heu, and G.-W. Kim, "Enhancing Rebar Endpoint Detection Performance through High-Resolution Rebar Image Generation with SRPAttention-GAN", The Journal of Korean Institute of Information Technology, Vol. 21, No. 11, pp. 39-52, Nov. 2023. [https://doi.org/10.14801/jkiit.2023.21.11.39]
• H. Chang, J.-S. Lee, and J.-H. Lee, "Aleevidating Class Imbalance Problems based on Mixup-based Noise Sampling of Diffusion Generative Model", Journal of Korean Institute of Intelligent Systems, Vol. 33, No. 3, pp. 222-227, Jun. 2023. [https://doi.org/10.5391/JKIIS.2023.33.3.222]
• A. Dosovitskiy, et al., "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale", arXiv preprint arXiv:2010.11929, , Oct. 2020. [https://doi.org/10.48550/arXiv.2010.11929]
• Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie, "A ConvNet for the 2020s", In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, Jan. 2022. [https://doi.org/10.48550/arXiv.2201.03545]