이번에 리뷰할 논문은 MemLoTrack : Enhancing TIR Anti-UAV Tracking with Memory-Integrated Low-Rank Adapation 입니다.

---

해당 논문을 이해하기 위해서는 LoRA에 대한 이해의 초석이 필요하다.

이미 잘 알고있다면 Introduction으로 바로 넘어가길 바란다.

LoRA란 무엇이지?

자연어처리 분야에서는 일반적인 데이터로 대규모 pre-training을 진행하고, 특정 Task에 맞게 Fine-Tuning하는 것이 중요하다. 하지만 모델의 크기가 점점 커지면서, 모델의 모든 파라미터를 다시 학습시키는 전체 미세 조정(Full fine-Tuning)은 현실적으로 불가능해지고 있다. 이를 해결하기 위해 LoRA(Low-Rank Adaptation) 방식을 연구진들은 제안하는데, 이는 사전 학습된 모델의 가중치는 그대로 Freeze한 채로, Transformer 아키텍처의 각 층에 Learn-able한 *Rank Decomoposition Matrices를 삽입하는 기술이다. 이를 통해 DownStream 작업에 필요한 학습 가능한 파라미터의 수를 획기적으로 줄였다고 한다.

자연어처리를 모르는 나에게는 이런 의문이 들었다.
"그냥 애초에 특정 Task에 맞게 Pre-training을 진행하면 안되나?"
현재 자연어 처리 분야에서는 일반적인 도메인의 대규모 데이터를 사용해 Pre-training을 진행한 뒤, 특정 Task나 Domain에 맞게 모델을 Fine-Tuning을 시키는 방식이 패러다임으로 자리잡고 있다.
그 이유로는 일반적인 데이터로 대규모 Pre-training을 거친 후 특정 Task에 맞게 Fine-Tuning을 하는 것이, 처음부터 특정 작업용 데이터만 가지고 학습하는 것보다 훨씬 더 큰 성능 향상(Significant Performance Gain)을 가져오기 때문이다. 다시 말해, 모델이 방대한 데이터를 통해 언어의 문법, 일반 상식 등 전반적인 이해도를 먼저 깊게 쌓은 상태에서 특정 Task를 배우는 것이 훨씬 더 효율적이라는 것이다.
바로, 갓난 아이가 글을 먼저 배우고, 수학, 과학을 배우는 것 처럼 말이다.

그럼 대규모 모델에 약간의 수정만 하지, 왜 "Full" Fine-Tuning을?

여기에는 현실적인 이유가 존재한다.
1. 성능 극대화 : 특정 DownStream작업에서 모델의 성능을 최대한으로 끌어올리기 위해 관행적으로 모델의 모든 파라미터를 다시 학습시키는 Full Fine-tuning을 수행해왔다.
2. 기존 '일부 수정 방식'의 한계 : 당연하게도 내가 생각한 방법이 존재했었다. (Adapter, Prefix-tuning 등..) 하지만, 기존의 기술들은 파라미터 전체를 학습시키는 Full Fine-tuning의 성능을 따라잡지 못하는 경우가 많았다. 즉, 학습 효율성을 얻는 대신 모델의 품질을 희생해야 하는 뼈아픈 Trade-off 인 것이다.
3. 부가적 패널티 발생 : 기존의 '가벼운 수정' 방식들은 성능 저하 뿐 아니라, 모델의 depth를 깊어지게 만들어 추론 과정에서 inference Latency를 유발하거나, 모델이 한 번에 사용할 수 있는 문장의 최대 길이를 줄여버리는 치명적인 단점들을 동반했다.


음.. 역시 다 이유가 있다.

Rank Decomposition Matrics?

*Rank Decomposition Matrics : 모델이 특정 Task를 새로 배울 때, 가중치 행렬 값들이 조금씩 변한다. 이 변화량을 $\Delta W$ 라고 하자.
거대 모델이 새로운 지식을 배울 때 뇌의 모든 매개변수가 다 복잡하게 변하는게 아니라, 아주 핵심적인 몇 가지 방향(Low-Rank)으로만 정보가 업데이트 된다는 것을 발견했다.

• 기존 방식 : $10, 000 \times 10,000$크기의 거대한 행렬 전체를 학습시킨다.
• 행렬 분해(Rank Decomposition) : 거대한 행렬 $\Delta W$를 직접 구하지 않고, 이를 두 개의 아주 얇은 행렬 $A$와 $B$의 곱으로 쪼갠다.(Decompose)
  - 행렬 $A$ : $10,000 \times r$
  • 행렬 $B$ : $r \times 10,000$
  • 여기서 $r$(Rank)은 8이나 16같은 아주 작은 숫자이다.
• $10, 000 \times 8 + 8 \times 10,000 = 160,000$개. 즉, 학습해야 할 파라미터가 1억개에서 16만개로 줄어들었다!

논문을 들어가기 전에, LoRA에 대해 살펴본 이유는 다음과 같다.
훌륭한 드론 추적을 위해 이번 논문에서 사용한 DINOv2에 LoRA 어댑터를 연결했기 때문이다. (DINOv2 + LoRA)

그럼 "A Simple Detector with Frame Dynamics"랑 다른게 뭐지? 그것도 YOLO Freeze 아닌가?

무거운 Pre-trained 모델을 Freeze하는건 똑같지만, 그 모델을 활용하는 방식과 학습 기법이 다르다.
이전에 다룬 논문에서는 YOLO 가중치를 Feature Extract를 위해 Freeze 했다. YOLO의 output을 바탕으로 Frame Dynamics등 새로운 모듈이 학습을 진행했다.

만약 이번 논문도 이전 YOLO 떄와 똑같은 기법을 썼다면, DINOv2에서 나온 최종 이미지 Feature만 받아다가 뒤에 붙은 메모리 모듈(MAL)만 학습시켰을 것이다.
하지만, 이 논문은 DINOv2의 Backbone 내부 깊숙한 곳에서부터 작동하는 방식 자체를 '드론 추적용'으로 Fine-Tuning 했다. 단지 파라미터가 너무 많아서 LoRA를 통해 아주 적은 가중치만 학습시킨 것이다.

• 이전 연구 : 모델 고정 + "결과물"만 활용 (No Fine-tuning on Backbone)
• 이번 연구 : 모델 고정 + 그 내부에 LoRA Adapter를 삽입 -> 모델의 "인지 능력 자체를 개조 (Parameter Efficient Fine-Tuning)

Introduction

우리는 TIR (Thermal Infrared) 기반 영상에서 드론을 검출하고, 추적하는 연구를 하고 있다.
하지만, 적외선(IR) 및 열화상(TIR) 기반 영상에서는 드론은 흐릿한 점 처럼 보이는데, 기존 추적기들은 드론이 카메라 밖으로 벗어나거나 장애물에 가려지면 추적 중이던 드론을 완전히 놓쳐버리는 한계 (Local Search의 한계)가 존재했다.

때문에 최근에는 ViT 기반 모델들이 많이 나오기 시작했는데, 대표적으로 이 논문처럼 LoRA를 결합한 가볍고 빠른 모델이 등장했다. 하지만, 아무리 똑똑한 모델이라도 Temporal Memory가 없으면 장애물 뒤로 드론이 숨었을 때 다른 엉뚱한 물체를 쫒아가는 Identity Drift가 발생한다.

과거 프레임의 드론 모습을 메모리 뱅크에 저장해서 현재 모습과 비교하게 만들자 라는 아이디어가 등장한다. 이 때 중요한 점은 아무 장면이나 메모리에 저장하면 Memory Contamination, 오탐지 및 오추적 장면까지 메모리에 업데이트 될 것이다.

Memory Contamination을 막기 위해, 연구진들은 Memory Bank (MB)에 이중 게이트 구조를 설계하였다. 이 두개의 게이트 모두 통과해야지만, MB에 저장될 수 있다.

• 첫 번째 Gate : Confidence Score
  - 드론일 확률이 80% 이상일 때만 통과한다.
• 두 번째 Gate : Kalman-Filter 기반 Motion Consistency Test
  - 과거의 속도와 방향을 통해 다음 프레임에 드론이 위치할 좌표를 미리 예측한다.

이러한 이중 게이트에 대한 설명은 밑에서 더 자세히 알아보겠다.
Training이 아닌 Test, Inference 시에는 MAL에서 현재의 Search Region Token을 Memory Token과 Cross-Attention하여 최소한의 Computational Overhead로 이러한 드리프트를 완화한다.

연구진들은 해당 연구에 기여한 주요 특징들을 다음과 같이 요약했다.

Memory-Integrated PEFT Tracker

LoRA 기반의 Fine-tuning 효율성을 유지하면서도 시간적 문맥을 활용하는 MAL을 구조 중 하나로 갖는 ViT 기반 아키텍처이다.

PEFT(Parameter-Efficient Fine-Tuning)은 쉽게 말해서 거대한 딥러닝 모델을 다루는 방법론 이라고 생각하면 되겠다. LoRA(Low-Rank Adaptation)은 그 방법론에 속하는 하나의 테크닉이라고 생각하면 되겠다.

> PEFT : 수백억 파라미터를 가진 거대 모델을 처음부터 다시 학습(Full Fine-Tuning)하려면 막대한 메모리와 시간이 필요한데, 이를 해결하기 위해 기존 모델의 가중치는 얼려두고(Freeze), 아주 적은 비율로 추가 파라미터만 훈련시키자는 방법론을 통칭하는 단어이다.

추후 구현 리뷰에서도 보겠지만, Parameter를 Freeze 한다는 것은 다음과 같다.

# PEFT 방법론의 적용 (백본 모델로는 DINOv2 사용)
for param in backbone.parameters():
	param.requires_grad = False
    
# LoRA 기법의 적용
# 얼려버린 기존의 qkv, fc1, fc2 옆에 학습 가능한 작은 행렬 (LoRA_A, LoRA_B)를 덧붙인 것이다.
original_qkv = block.attn.qkv
block.attn.qkv = LoRALinear(original_qkv, config=self.config)

위와 같이 사용하는데.. qkv라는 이름으로 트랜스포머의 인자(?)가 뭉쳐진 것이 의아할 것이다.

DINOv2의 Transformer block은 아래와 같이 생겼다.
Block
├─ LayerNorm
├─ Attention
│            ├─ qkv (Linear)
│            └─ proj (Linear)
├─ LayerNorm
└─ MLP
           ├─ fc1
           └─ fc2
Transformer의 핵심 연산은 Self-Attention으로, 입력 토큰을 Q, K, V 이렇게 세 가지로 변환해서 계산하는건 다 알 것이다.
$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

• Q : 현재 토큰이 찾고자 하는 정보
• K : 다른 토큰이 가진 특징
• V : 실제로 전달될 정보
  그건 그렇다 쳐도 도대체 왜 3개를 따로따로 만들지 않고 하나의 Linear Layer로 묶냐는 것이다.

qkv = Linear(embed_dim, 3*embed_dim)

DINOv2에서는 내부 구현체가 연산 최적화를 위해 Q, K, V 가중치 행렬을 하나의 거대한 행렬로 병합해서 사용하기 때문이다.
거의 모든 트랜스포머 모델들이 이렇게 QKV 통합 레이어 방식을 사용하는데, 이는 GPU 커널 실행 시 오버헤드가 감소하고 행렬곱셈의 병렬 처리를 극대화 할 수 있다.

논문에서는 Q,K,V를 독립적인 가중치 행렬 $W_q, W_k, W_v$로 설명하지만, 실제 PyTorch 구현에서는 GPU의 병렬 연산 속도를 극대화하고 커널 실행 오버헤드를 줄이기 위해 $W_{qkv}$라는 하나의 거대한 Linear Layer로 병합하여 한 번에 계산한 뒤, 결과 텐서를 3등분 하여 사용한다고 알고 있으면 되겠다.

DINOv2의 내부 구현체가 연산 최적화를 위해 Q, K, V 가중치 행렬을 하나의 거대한 행렬로 병합시켜서 사용하기 때문에 LoRA 연산 시에 이렇게 병합된 행렬 구조에 맞게 설계한 것이다.

Dual-Gated Memory Write Polity

추론 단계에서 FIFO Memory Bank에는 Confidence Score와 Kalman-Filter 기반 Motion Consistency 모두 통과해야 들어갈 수 있다. 이를 통해 OC, OV, DBC 상황에서 Memory Contamination을 완화할 수 있게 된다.

State-of-the-Art Robustness on Anti-UAV410

MemLoTrack은 Anti-UAV410에서 SOTA를 달성하였다. FocusTrack이 더 우수한 정밀도 기반 지표를 보여주지만, Robustness 측면에서 더 주목할 만하다는 것이다.

Model	Size	AUC	$P$	$P_{norm}$	SA
MemLoTrack	224	63.6 / 66.9	82.7 / 86.7	79.8 / 66.6	64
FocusTrack	256	62.8	86.2	82.8	63.9
LoRAT-B	224	62.2 / 66.4	80.9 / 85.6	78.1 / 64.8	62.5

---

Related Work

2.1. Transformer-Based SOT and Parameter-Efficient Adaptation

One-stream ViT 기반 Tracker와 LoRA 기반 Fine-tuning. 현대의 단일 객체 추적에서는 Transformer 가 백본 모델로 표준이 되었다.
이는 Self-Attention 덕분인데, 장거리 종속성 모델링과 Template-Search Token의 조기 융합을 가능하게 하기 때문이다.

유사한 연구로 LoRAT은 백본 모델로 ViT를 고정하고, Attention block에 LoRA를 삽입하며, Template과 Search Region 토큰 구별하기 위한 Token-Type Embedding을 도입하였다. PEFT 하에 Inductive bias를 피하기 위해 Convolutional Head를 MLP로 대체하였다.

2.2. Anti-UAV Tracking in TIR : Benchmarks, Protocols, and Specialized Designs

시각 기반 Anti-UAV 추적의 주요 벤치마크는 RGB-T를 다루는 Anti-UAV300 데이터셋과 TIR 전용으로 Anti-UAV410이 있다.

이러한 데이터셋에는 드론이 화면에서 갑자기 사라졌다가 나타나는 특성을 반영하기 위해 IoU 뿐 아니라 타겟의 존재 여부를 맞추는 능력을 포함하는 SA를 주요 평가 지표로 사용한다.

그 외에도 전반적인 추적 성공률을 나타내는 Sucess AUC, 중심점 위치의 정밀도를 보는 $P / P_{norm}$, 가려짐(OC), 시야 이타 (OV), 동적 배경 간섭 (DBC) 등 다양한 악조건 상황 속에서 세부 성능을 평가한다.

RGB 카메라는 타겟의 질감과 같은 세밀한 시각적 특징을 포착할 수 있겠지만, 조명 변화나 악천후에 매우 취약하다.
반면 열화상(TIR)은 악천후에 강하지만, 드론이 질감이 전혀 없는 흐릿한 텍스쳐로 보인다는 치명적인 단점이 있다.

특히 드론의 주변 환경과 열 패턴이 비슷하여 섞여버리는 열 교차(Thermal Crossover) 현상이 잦다. 따라서 TIR 환경에서는 '생김새' 하나로만 의존해서는 안되며, 시간의 흐름에 따라 누적된 과거의 증거 (temporal context)와 대상의 물리적 움직임 단서를 결합해야만 대상을 놓치지 않을 수 있다.라고 논문은 말한다.

MemLoTrack은 ViT의 가중치를 건드리지 않고 Self-Attention Block에만 LoRA를 적용하는 PEFT 방식의 One-stream Tracker의 기본 뼈대를 삼는다. (물론 ViT, Backbone으로 DINOv2)

Feature를 추출하는 Encoder와 최종 결과를 내는 Prediction Head 사이에 명시적인 Memory Attention Layer (MAL)을 새롭게 추가 하여 과거의 정보와 현재를 연결하도록 설계하였다.

추론 시에 사용할 메모리 시스템을 훈련 과정에서도 자연스럽게 학습할 수 있도록 하였는데, 학습 시에는 Template Frame과 현재의 Search Region Frame 사이의 시간대에서 정확히 7장의 메모리 프레임을 샘플링하여 최적화를 진행한다.

추론 시에는 FIFO 방식으로 MB를 실제 가동하여 토큰들을 저장하고 이 MB가 MAL 모듈에 Cross-Attention을 위한 K, V를 공급한다.

---

Method

MemLoTrack은 Template $z$와 Search Region Frame $x_t$를 공동으로 인코딩하는 One-stream ViT Tracker이다.
(Template Frame은 Background와 Foreground로 나뉘는데, Background는 템플릿 이미지 안에서 목표물이 아닌 주변 배경에 해당하는 패치들을 의미하며, Foreground는 실제 목표물이 포함된 영역의 패치들을 의미한다.)

Backbone 모델로는 DINOv2 ViT (패치 크기 : 14, Embedding = 768)을 사용한다.
논문에는 다음과 같이 나와있지만, 혼동하지 않도록 주의하자.

LoRA adapters (rank $r = 64$) into all linear projectinos of attention and MLP blocks.

LoRA는 Transformer Encoder 블록에만 삽입되며, 초기 패치 임베딩을 담당하는 Linear Projection 블록에는 들어가지 않는다.

Figure 1 맨 아래에 파란색 블록인 Linear pProjection은 이미지를 14 x 14 패치로 잘라 초기 토큰을 만드는 특정한 계층의 이름으로, DINOv2의 가중치를 그대로 사용하여 완전히 Frozen된 상태이다.

논문에서 말한 Linear Projections는 특정 블록 이름이 아니라 Encoder 내부에서 일어나는 수학적인 선형 투영 연산 자체를 뜻한다.

Transformer Encoder의 내부를 뜯어보면 다음과 같은 연산들이 있다.

• Attention Block 내부 (MSA) : 입력된 데이터를 Q, K, V로 변환하기 위해 가중치 행렬을 곱하는데, 이 과정이 바로 Linear Projection 연산이다. LoRA는 바로 여기에 삽입된다.

• MLP 블록 내부 : 신경망 층 사이에서 차원을 늘리거나 줄이는 과정 역시 Linear Projection 연산이며, 여기에도 마찬가지로 LoRA가 삽입왼다.

결론적으로 LoRA 어댑터는 인코더 바깥에 있는 초기 Linear Projection 블록이 아닌, Transformer Encoder 내부의 Attention과 MLP 내부의 Linear Projection 연산 과정에만 삽입된다.

모델 Training 시에느느 $(z, x_t)$를 공동으로 인코딩하여 Search Query $X_t$를 얻는다. 그런 다음 $z$와 $x_t$ 사이에 있는 7개의 메모리 프레임을 샘프링한다. (샘플링의 상세 과정은 이후에 다루도록 한다.)

샘프링된 각 프레임은 Search Template과 마찬가지로 동일한 패치 임베딩, grid에 대한 위치와 메모리 타입 임베딩 또한 거친다.

이러한 Memory Token들은 인코더를 우회하고 MAL의 K값으로 쌓이는 반면 $X_t$는 쿼리로 사용된다. Training에는 MB나 Gate가 사용되지 않는다. 헤드는 Search Region 의 토큰 그리드에서 작동하며, 각 셀에 대해 Classification Map과 네 개의 거리를 생성한다. (이 또한 나중에 상세히 다루도록 하겠다.)

3.2. Inputs and Tokens

각 Training insatnce는 Template $z \in R^{3\times 112\times 112}$,Search Frame $x_t \in R^{3 \times 224 \times 224}$그리고 Search 해상도 $L_{mem}$Memory Frame으로 구성된다.

Learnable한 토큰 유형 임베딩으로는 $e_t, e_b, e_s, e_m$을 사용한다.

Memory Frame은 Search Frame과 동일한 해상도를 가지며, 3중 임베딩 시스템으로 Final Input Token이 만들어진다.

• PE (Patch Embedding)
• P(i, j) : Positional Embedding
• Token Type Embedding : Transformer가 각 토큰 출신을 확인한다.
  
  
  Token Type Embedding은 입력되는 이미지들이 템플릿(배경)인지, Search Region인지, 과거 메모리 프레임인지 구분하기 위한 식별표이며, Global Redetection이 아닌 Restriced Search Region 내에서만 Local하게 추적을 수행하기 때문에 현재 프레임에서 목표가 어디로 이동했는지 확인하기 위한 Search Frame이라고 이해하면 된다.

3중 임베딩 시스템으로 만들어진 Final Input Tokens 각각은 입력된 이미지를 14x14 픽셀 크기의 Patch로 자른 뒤에 만들어진 것인데, 2차원 그리드 형태로 배열이 된 이 개별 토큰들을 Transformer에 넣기 위해 1차원으로 Flattening하여 Template Sequence와 Search Sequence를 만든다.

지금까지 설명한 Training pipeline은 위와 같다.

Encoder output

$T_z$와 $T_{xt}$를 연결하고 (각각 템플릿 시퀸스와 Search Region Sequence) 인코더를 통과시키면 MAL에 의해 Search Query $X_t$가 생성된다. (B x $N_x$ x d)

3.3. Training-Time Memory Frame Sampling

Training 프로세스를 Inference시 메모리 메커니즘과 일치시키기 위해 Streaming 기반 Memory Sampling을 어떻게 하는지 보여준다.

각 Training instance에 대해 시간 $t_z$에서의 Template Frame $z$와 시간 $t_x$ 사이의 시간 간격을 기반으로 Memory Frame 집합을 설정한다.

구체적인 과정은 아래 그림을 보며 이해하도록 하자.

샘플링 로직은 $L_{\text mem}$에 대한 구간 길에 따라 두 가지로 나뉜다. (논문에서 구현한 $L_{\text mem}$의 값은 7이다.)

1) 구간이 충분할 때
$(|t_x - t_z| > L_{\text mem})$ : 다양한 시간적 맥락 정보를 얻기 위해 중간 기간에서 $L_{\text mem}$개의 고유한 프레임을 무작위로 추출한다.
2) 구간이 부족할 때
$(|t_x - t_z| \leq L_{\text mem})$ : 사용 가능한 모든 중간 프레임을 활용하고 남은 슬롯은 마지막 유효 프레임을 복제하여 채운다.

이 프레임들은 현재 미니 배치에서만 사용되고, 시간 경과에 따른 영구 버퍼나 배치 간에 저장되지는 않는다.

추론 중에는 MemLoTrack이 이중 게이트가 만족될 때마다 시간 경과에 따라 임베딩 후 토큰 $E_t$를 누적하는 $L_{\text mem}$ 크기의 FIFO 메모리 뱅크를 유지한다.

3.4. Memory Attention Layer (MAL)

MAL은 검색 쿼리 $X_t$로부터 쌓인 메모리 토큰 $M_t$로의 cross-attention을 수행하고 메모리 조건부 표현을 출력한다.

$\text{Attn}(X,M) = \text{softmax}(\frac{(XW^Q)(MW^K)^\top}{\sqrt d_h})(MW^V)W_O$

예를 들어 시퀸스의 시작 부분에서는 메모리 스택이 비어있는 경우, MAL이 우회되고 쿼리 토큰 $X_t$가 직접 MLP 헤드로 전달된다.

3.5. Inference-Time Memory Bank(MB) with Dual Gating

누누이 말하던 Dual Gating에 대해서 소개한다.
Inference에서는 크기 $L_{\text mem}$짜리 FIFO 메모리 뱅크를 유지한다. 프레임 $t$를 처리하고, 현재 항목을 쓸지 말지(MB에 넣을지 말지) 결정하기 위해 두 개의 게이트가 적용된다.

1. Confidence Gate
Confidence Score가 특정 임계값을 초과하면 통과
2. Motion Consistency Gate
Predicted BBox와 Decoded BBox 사이의 마할라노비스 거리가 카이제곱 임계값보다 작으면 통과

위 두 게이트를 모두 통과하면 해당 토큰 $E_t$는 먼저 Affine 변환을 거쳐 현재 쿼리 그리드에 공간적으로 정렬된 후 MB에 추가된다.

이 정렬은 메모리 토큰과 현재 쿼리가 표준 좌표계를 공유하도록 보장한다.
$E_t = \text{Align}_t(PE(x_t)) + P_{16\times 16} + e_{\text{mem}} \in \mathbb{R}^{B \times N_x \times d}$

메모리 뱅크가 가득 차면, 가장 이전에 저장된 값이 제거된다. (FIFO...).

3.6. Loss and Optimization

MemLoTrack의 Learnable한 요소는 LoRA Adapter, Memory Attention Layer(MAL), Prediction Head는 e2e 방식으로 공동으로 최적화 된다.

총 손실 함수 $L$은 분류를 위한 BCE 항과 BBox 회귀를 위해 일반화된 GIoU 항의 가중 합으로 공식화 된다.

$L = \lambda_{\text{cls}} \cdot \text{BCE}(S_t, S_t^*) + \lambda_\text{box} \cdot (1-\text {GIoU}(b_t, B_t^*))$

Training과 Inference 모두 Template과 Search Region의 크기는 각각 $z = 112 \times 112$, $x = 224\times224$로 구성된다.

3.7. Prediction Head and Targets

Head는 인코더에 의해 생성된 $16\times16$ 토큰 그리드에 대해 각 토큰별로 적용되는 경량 MLP이다. (각 토큰은 $224 \times 224$ Search window 내 $14 \times 14$ 픽셀 크기에 해당한다.)

각 그리드 셀에 대해 Classification Score $S_t$와 네 개의 거리 ($l, t, r, b$)를 출력하는데, BBox는 이러한 거리를 셀 중심 주변 사각형으로 변환한다.

$x_L = c_x - l, x_R = c_x + r, y_T = c_y - t, y_B = c_y$

Results
작성 중...

---

Implementation with PyTorch

Network

class DINOv2(nn.Module):
    def __init__(self, type_embedder, config):
        super().__init__()
        
        self.backbone = backbone
        for param in backbone.parameters():
            param.requires_grad = False  # Backbone Freezing

        self.type_embedder = type_embedder
        self.type_embedding = TypeEmbedding()
        self.config = config

        for block in self.backbone.blocks:
            block.attn.qkv = LoRALinear(block.attn.qkv, config = self.config)
            block.mlp.fc1 = LoRALinear(block.mlp.fc1, config = self.config)
            block.mlp.fc2 = LoRALinear(block.mlp.fc2, config = self.config)
        self.mal = MAL(embed_dim=768, num_heads=12)
        self.head = PredictionHead(in_channels=768, hidden_channels=192)

    def get_initial_tokens(self, x):
        # Patch embedding
        tokens = self.backbone.patch_embed(x)
        
        # Positional Embeding
        pos_embed = self.backbone.pos_embed[:, 1:, :]
        output = tokens + pos_embed
        
        return output

class LoRALinear(nn.Module):
    # LORA Adapter는 Trainable하게 !! 
    def __init__(self, original_layer, config):
        super().__init__()

        if isinstance(original_layer, LoRALinear):
            original_layer = original_layer.original_layer

        #self.original_qkv = original_qkv
        self.config = config

        self.original_layer = original_layer
        self.in_features = original_layer.in_features
        self.out_features = original_layer.out_features
        self.rank = config['r']
        self.lora_alpha = config['lora_alpha']
        self.lora_dropout = nn.Dropout(config['lora_dropout'])
        self.bias = config['bias']
        self.target_modules = config['target_modules']
        self.scaling = self.lora_alpha / self.rank

        in_features = self.original_layer.in_features
        out_features = self.original_layer.out_features

        self.LoRA_A = nn.Parameter(torch.zeros(in_features, self.rank))
        self.LoRA_B = nn.Parameter(torch.randn(self.rank, out_features))

        nn.init.kaiming_uniform_(self.LoRA_A, a=math.sqrt(5))
        nn.init.zeros_(self.LoRA_B)

class PredictionHead(nn.Module): # [Batch, 768, 16, 16] -> [Batch, 192, 16, 16]
    def __init__(self, in_channels=768, hidden_channels=192): # 256으로도 해보기
        super().__init__()
        self.ClassificationMLP = nn.Sequential(
            nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(hidden_channels),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(hidden_channels, 1, kernel_size=3, padding=1), # [B, 192, 16, 16] => [B, 1, 16, 16] 1채널만 필요함 (점수값)
            nn.Sigmoid()
        )

        self.RegressionMLP = nn.Sequential(
            nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(hidden_channels),
            nn.ReLU(inplace=True),

            nn.Conv2d(hidden_channels, 4, kernel_size=3, padding=1), # 4방향 Output
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        B, N, C = x.shape
        H = W = int(math.sqrt(N))
        assert H * W == N, "Search tokens must be a square grid"

        x = x.transpose(1, 2).contiguous().view(B, C, H, W)
        score_map = self.ClassificationMLP(x) # Output
        size_map = self.RegressionMLP(x) 

        return score_map, size_map

class KalmanGate(object):
    '''
    https://www.geeksforgeeks.org/python/kalman-filter-in-python/
    참고하였음.
    '''
    # 1. 드론의 현재 상태 벡터로 정의 ->6차원 벡터 
    # 2. Prediction Head가 측정한 z_t
    def __init__(self, initial_bbox):
        #kf = KalmanFilter(dim_x=3, dim_z=1)
        self.kf = KalmanFilter(dim_x = 8, dim_z = 4)

        self.kf.F = np.array([
            [1, 0, 0, 0, 1, 0, 0, 0], # cx_new = cx + vx
            [0, 1, 0, 0, 0, 1, 0, 0], # cy_new = cy + vy
            [0, 0, 1, 0, 0, 0, 1, 0], # w_new = w + vw
            [0, 0, 0, 1, 0, 0, 0, 1], # h_new = h + vh
            [0, 0, 0, 0, 1, 0, 0, 0], # vx_new = vx (등속도 운동 가정)
            [0, 0, 0, 0, 0, 1, 0, 0], 
            [0, 0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 0, 1]
        ])

        self.kf.H = np.array([
            [1, 0, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0]
        ])
        # x(x0) : initial state estimate
        self.kf.x[:4, 0] = initial_bbox

        # Initial Error Covariance (Pos)
        self.kf.P *= 10.0
        
        # Initial Error Covariance (Velocity)
        self.kf.P[4:, 4:] *= 1000.0

        # Measurement Noise Covariance
        self.kf.R *= 10.0

        # Process NOise Covariance
        self.kf.Q *= 0.01


    def predict(self):
        self.kf.predict()
        return self.kf.x
    
    def update(self, z):
        S = np.dot(self.H, np.dot(self.P, self.H.T)) + self.R
        K = np.dot(np.dot(self.P, self.H.T), np.linalg.inv(S))
        y = z - np.dot(self.H, self.x)
        self.x = self.x + np.dot(K, y)
        I = np.eye(self.P.shape[0])
        self.P = np.dot(I - np.dot(K, self.H), self.P)
        return self.x
    
    def gate(self, z_t, threshold=9.49):
        '''
        Threshold 값 다시 정해야 함. 
        마할라노비스 거리 기준이니까 신뢰구간 구하기
        '''
        self.kf.predict() # kalman으로 업데이트

        z = np.array(z_t).reshape(4,1)
        y = z - np.dot(self.kf.H, self.kf.x)

        S = np.dot(self.kf.H, np.dot(self.kf.P, self.kf.H.T)) + self.kf.R
        S_inv = scipy.linalg.inv(S)
        mahalanovis_dist = np.dot(y.T, np.dot(S_inv, y))[0, 0]

        is_passed = (mahalanovis_dist < threshold)

        if is_passed: # 칼만 내부 상태 업데이트로 속도 학습하기
            self.kf.update(z)

        return is_passed, mahalanovis_dist

class MemorySampling(nn.Module): # Only During Training
    def __init__(self, max_size=7):
        super().__init__()
        self.max_size = max_size

    def forward(self, template_idx, search_idx, frame):
        candidates = frame[template_idx + 1 : search_idx]
        num_candidatees = len(candidates)

        sampled_memory = []

        if num_candidatees >= self.max_size:
            sampled_memory = random.sample(candidates, self.max_size)

        elif num_candidatees > 0:
            sampled_memory = list(candidates)
            empty_size = self.max_size - len(sampled_memory)
            last_frame = candidates[-1]
            sampled_memory.extend([last_frame] * empty_size)
        
        else:  # No Interval : Full with Template frame
            template_frame = frame[template_idx]
            sampled_memory = [template_frame] * self.max_size

        return sampled_memory

class MAL(nn.Module):
    def __init__(self, embed_dim = 768, num_heads=12):
        super().__init__()
        self.cross_attention = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.MLP = nn.Sequential( # Ordinary Transformer's MLP (FFN) : embed_dim * 4
            nn.Linear(embed_dim, embed_dim*4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim)
        )
    def forward(self, search_token, memory_token):
        if memory_token is None or memory_token.size(1) == 0: # First frame : just return
            return search_token
        
        attention_output, _ = self.cross_attention(query = search_token, key=memory_token, value = memory_token)
        x = self.norm1(search_token + attention_output) # Serach token이 Encoder에서 나온거
        MLP_output = self.MLP(x)
        final_ouptut = self.norm2(x + MLP_output)

        return final_ouptut

class TrackerDINOv2():
    def __init__(self, initial_bbox):
        super().__init__()

        # Model initialization
        self.net = DINOv2(TypeEmbedding(), LoRAConfig)
        self.type_embedder = TypeEmbedding()
        self.template = None
        self.kalman = KalmanGate(initial_bbox)
        self.memory_bank = MemoryBank(kalman_gate=self.kalman)
        self.transform = T.Compose([T.ToPILImage(), T.ToTensor(), T.Normalize(mean = [0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

    def init_with_bbox(self, frame, bbox):
        self.template = self.extract_template(frame, bbox)
        self.prev_bbox = bbox

    def extract_template(self, frame, bbox):
        x, y, w, h = bbox
        cx = x + w / 2
        cy = y + h / 2
        s_z = SiamFC_crop_size(w, h)
        template = self.SiamFC_crop(frame, cx, cy, s_z, 112)

        template = self.transform(template).unsqueeze(0)

        return template

    def extract_search(self, frame, prev_bbox):
        x, y, w, h = self.prev_bbox
        cx = x + w / 2
        cy = y + h / 2

        s_z = SiamFC_crop_size(w, h, context=0.5)
        s_x = s_z * 2.0 # 2배 해서 면적 4배

        patch = self.SiamFC_crop(frame, cx, cy, s_x, out_size=224)
        search_tensor = self.transform(patch).unsqueeze(0)
        return search_tensor

    def SiamFC_crop(self, img, cx, cy, size, out_size):
        # Tracking에서 사용할 template / Search patch를 이미지에서 추출
        # cv2.getRectSubPix() : 중심 좌표 기준으로 sub-pixel crop을 수행하는 함수
        # 일반 crop은 center 기준 crop이 어려움. -> image boundary 처리 필요함

        patch = cv2.getRectSubPix(
            img,
            (int(size), int(size)),
            (cx, cy)
        ) # bbox 중심 기준으로 crop
        patch = cv2.resize(patch, (out_size, out_size)) # 모델 입력 크기로 resize

        return patch
    
    def prediction_2_Box(self, score_map, regression_map):
        B, _, H, W = score_map.shape
        assert B == 1, "PredictionHead의 output은 반드시 1개의 배치만 있어야 함."
        flatten_scoreMap = score_map.view(B, -1)
        conf, idx = flatten_scoreMap.max(dim=1)
        
        max_score_pos = int(idx[0].item())
        confidence_score = float(conf[0].item())

        grid_x = max_score_pos % W
        grid_y = max_score_pos // W
        

        '''
        16 x 16에서 생각해보면 0~15 idx까지 row 인것 생각하기  
         xxxxxxxxxxxxxxxxx
        + 패치가 16x16 사이즈인 것이랑, 실제 원본 이미지 224x224 차이 생각하기
        '''
        grid_x =  (max_score_pos % 16) 
        grid_y =  (max_score_pos // 16)
        
        # Regression Map에서 4방향 거리를 뽑아오기
        # Regression_map 의 shape : [B, 4, 16, 16]
        l = float(regression_map[0, 0, grid_y, grid_x].item())
        t = float(regression_map[0, 1, grid_y, grid_x].item())
        r = float(regression_map[0, 2, grid_y, grid_x].item())
        b = float(regression_map[0, 3, grid_y, grid_x].item())

        stride = 14.0 # Patch Size
        w = (l + r) * 14
        h = (t + b) * 14

        cx = grid_x * 14    
        cy = grid_y * 14

        org_bbox = [cx, cy, w, h]
        return confidence_score, org_bbox

def SiamFC_crop_size(w, h, context=0.5):
    p = (w + h) / 2 * context # contect padding : 배경 정보를 얻기 위함이며, appearance와 scale 변화에 강해지기 위해
    s = np.sqrt((w+p) * (h+p))
    
    return s

Execution Codes

def main():
    cfg = TrainConfig()
    os.makedirs(cfg.save_dir, exist_ok=True)
    set_seed(cfg.seed)

    wandb.init(
        project = "MemLoTrack-DINOv2_train",
        name="Experiment_01",
        config=cfg.__dict__
    )

    device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
    print(f"현재 사용 장치 :{device}")

    model = DINOv2(TypeEmbedding(), LoRAConfig)
    model.to(device)

    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable_params, lr=cfg.lr, weight_decay=cfg.weight_decay)

    train_dataset = AntiUAVDataset(cfg.dataset_root, split=cfg.split)
    val_dataset = AntiUAVDataset(cfg.dataset_root, split='val')

    train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, collate_fn=collate_fn)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)

    best_loss = float('inf')
    best_iou = 0.0

    for epoch in range(cfg.epochs):
        print(f"\n---[Epoch {epoch+1}/{cfg.epochs}] Train ---")
        train_loss = train(model, train_loader, optimizer, device, cfg)
        
        print(f"---[Epoch {epoch+1}/{cfg.epochs}] Validation ---")
        val_loss, val_iou = validate(model, val_loader, device, cfg)
        
        wandb.log({
            "Train Loss": train_loss,
            "Val Loss": val_loss,
            "Val IoU": val_iou
        })

        if val_iou >= best_iou:
            print(f"New Best IoU Recorded : {best_iou:.4f} -> {val_iou:.4f}")
            best_iou = val_iou

            ckpt_path = os.path.join(cfg.save_dir, "best_MemLoTrack.pt")
            torch.save({
                "epoch": epoch+1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "best_iou": best_iou
            }, ckpt_path)
            
    wandb.finish()

if __name__ == "__main__":
    main()

더욱 자세한 코드가 궁금하다면?

깃허브 참고해주세요~!