https://arxiv.org/abs/1912.03263

JEM (Joint Energy based Models)

Your Classifier Is Secretly an energy based model and you should treat it like one

이 논문에서는 Standard discriminative classifier of $p(y|\mathbf{x})$를 joint distribution $p(\mathbf{x},y)$에 대한 energy based model을 통해 재 해석함으로써 Energy based model(EBMs)를 사용하여, discriminative 문제에서의 생성모델이 가지고 있는 잠재력에 대한 해석을 진행한 논문이다.

지난 수십년동안, generative model은 downstream problem (such as semi-supervised learning, imputation of missing data, and calibration of uncertainty 등)에서 좋은 영향을 줄 것이라고 기대되어져 왔지만, 대부분의 연구는 이러한 deep generative model들을 연구할 때, 이러한 문제들은 뒷전이고 qualitative sample을뽑는 것과 validation set의 log-likelihood에 초점이 맞춰져 있었다고 함.

최근에 와서, generative model 방법론과 각각의 문제에 대한 solution들에 대한 performance gap이 굉장히 크다고 알려졌다. 하나 가능성 있는 설명은 대부분의 downstream task들이 discriminative in nature 성격을 띄고 있고, SOTA generative model은 이러한 SOTA discriminative architecture에 의해 꽤 강력하게 발산되기 때문이다. 그래서 단지 classifier만을 훈련하는 것 만으로도, generative model의 performance가 discriminative model의 performance보다 현저히 많이 떨어진다. 따라서, 잠재적인 모델의 생성능력으로 부터 얻을수 있는 이점이 discriminative performance의 감소보다 낫지 않다. 최근의 연구는 generative model의 discriminative performance를 향상시키기 위한 시도가 있다. 그런데 이러한 방법들은 여전히 순수한 discriminative를 밑돌고 있다.

이 논문은 downstream discriminative problem에서 energy based model의 사용을 지지한다. Energy based model은 downstream discriminative problems에서 생성모델의 잠재력을 깨닫게 도와준다. 지금까지 EBM으로 작업하는 작업은 어려웠지만, EBM은 다른 generative model들 보다 discriminative framework에 좀 더 자연스럽게 맞출 수 있으며, modern classifier architectures를 사용하기에 용이하다. Figure 1은 EBM 아키텍처를 소개한다. classifier의 logit을 joint density of data points와 labels, 그리고 density of data point alone을 통해 재해석한 것이다.

이 논문의 main contribution은 다음과 같다.

1. data와 label의 joint modeling을 위한 직관적이고 새로운 framework를 제안한다.
2. 이러한 모델이 이전의 SOTA hybrid models(both generative and discriminative modeling)에 비교해 outperform함을 보인다.
3. generative modeling의 결합을 통해 모델의 calibration, oodd, adversarial robustness가 향상됨을 보이고, hand-tailored method보다 여러 테스크에서 좋다는 것을 보임.

Energy based model은 $\mathbf{x}\in\mathbb{R}^D$ 에 대한 어떠한 확률 밀도 함수 $p(\mathbf{x})$에 대해서도 표현이 가능하다.

여기서 $E_\theta(\mathbf{x}):\mathbb{R}^D\rightarrow\mathbb{R}$는 energy function으로 불리며, 각각의 point를 scalar로 매핑하여준다. 그리고 $Z(\theta)=\int_\mathbf{x}\exp(-E_\theta(\mathbf{x}))$는 normalizing constant이며 partition function으로 불린다. 따라서, $\mathbf{x}$를 입력으로 받아 scalar를 내뱉는 아무 function을 사용해서 EBM으로 parametrize가 가능하다. (EBM의 특징)

대부분의 $E_\theta$를 선택할 때, $Z(\theta)$에 대해서는 계산할 수 없는데, 이것은 결국 normalized densities를 추정하거나 parameter $\theta$에 대한 standard maximum likelihood를 추정하는 방법들이 불가능한 것을 의미한다. (EBM의 특징)

일반적으로 $p_\theta(\mathbf{x})$를 $p_d(\mathbf{x})$에 근사하기 위해, Kullback-Leibler (KL) divergence between $p_\theta$ and $p_d$를 최소화 하는 방법을 사용하며, 이는 $p_\theta$를 $p$에 대한 log-likelihood의 기댓값을 maximize하는 것이다.

Stochastic gradient ascent를 사용하여 문제를 해결할 수 있기 때문에, $\theta$에 대하여 미분해주면 다음과 같다.

• $\log p_\theta(\mathbf{x})$유도
  $\log p_\theta(\mathbf{x})=\log{\exp(-E_\theta(\mathbf{x}))\over{Z(\theta)}}=\log[\exp(-E_\theta(\mathbf{x}))]-\log[Z(\theta)]\\=-\log Z(\theta)-E_\theta(\mathbf{x})$
  $\nabla_\theta \log p_\theta(\mathbf{x})=-{1\over{Z(\theta)}}\nabla_\theta Z(\theta)-\nabla_\theta E_\theta(\mathbf{x})\\=-{1\over{Z(\theta)}}\nabla_\theta\int_\mathbf{x}\exp(-E_\theta(\mathbf{x}))-\nabla_\theta E_\theta(\mathbf{x})\\={1\over{Z(\theta)}}\int_{\mathbf{x}^\prime}\exp(-E_\theta(\mathbf{x}^\prime))\nabla_\theta E_\theta(\mathbf{x}^\prime)-\nabla_\theta E_\theta(\mathbf{x})\\=\int_{\mathbf{x}^\prime}{\exp(-E_\theta(\mathbf{x}^\prime))\over{Z(\theta)}}\nabla_\theta E_\theta(\mathbf{x}^\prime) -\nabla_\theta E_\theta(\mathbf{x})\\=\mathbb{E}_{p_\theta(\mathbf{x}^\prime)}[\nabla_\theta E_\theta(\mathbf{x}^\prime)]-\nabla_\theta E_\theta(\mathbf{x})$

이 과정에서 어려운 부분은 이렇게 미분을 했을 때 $p_\theta(\mathbf{x})$에서 sample을 뽑는 것이 불가능하기 때문에, 이 gradient 추정법을 사용하기 위해서는 MCMC와 같은 샘플링 method를 사용할 수 밖에 없다. 이러한 방법은 몇몇 earliest EBM의 훈련방법으로 사용되었었다. 예를들어 Restricted Boltzmann Machines은 block Gibbs sampler를 사용해 위 expectation에 대해 추정했다.

비록 오랜기간동안 많은 발전이 있지 않았지만, 이 방법을 사용하여 large-scale EBM의 훈련에 사용하는 방법들이 최근에 등장했다. 위의 log-likelihood의 기댓값을 근사하는 성공적인 방법들 중에, Stochastic Gradient Langevin Dynamics (SGLD)가 다음 샘플링 방법을 통해 샘플링 한다.

이들이 갖는 의미는 기존 classifier를 energy based model관점에서 해석하면서 살펴보도록 하자.

What your classifier is hiding

최근 machine learning에서 분류 문제는 K개의 클래스에 대해 예측하는 것이 일반적이며, 다음과 같은 모델을 사용한다. $f_\theta:\mathbb{R}^D \rightarrow \mathbb{R}^K$ 는 데이터 포인트 $\mathbf{x}\in\mathbb{R}^D$를 $K$개의 logits(real-valued number)으로 매핑해준다. 이러한 logits는 Softmax transfer function을 통해 categorical distribution으로 parametrize할 수 있다.

여기서 $f_\theta(\mathbf{x})[y]$는 $y^\text{th}$ 인덱스에 해당하는 $f_\theta(\mathbf{x})$값을 의미한다.

논문에서 제시하는 중요한 관측은 이러한 $f_\theta$로 부터 얻은 logits에 대해 $p(\mathbf{x},y)$와 $p(\mathbf{x})$로 re-interpret 할 수 있다는 점이다. $f_\theta$를 변화하지 않고, logits를 joint distribution에 해당하는 energy based model을 정의하는데 재사용 할 수 있다.

여기서 $Z(\theta)$는 unknown normalizing constant 이며 $E_\theta(\mathbf{x},y)=-f_\theta(\mathbf{x})[y]$를 의미한다.

$y$에 대해 marginalizing을 통해, $\mathbf{x}$에 대한 unnormalized density 또한 얻을 수 있다.

이를 통해, $E_\theta(\mathbf{x})$또한 logits을 재사용하여 구할 수 있게 된다.

• $E_\theta(\mathbf{x})$ 유도
  $p_\theta(y|\mathbf{x})={\exp(f_\theta(\mathbf{x})[y])\over{\sum_{y^\prime}\exp(f_\theta(\mathbf{x})[y^\prime])}}={p_\theta(\mathbf{x},y)\over{p_\theta(\mathbf{x})}}={\exp(-E_\theta(\mathbf{x},y))\over{\exp(-E_\theta(\mathbf{x}))}}$
  $\exp(-E_\theta(\mathbf{x}))=\sum_y\exp(f_\theta(\mathbf{x})[y])\\E_\theta(\mathbf{x})=-\log\sum_y\exp(f_\theta(\mathbf{x})[y])$

임의의 scalar에 의해 바뀌는 logits $f_\theta(\mathbf{x})$이 모델에 영향을 끼치지 않는 일반적인 classifier와는 다르게, 위 framework에서 데이터 포인트 $\mathbf{x}$에 의해 바뀌는 logits은 $\log p_\theta(\mathbf{x})$에 영향을 주게된다. 무슨 말이냐면 , 예를들어, softmax의 input이 (10, 20)일때의 output과, (10+10,20+10)일때의 output은 같다. 그러나 $\log p_\theta(\mathbf{x})$는 영향을 받게 된다.

즉, 논문에서는 데이터와 label사이의 joint density $p_\theta(\mathbf{x},y)$뿐만 아니라, input example에 대한 밀도함수 $p_\theta(\mathbf{x})$를 정의하기 위해 logits안에 존재하는 숨겨진 extra 자유도를 사용하고 있다.

최종적으로 $p_\theta(y|\mathbf{x})$를 $p(\mathbf{x},y)/p_\theta(\mathbf{x})$를 통해 계산함으로써 normalizing constant를 없애고, standard Softmax parameterization을 산출해내는 것이 일반적인 standard discriminative model에 대한 내용이며, 이 논문에서는, EBM을 바탕으로 standard discriminative model사이에 있는 숨겨진 generative model을 발견해낸 것이다. 이를 JEM으로 부른다.(Joint Energy-based Model).

Optimization

새롭게 해석한 classifier architecture의 이점을 살리기 위한 optimization방법에 대한 설명.

논문에서 제시하는 모델의 $p(y|\mathbf{x})$의 parameterization은 $y$에 대해 normalized되어있고, likelihood를 maximize하여 쉽게 계산할 수 있다. 근데 EBM은 unnormalized되어 있어서 maximizing 하는 방법이 쉽지 않다. 따라서 다음의 likelihood에 대해 gradient 를 추정할 수 있을 것이다.

결국 해석해보면, 가장 첫번째 텀인 Cross entropy loss term같은 경우가 우리가 일반적으로 classification에서 사용하는 loss에 대한 gradient를 의미할 것이며 손 쉽게 계산이 가능하다(cross entropy loss). 여기에, JEM은 joint energy based modeling을 통해, generative term인 뒤에 두 텀이 추가된다는 것이 핵심적인 내용이다.

이 때, 추가된 generative term에서 두번째 텀에 해당되는 $\nabla_\theta \mathbb{E}_{p_\theta(\mathbf{x}^\prime)}[\nabla_\theta E_\theta(\mathbf{x}^\prime)]$의 경우 negative sample $\mathbf{x}^\prime$에 대한 $E_\theta(\mathbf{x}^\prime)$값을 올려주는 방향(gradient ascent)으로 gradient가 형성된다. 즉, obejctive의 gradient에서, negative sample에 $\mathbf{x}^\prime$에 대한 energy score는 높아지도록 학습이 형성된다. 반면에 뒤에 있는 텀은 실제 샘플 $\mathbf{x}$에 대해서는 energy score를 낮춰주는 방향의 gradient가 형성되며, positive sample $\mathbf{x}$에 대한 energy score가 낮아지도록 학습이 형성됨.

샘플의 관점에서 보았을 때는, 샘플은 marginal probability $p_\theta(\mathbf{x}^\prime)$을 커지는 방향인 , $\log\sum_{y}\exp(f_\theta(\mathbf{x}^\prime)[y])$의 값이 커지는 방향으로 샘플이 생성되어야 한다. 이것은 반대로, 샘플을 생성하는 gradient의 방향은 energy score가 작아지는 방향으로 학습이 된다는 의미이다. 이후 CEM에서 AT를 보면 좀 더 이해가 될 것임.

$\log p_\theta(\mathbf{x},y)$에 대해 위에서 했던 방식의 추정은 MCMC sampler를 사용했을 때, 편향된다. 저자들의 목적은 EBM training을 standard classificiation setting에 통합하고자 하는 것이다. 따라서, 이러한 이유에서, 저자들은 바로 위에 식과 같이 decompose하여 훈련하는 것을 제안하고, 이것은 일반적인 cross-entropy 를 사용하여 $p(y|\mathbf{x})$를 optimize하고 $\log p(\mathbf{x})$를 $\log\sum_y\exp f_\theta(x)[y]$에 대한 gradient로 SGLD를 진행한다.