Research Paper Review(D-FINE) : Redefine Regression Task in DETRs as Fine-Grained Distribution Refinement

 오늘은 D-FINE이라는 논문에서 소개한 model에 대해서 공부해보자. 개인적으로 정말 정말 잘 만든 아름다운 모델이라고 생각한다. 

https://arxiv.org/abs/2410.13842

D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement

Preliminaries 

 D-FINE 논문은 기존 DETR 구조에 box 경계를 확률분포로 예측하는 방식을 처음으로 적용한 모델이다. 왜 스칼라 기반 경계 예측보다 분포 기반 경계 예측이 더 유리한지 이해하려면 아래 글을 먼저 읽어보길 바란다. 

https://seokmin-hardstudy.tistory.com/24

How to mitigate the problem of incomplete human annotations (object detection example)

DETR 구조를 변형한 논문이므로, DETR paper도 읽어야 한다. 밑의 블로그에 정리되어있다. 

https://seokmin-hardstudy.tistory.com/20

Research Paper Review(DeTR) : End-to-End Object Detection with Transformers

Background

 box 경계를 확률분포로 예측하는 방식은 human annotation으로 인해 발생하는 불완전성을 완화하는 장점이 있기 때문에, 많은 연구자들이 DETR류 모델에도 이를 적용하려는 시도를 해왔다. 그러나 기존 DETR 구조를 그대로 유지한 채 단순히 스칼라 기반 경계 예측을 분포 기반 예측으로 바꾸면, 분포 예측 자체가 스칼라 예측보다 훨씬 어려운 문제이기 때문에 모델이 제대로 학습되지 않는 문제가 발생했다. 따라서 분포 기반 box 경계 예측이 가능하도록, 예측 난도를 낮추는 추가적인 설계가 반드시 필요했다.

Method

 저자들이 제안한 해결책은 다음과 같다. 먼저 decoder의 첫 번째 layer에서는 기존 방식대로 스칼라 값으로 box 경계를 예측한다. 그리고 이후 layer들에서는 첫 번째 layer에서 얻은 스칼라 경계를 기준점으로 삼아, 경계를 distribution 형태로 refinement하도록 설계했다. 즉, 처음부터 distribution을 직접 예측하는 어려운 문제를 던지는 대신, 스칼라 예측 → distribution 기반 refinement로 이어지는 단계적 난이도 설계를 통해 distribution 기반 box 경계 예측이 가능하도록 만든 것이다. 이제 그 구체적인 구현 방식을 살펴보자.

Method : Fine-grained Distribution Refinement (FDR)

 box 경계를 예측하는 전체 흐름을 스칼라 예측 → distribution 기반 refinement 관점에서 살펴보자. 먼저 decoder의 첫 번째 layer에서는 기존 방식과 동일하게 스칼라 값으로 box 경계를 예측한다. 이때 결과는 b0 = {x, y, W, H}이며, (x, y)는 box 중심을, W, H는 width와 height를 의미한다. 첫 번째 layer의 목표는 아주 정밀한 경계를 맞추는 것이 아니라 coarse한 초기 박스를 만드는 것이므로, top·bottom은 중심에서의 거리가 동일하다고 보고 H만 예측하고, left·right도 동일하다고 보고 W만 예측한다. 따라서 첫 layer 기준 box 중심으로부터 top, bottom, left, right까지의 거리는 각각 {H/2, -H/2, -W/2, W/2}로 표현된다.

이후 layer들에서는 box 중심 (x, y)는 그대로 고정한 채, top, bottom, left, right 네 방향의 거리를 각각 distribution 기반으로 refinement하여 박스를 점점 더 정교하게 만든다. 이제 decoder layer 1 이후의 layer에서 box 중심으로부터 네 방향까지의 거리를 어떻게 refinement하는지, 그 과정을 수식으로 살펴보도록 하자.

<refinement expression>

각 layer에서는 refinement를 적용해, box 중심 (x, y)로부터 top, bottom, left, right 방향까지의 거리를 예측한다. 이를 문자 형태로 표현하면 d_l = {top, bottom, left, right}이다. 이때 d_0는 decoder 첫 번째 layer에서 예측한 초기 거리값을 의미한다. 즉, 첫 layer의 coarse 예측이었던 {H/2, -H/2, -W/2, W/2}가 d_0에 해당한다. 이후 각 layer에서 계산되는 d_l은 이 d_0에 refinement 항을 더해 얻는 값이다.

이제 refinement 항만 따로 살펴보자. refinement는 네 방향 각각에 대해 얼마만큼 보정할지를 나타내므로, 최종적으로 {delta_top, delta_bottom, delta_left, delta_right} 형태의 4차원 vector로 구성된다.

<only refinement term>

해석을 더 명확하게 하기 위해, {H, H, W, W} 항을 sigma 함수 내부로 이동시키고 Pr_l(n) 또한 풀어서 적었다. 이때 식의 결과는 총 4개의 원소로 구성되며, 각각은 top, bottom, left, right 방향을 의미한다. 모든 연산은 element-wise로 수행되기 때문에, 한 방향에 대한 수식만 이해하면 나머지 세 방향도 동일한 방식으로 해석할 수 있다. 따라서 이하에서는 top 항만을 예시로 들어 해석하도록 하겠다.

<top-only refinement term>

 이 식은 확률분포의 기댓값을 구하는 식이다. [H x W(n)]은 distribution의 정의역(refinement할 정도 ex : -9,-8,..,0,..8,9)을 의미하고, Pr_l(n)은 그 정의역이 확률분포에서 갖는 확률값이다. 즉, Pr_l(n) 은 “refinement를 n번째 정의역만큼 적용할 가능성”을 나타내는 값이며, 모델이 각 정의역을 얼마나 타당한 refinement 후보로 판단하는지를 확률적으로 표현한 것이다. 따라서 기댓값을 구하면, refinement할 정도인 것을 알 수 있다. 

 

이제 모델을 조금 더 깊게 이해하기 위해서, W(n)은 어떻게 생겼는지와 Pr_l(n)은 어떻게 만들어내는지, 훈련할 때 refinement에 대한 loss는 어떻게 주는지를 알아보도록 하자. 

 먼저 W(n) 의 역할부터 살펴보자. 여기서 n은 1부터 N까지의 정수이며, refinement 확률분포의 정의역(domain)은 layer 1에서 예측한 H 값에 W(n)을 곱한 형태, 즉 H x W(n) 으로 구성된다.

정의역을 설계하는 가장 자연스러운 방식은 다음과 같다.

• n이 작은 값일수록 W(n) 이 음수가 되어, H x W(n) 이 음의 방향으로 최대 수준의 refinement를 의미하도록 한다.
• n이 중간 값(예: N/2) 에 도달하면 W(n) = 0, 즉 refinement가 없는 상태를 나타낸다.
• n이 큰 값일수록 W(n) 이 양수가 되어, H x W(n) 이 양의 방향으로 최대 수준의 refinement를 나타낸다.

이 논문에서 사용한 정의역 역시 이러한 구조를 그대로 따른다. 즉, W(n) 은 refinement 범위를 음수 → 0 → 양수로 대칭적이고 균일하게 펼쳐주는 스케일링 인자 역할을 하며, 모델이 “얼마나 멀리 조정(refine)”할지를 선택할 수 있는 연속적인 후보군을 형성한다.

<W(n), a = 0.5이다.>

논문에서 사용한 W(n) 을 살펴보면, 앞서 설명한 기본 구조를 따르면서도 특징적인 설계가 하나 더 있다. 바로 W(n)의 절대값이 큰 구간에서만 기울기가 크게 증가한다는 점이다.

이런 설계를 채택한 이유는 다음과 같다.

• refinement가 크게 필요한 경우에는, 정의역이 멀리 떨어진 큰 값을 포함해야 하므로, W(n)이 급격히 변화하여 큰 보정량을 표현할 수 있도록 한다.
• 반대로 refinement가 미세하게만 필요한 경우에는, W(n)이 완만하게 변하도록 하여 정의역이 촘촘하게 배치되도록 만들어, 더 정교한(fine-grained) refinement가 가능해진다.

즉, W(n)은 “큰 폭의 보정”과 “미세한 보정”을 동시에 균형 있게 다룰 수 있도록 설계된 스케일링 함수라고 볼 수 있다.

 

이제 Pr_l(n) 이 network에서 어떻게 생성되는지 살펴보자. 가장 단순한 방식은 각 decoder layer에서 query embedding을 Linear Layer에 통과시켜 logit을 만들고, 그 logit에 softmax를 적용해 확률분포 Pr_l(n) 을 얻는 것이다. 실제 논문 역시 이 기본 구조를 따른다. 하지만 중요한 차이가 하나 있다. 각 layer는 현재 layer에서 생성한 logit만 사용하지 않고, 이전 layer가 만들어둔 logit과 더한 뒤, 그 합을 softmax에 통과시켜 최종 분포를 생성한다. 즉, logit 차원에서 residual refinement를 수행하는 구조다.

이 residual 구조가 왜 필요한지 직관적으로 이해하려면 다음 비유가 도움이 된다.

학교 동아리에서 두 학생이 함께 하나의 알고리즘을 만든다고 생각해보자.

• 학생 A가 “어떻게 생각했는지에 대한 논리”만 설명하고(= query embedding을 transformer layer에 통과시켜 update하는 것에 해당), 실제 작업물(logit)은 넘겨주지 않으면, 학생 B는 처음부터 다시 작업해야 한다.
• 반면 학생 A가 논리도 설명해주고, 자신이 작성해둔 중간 작업물(logit)까지 같이 넘겨준다면, 학생 B는 그 결과물을 기반으로 더 refinement 하면서 알고리즘을 발전시킬 수 있다.

D-FINE의 refinement에서 일어나는 일도 동일하다. 각 layer는 query embedding을 update하여 “내가 지금 어떤 생각으로 refinement하고 있는지에 대한 힌트(논리)”만 다음 layer에 넘기는 것이 아니라, 이전 layer의 logit(작업물) 또한 함께 전달하여,
다음 layer가 기존 분포를 이어받아 좀 더 쉽게 다듬을 수 있도록 만든다. 즉, query embedding update는 ‘사고 과정의 전달’,
logit 누적은 ‘중간 작업물의 전달’에 해당하며, 이 두 가지가 함께 이루어지기 때문에 refinement가 layer-by-layer로 안정적이고 점진적으로 이루어진다. (글 처음에 distribution을 학습하는건 어려운 task이기 때문에 최대한 쉽게 문제를 만들어줘야 한다는 것을 reimend 하자)

<Pr_l(n)>

 

마지막으로, refinement loss가 어떻게 정의되는지 살펴보자. 가장 기본적인 접근은 각 layer에서 정답 refinement distribution과 모델이 예측한 refinement distribution 사이의 Cross-Entropy Loss를 적용하는 방식이다. 정답 refinement distribution(GT distribution)은 다음과 같이 구성한다. refinement해야 할 scalar ground truth 값과 가장 가까운 두 정의역을 선택한 뒤, 두 정의역과의 거리에 반비례하여 확률을 배정한다. 예를 들어 GT가 7.3이고 가장 가까운 두 정의역이 7과 8이라면, 정의역 7에는 70%, 정의역 8에는 30%의 확률을 주는 식이다.

그런데 이 논문에서는 한 가지 중요한 세부 사항을 추가한다.
만약 decoder의 첫 번째 layer가 예측하는 초기 box 경계값이 충분히 안정되기 전에 refinement loss를 강하게 적용하면 문제가 발생한다는 것이다. 이 경우,

• refinement가 감당해야 하는 역할이 지나치게 커져 학습이 불안정해지고,
• 첫 번째 layer가 예측하는 초기 box 경계값이 아직 제대로 자리 잡지 않았기 때문에 그 시점에 학습된 refinement 정보는 이후 정착된 후의 refinement정보와 맞지 않아 결국 폐기되어야 하며,
• 그 과정에서 모델은 불필요한 correction을 반복해야 하므로 학습이 더 길어지고 수렴 속도도 느려진다.

따라서 저자들은 초기 coarse prediction이 어느 정도 안정된 이후에 refinement loss를 적용하는 전략을 선택했다.
이렇게 하면 refinement distribution이 담당하는 문제가 훨씬 단순해지고, 모델이 각 layer에서 수행해야 하는 보정의 크기도 줄어들어 더 빠르고 안정적으로 학습을 수렴하게 된다. 이제 수식을 살펴보자. 

<refinement loss>

여기서 l 은 decoder의 layer index, k 는 query index를 의미한다. 식 안에 등장하는 두 개의 CE 항의 합은, ground truth refinement distribution과 모델이 예측한 refinement distribution 사이의 Cross-Entropy를 계산하는 부분이다. 이때 W(n←) 과 W(n→) 는 scalar ground truth 값과 가장 인접한 두 정의역을 의미한다. 그리고 w_(n←), w_(n→) 는 앞서 설명한 방식대로, scalar g.t 값(pi) 를 기준으로 두 정의역까지의 거리에 반비례하도록 할당된 확률이다. 즉, w_(n←)와 w_(n→)가 바로 g.t refinement distribution을 구성하는 두 확률 값이다. 또한 논문에서는 refinement loss를 효과적으로 적용하기 위해 IOU 기반의 weighting을 앞에 추가한다. 이는 초기 coarse prediction이 일정 수준 이상 안정된 이후에만 refinement loss가 본격적으로 영향력을 갖도록 하기 위한 설계다. 정답 box와 첫 번째 layer의 예측 box 사이의 IOU가 충분히 커졌을 때만 refinement loss의 가중치가 커지기 때문에, 학습 초기 단계에서 refinement가 지나치게 큰 책임을 떠안아 학습을 방해하는 문제를 자연스럽게 방지할 수 있다. 이러한 구조가 가능한 이유는, 첫 번째 layer에서 예측한 box와 ground truth box 사이의 별도 IoU loss가 이미 존재하기 때문이다. 즉, 초기 coarse prediction은 IoU loss를 통해 먼저 안정화시키고, 그 이후 refinement loss가 점진적으로 작동하도록 설계함으로써 학습 안정성과 수렴 속도를 모두 개선한 것이다.

 

Model structure를 그림으로 나타내면 다음과 같다. 

<model structure>

Method : Global Optimal Localization Self-Distillation 

 논문에서는 효율적인 학습을 위해 한 가지 distillation 기법을 추가로 제안한다. 바로, 마지막 layer의 refinement distribution을 이전 layer들의 refinement distribution과 KL divergence로 맞춰주는 방식, 즉 GO-LSD(Global Optimal Localization Self-Distillation)이다.

핵심 아이디어는 간단하다. 마지막 layer가 가장 많은 localization 정보를 가지고 있으므로, 그 layer에서 얻어진 refinement distribution은 가장 정답에 가깝고, 동시에 soft label 특성이 잘 살아 있는 “좋은 분포”이다. 그렇다면, 이 최종 분포를 앞쪽 layer들에게 distill해주면, shallow layer들도 더 빠르게 정확한 방향으로 수렴할 수 있게 된다. 즉, GO-LSD는 “마지막 layer가 스스로 teacher 역할을 하고, 앞쪽 layer들이 student 역할을 하는 self-distillation”이라고 이해하면 된다.

<distillation loss>

 

구체적인 과정은 다음과 같다. 먼저 모든 layer에서 각각 독립적으로 Hungarian matching을 수행해 layer별 query–ground truth 매칭을 얻는다. 이후 마지막 layer와 특정 layer의 query가 같은 ground truth에 매칭된 경우, 두 분포를 matched set으로 묶어 KL(Pr_l(n), Pr_L(n))을 적용한다. 이때 weighting으로 사용되는 값은 IoU 기반 α_k이며, 이는 해당 query가 얕은 layer와 깊은 layer 모두에서 동일한 object를 보고 있다고 확신할 수 있는 상태이기 때문에, 마지막 layer의 refinement distribution을 강하게 distill해도 안정적이라는 점을 반영한다.

그리고 D-FINE에서는 중요한 디테일을 하나 더 도입한다. 바로 unmatched query에도 distillation을 적용한다는 점이다. 일반적으로는 layer 간 ground truth 매칭이 서로 다르면 KL을 적용하지 않지만, DETR 구조에서는 1대1 매칭을 사용하기 때문에, shallow layer와 deep layer 간의 GT 매칭이 쉽게 흔들리는 assignment shake가 발생한다. 즉, 동일한 query가 shallow layer에서는 GT A에 매칭되었다가, 마지막 layer에서는 GT B에 매칭되는 일이 빈번하다. 하지만 이것이 “unmatched query가 잘못된 query”임을 의미하지는 않는다. 단지 shallow layer가 아직 representation을 충분히 정제하지 못한 상태이거나, 단순히 Hungarian matching 기준이 layer마다 조금씩 달라져 생기는 자연스러운 현상일 뿐이다.

이 때문에 unmatched query에 distillation을 전혀 적용하지 않으면, 훈련 시 distillation이 적용되는 query 수가 극도로 줄어들어 shallow layer가 마지막 layer의 “정답에 가까운 soft distribution 구조”를 학습할 기회를 잃게 된다. 이를 방지하기 위해 논문에서는 unmatched query에도 distillation을 적용하되, 그 강도를 약하게 조절한다. 이때 사용되는 weighting이 바로 classification confidence 기반 β_k이다. ( β_k안의 conf_k가 보통 α_k 안의 IOU_k보다 절대값이 작다) 

β_k안의 Conf_k는 마지막 layer에서 예측된 classification probability를 기반으로 계산된 값이다. 즉, 마지막 layer가 “이 query는 foreground object를 보고 있다”고 강하게 확신할수록 Conf_k가 높아지고, 반대로 마지막 layer가 해당 query를 background라고 판단하면 Conf_k는 낮아진다. 따라서 Conf_k는 마지막 layer 기준으로 ‘이 query가 의미 있는 object를 보고 있는가?’를 판단하는 신뢰도 지표라고 볼 수 있다.

만약 unmatched query라도 Conf_k가 높다면, 이는 마지막 layer 관점에서 “이 query는 의미 있는 foreground object를 어느 정도는 보고 있다”고 판단되는 상태이다. 이런 query는 shallow layer에서 GT assignment는 흔들렸을지라도, 마지막 layer의 refined distribution이 가진 soft 구조(분포의 shape)를 일부 전달받는 것이 유효하다. 반면 Conf_k가 낮은 unmatched query는 마지막 layer 기준에서도 object인지조차 확실하지 않은 상태이므로, distillation을 강하게 적용하면 gradient noise와 불안정성이 발생한다. β_k는 Conf_k가 낮을 때 자연스럽게 0에 가깝게 설정되므로, 이러한 query는 distillation 영향을 거의 받지 않게 된다.

 

bulle point로 정리해보자. 

• 마지막 layer와 이전 layer가 같은 object를 얘 무조건 object다 심지어 우리가 생각한 box boundary도 비슷해-> distill 세게
• 마지막 layer는 확신하는데 이전 layer는 의견이 조금 다를 때 -> distill 약하게 
• 마지막 layer도 확신이 없을 때 -> distll 안함 

 

Experiments

 

 latency와 AP socre 모두 우수한 것을 알 수 있다. 

 

Comment

 진짜 좋은 idea인 것 같다. 약간 논리 정리를 해보자면

• distribution 기반으로 학습하기 위해서 문제를 쉽게 만들자.
• deeper layer로부터 만든 distribution을 shallow layer의 distribution과 self-distillation을 함으로써 deeper layer의 지식을 shallow layer로 전달할 수 있다. 

끝!!!!!!!!