기술 요약: 손실의 정렬: 효율적인 사전 학습을 위한 그래디언트 기반 복합 손실 가중치

문제 정의

현대 딥러닝 시스템은 종종 대규모의 라벨이 없는 데이터를 이용한 사전 학습 (pretraining) 과 이를 하위 작업 (downstream tasks) 에 적응시키는 2 단계 패러다임을 채택합니다. 많은 사전 학습 파이프라인에서 목적 함수는 단일 손실이 아니라 다음 토큰 예측, 대비 손실 (contrastive losses), 또는 회귀 목표와 같은 여러 손실 항의 가중 합으로 구성됩니다. 이러한 항들의 상대적 가중치는 학습된 표현의 품질과 하위 작업 성능에 지대한 영향을 미치는 중요한 하이퍼파라미터입니다.

현재 이러한 가중치 선택은 수동 조정, 무작위 탐색, 또는 베이지안 최적화에 의존합니다. 이러한 방법들은 수많은 독립적인 학습 실행을 필요로 하므로 대규모 사전 학습에는 계산적으로 불가능합니다. 또한, 균일한 가중치 부여와 같은 단순한 기본값은 서로 다른 손실 항 간의 다양한 기여도와 충돌을 고려하지 못해 종종 실패합니다. 본 논문은 반복적인 전체 학습 실행의 비용을 치르지 않고 하위 작업 성능을 극대화하기 위해 복합 사전 학습 목표의 선형 손실 가중치를 조정할 수 있는 효율적인 방법의 필요성을 다룹니다.

방법론: 그래디언트 정렬 사전 학습 (GraP)

저자들은 온라인으로 손실 가중치를 학습하는 그래디언트 기반 2 단계 최적화 (bilevel optimization) 방법인 그래디언트 정렬 사전 학습 (Gradient-aligned Pretraining, GraP) 을 제안합니다. 핵심 아이디어는 복합 사전 학습 그래디언트를 하위 목표의 그래디언트와 정렬시키는 것입니다.

1. 2 단계 최적화 공식화

문제는 바깥쪽 루프가 하위 손실 $L_{down}$을 최소화하기 위해 손실 가중치 $w$를 최적화하고, 안쪽 루프가 복합 사전 학습 손실 $L_{pre}(w)$를 최소화하기 위해 모델 파라미터 $\theta$를 최적화하는 2 단계 최적화 문제로 설정됩니다.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{subject to} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. 그래디언트 기반 업데이트 규칙

단일 단계 SGD 와 같은 단절된 역전파 (truncated backpropagation) 접근법을 사용하여 저자들은 가중치에 대한 업데이트 규칙을 유도합니다. 가중치에 대한 하위 손실의 그래디언트는 하위 그래디언트와 개별 사전 학습 그래디언트 간의 내적 (inner product) 에 비례합니다.
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
여기서 $g_{down}$은 하위 그래디언트이고 $g_k$는 $k$번째 사전 학습 손실의 그래디언트입니다. 목표는 복합 그래디언트 $\sum w_k g_k$와 하위 그래디언트 간의 정렬을 극대화하는 것입니다.

3. 임베딩 공간 근사

간단한 구현은 $g_k$를 얻기 위해 모델을 통해 $K$개의 전체 역전파를 계산해야 합니다. 계산 비용을 줄이기 위해 GraP 는 전체 파라미터 공간이 아닌 공유 백본 임베딩 (shared backbone embedding) $z$에 대한 그래디언트를 계산합니다.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
손실들이 동일한 임베딩 위에 있는 별도의 헤드를 통해 적용되므로, 이러한 그래디언트를 계산하는 것은 전체 역전파보다 훨씬 저렴합니다. 최적화 목표는 이 공유 표현 공간에서 정렬을 극대화하는 것으로 바뀝니다.
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. 정규화 전략

본 논문은 가중치 자체 (예: $\sum w_k = 1$) 나 그 $L_2$ 노름을 정규화하는 것을 반대합니다. 이러한 접근법은 복합 그래디언트의 최적 방향을 왜곡할 수 있기 때문입니다. 대신 GraP 는 복합 그래디언트 자체의 $L_2$-노름 ($\|\sum w_k g_k\| = 1$) 에 제약을 부과합니다. 이는 최적화가 업데이트 벡터의 크기보다는 방향에 초점을 맞추도록 하여, 지역적 목표 하에서 최적 정렬 방향을 회복시킵니다.

5. 알고리즘 효율성

결과적으로 도출된 알고리즘은 반복당 백본을 통해 단순한 순방향 전달 (forward pass) 하나와 역방향 전달 (backward pass) 하나만 요구합니다. 헤드를 위해 계산된 그래디언트는 손실 가중치 업데이트와 백본 파라미터 업데이트 모두에 재사용됩니다. 이는 하이퍼파라미터 조정의 오버헤드를 단일 학습 실행 대비 약 30% 수준으로 줄여주며, 베이지안 최적화의 수배에 달하는 비용과 비교됩니다.

주요 기여

1. 문제 정의: 저자들은 하위 성능을 명시적으로 목표로 하는 2 단계 최적화 문제로 사전 학습 손실 가중치 선택 문제를 공식화했습니다.
2. 정규화 통찰: 그들은 정규화 전략을 분석하고, 가중치 자체를 제약하는 것보다 업데이트 벡터 (복합 그래디언트) 의 노름을 제약하는 것이 최적 정렬 방향을 보존하여 더 우수함을 입증했습니다.
3. 효율적인 알고리즘: 저자들은 하이퍼파라미터 수에 대한 순방향 - 역방향 단계 수와 관련하여 $O(1)$ 복잡도를 가진 GraP 알고리즘을 제안했습니다. 이는 공유 표현 공간에서의 그래디언트 정렬을 활용하여 여러 번의 역전파를 피합니다.
4. 실증적 효율성: 이 방법은 베이지안 탐색과 같은 표준 하이퍼파라미터 최적화와 비교하여 동등하거나 더 나은 성능을 달성하면서도 거의 50 배 더 효율적이며, 전통적인 방법이 계산적으로 불가능한 환경에서도 조정을 가능하게 합니다.
5. 중복성 식별: 이 접근법은 중복된 손실 항을 식별하고 가중치를 낮추는 능력 (예: 비전 작업에서 특정 대비 손실에 거의 0 에 가까운 가중치 부여) 을 보여주며, 이는 자동 손실 선택 메커니즘으로 효과적으로 작용합니다.

실험 결과

이 방법은 이벤트 시퀀스 모델링과 자기 지도 컴퓨터 비전 두 도메인에서 평가되었습니다.

이벤트 시퀀스 모델링

• 데이터셋: Churn, AgePred, AlfaBattle, MIMIC-III, Taobao.
• 설정: 사전 학습은 여러 필드에 걸친 다음 토큰 예측 (NTP) 과 대비 손실 (contrastive loss) 의 복합 목표에 포함되었습니다.
• 성능: GraP 는 세심하게 조정된 베이스라인 (GradNorm, DWA, MGDA, PCGrad 포함) 과 Optuna 로 조정된 가중치와 일치하거나 더 나은 성능을 보였습니다. RNN, Transformer, HT-Transformer 백본 전반에 걸쳐 Churn 및 MIMIC-III 데이터셋에서 일관되게 최상의 결과를 달성했습니다.
• 효율성: 수십 개의 독립 실행이 필요한 베이지안 최적화 (Optuna) 와 비교하여 GraP 는 단일 학습 실행과 적절한 오버헤드 (평균 처리량 44% 감소) 로 유사한 성능을 달성했습니다.

자기 지도 컴퓨터 비전

• 설정: All4One 프레임워크 (Barlow Twins, BYOL, NNCLR 목표 결합) 를 사용하여 CIFAR-10, CIFAR-100, ImageNet-100 에서 평가되었습니다.
• 성능: 목표가 종종 잘 균형을 이루는 이 설정에서 GraP 는 기본 All4One 파라미터 및 기타 다중 작업 학습 방법과 비교 가능한 성능을 달성했습니다.
• 손실 선택: 주목할 만한 점은 GraP 가 Attn-NNCLR 손실에 거의 0 에 가까운 가중치를 부여했다는 것입니다. 이 손실을 수동으로 제거해도 유사한 성능이 나타나, 이 방법이 저영향 또는 중복 목표를 식별할 수 있음을 검증했습니다.

제거 실험 (Ablation Studies)

• 레이블 누출: 가중치 조정에 사용된 레이블 데이터의 비율을 변화시킨 실험은 견고성을 보여주었으며, 레이블은 백본에 대한 직접적인 감독이 아닌 약한 안내 신호를 제공함을 시사합니다.
• 정규화: 제안된 복합 그래디언트 정규화는 표준 가중치 합 정규화보다 우수한 성능을 보였습니다.
• 보조 학습: GraP 로 사전 학습한 후 지도 미세 조정을 수행한 것은 제한된 레이블 데이터에서의 과적합 감소로 인해 결합 보조 학습보다 우수한 성능을 보였습니다.

중요성과 주장

본 논문은 그래디언트 기반 손실 가중치 최적화가 고전적인 다중 작업 학습 및 검색 기반 하이퍼파라미터 최적화에 대한 실용적이고 효율적인 대안이라고 주장합니다. 그 주요 중요성은 전통적인 방법이 계산적으로 불가능한 여러 손실 항을 가진 대규모 사전 학습 시나리오에서 하이퍼파라미터 조정을 실현 가능하게 만든다는 점에 있습니다.

저자들은 GraP 가 여러 번의 학습 실행을 필요로 하지 않으므로 계산 예산이 제한된 환경에 적합하다고 강조합니다. 이 방법은 각 손실에 대해 별도의 헤드를 가정하지만, 사전 학습 업데이트 방향을 하위 목표와 정렬시킴으로써 상충되는 목표 간의 트레이드오프를 효과적으로 처리합니다. 이 연구는 복잡한 사전 학습 파이프라인에서 손실 가중치의 자동적이고 온라인적인 정렬이 고비용의 수동 조정이나 블랙박스 검색을 통해 얻은 것만큼이나, 혹은 그 이상으로 전이되는 표현을 제공할 수 있음을 시사합니다.