태스크 산술에서의 가중치 얽힘 해제 이해 및 강제화

태스크 산술이 왜 성공하는지에 대한 근본적인 이론적 설명을 제공하고, 이를 바탕으로 모델 병합 시 발생하는 태스크 간 간섭 문제를 해결하는 새로운 정규화 기법을 제시한다. OrthoReg를 통해 추가 학습 없이도 여러 모델의 능력을 더 정교하게 결합하거나 특정 능력을 깨끗하게 제거할 수 있다.

태스크 산술의 핵심은 여러 모델의 지식을 단순히 더하는 것인데, 기존에는 왜 이 덧셈이 서로의 지식을 파괴하지 않는지 명확하지 않았다. 연구진은 이를 딥러닝의 기초인 가중치 행렬 구조에서 찾았다. 선형 레이어의 가중치 행렬 W가 입력 데이터의 특정 특징에만 반응하도록 분화되어 있다면, 서로 다른 태스크는 W의 서로 다른 열(column)들을 사용하게 된다. 이것이 바로 Task-Feature Specialization(TFS)이다.

만약 두 태스크가 사용하는 가중치 열들이 겹치지 않는다면(직교한다면), 한 태스크를 위해 가중치를 수정해도 다른 태스크의 출력에는 영향을 주지 않는다. 즉, Gradient Descent 과정에서 발생하는 가중치 업데이트가 서로 독립적인 방향으로 이루어지는 것이다. 하지만 실제 모델에서는 이러한 특징 분리가 완벽하지 않아 태스크 간 간섭이 발생한다.

논문은 이 문제를 해결하기 위해 가중치 업데이트 시 행렬의 각 열들이 서로 수직이 되도록 강제하는 OrthoReg를 제안한다. 이는 마치 여러 사람이 하나의 도화지에 그림을 그릴 때, 각자에게 할당된 구역이 겹치지 않도록 선을 긋는 것과 같다. 결과적으로 ViT-L-14 모델에서 기존 방식 대비 성능을 최대 4.16%p 향상시키는 등 태스크 병합의 효율성을 극대화했다.

전체적으로 Neural Tangent Kernel(NTK) 선형화 가설을 기반으로 태스크 산술의 메커니즘을 분석한다. 핵심은 가중치 업데이트 행렬 delta W의 구조를 제어하는 것이다.

OrthoReg 정규화 항은 각 레이어 l의 업데이트 행렬 delta W(l)에 대해 계산된다. [delta W(l)의 전치 행렬과 delta W(l)을 곱하여 Gram 행렬을 구하고] -> [이 결과에서 단위 행렬 I를 뺀 값의 Frobenius norm의 제곱을 계산하여] -> [최종 손실 함수에 더해준다]. 이 연산의 의미는 delta W의 각 열 벡터들이 서로 직교하면서 동시에 단위 길이를 갖도록 유도하는 것이다.

학습 시 전체 손실 함수는 '기존 태스크 손실 + lambda * OrthoReg 손실'로 구성된다. [lambda라는 하이퍼파라미터를 입력으로 받아] -> [정규화의 강도를 조절하며] -> [가중치 업데이트 방향이 기존 특징들과 겹치지 않게 제한한다]. 이를 통해 모델은 새로운 태스크를 학습하면서도 기존 지식의 영역을 침범하지 않는 최적의 업데이트 경로를 찾게 된다.

CLIP 기반의 ViT 모델들(B-32, B-16, L-14)을 대상으로 8개의 이미지 분류 태스크를 병합하는 실험을 수행했다. ViT-L-14 모델에서 단순 파인튜닝 대비 정확도가 84.07%에서 88.23%로 크게 향상되었으며, 이는 개별 태스크 전용 모델들의 평균 성능(100.08% Normalized Accuracy)과 대등한 수준이다.

태스크 제거(Task Negation) 실험에서도 탁월한 효과를 보였다. 특정 태스크의 벡터를 뺐을 때, OrthoReg를 적용한 경우 타겟 태스크의 정확도는 더 급격히 떨어지면서도(제거 성능 향상), ImageNet과 같은 대조군 태스크의 성능은 더 잘 보존되는 것으로 나타났다.

Ablation Study를 통해 정규화 강도 lambda에 대한 민감도를 분석한 결과, lambda 값이 커질수록 가중치 간 직교성이 뚜렷해지며 모델 병합 성능이 안정적으로 우상향하는 경향을 확인했다. 또한 LoRA와 같은 매개변수 효율적 파인튜닝(PEFT) 기법과 결합했을 때도 일관된 성능 향상을 보였다.

본 연구는 태스크 산술의 성공 요인을 Task-Feature Specialization(TFS)으로 규정하고, 이를 수학적으로 공식화했다. 정의에 따르면, 특정 태스크 t에 대해 모델 출력이 가중치 행렬 W의 특정 열 k의 활성화에 민감하지 않다면 해당 인덱스는 specialized feature set에 포함되지 않는다. 두 태스크의 feature set이 공집합일 때 완벽한 가중치 얽힘 해제가 가능함을 증명했다.

이론적 분석을 통해 TFS가 가중치 벡터 직교성(WVO)의 충분조건임을 보였다. 즉, W의 Gram 행렬이 블록 대각(block-diagonal) 구조를 가질 때 태스크 간 간섭이 최소화된다. 실제 CLIP ViT-B/16 모델 분석 결과, 대부분의 레이어에서 가중치 벡터 간 각도가 90도에 매우 가깝게 집중되어 있음을 확인하여 이론의 타당성을 뒷받침했다.

제안된 OrthoReg는 TTA(Tangent Task Arithmetic)와 달리 자코비안(Jacobian) 계산이 필요 없어 메모리 효율성이 높다. TTA는 훈련 시간과 메모리를 2~3배 더 소모하는 반면, OrthoReg는 표준 파인튜닝 대비 미미한 오버헤드만으로도 더 높은 성능을 달성한다. 이는 가중치 업데이트의 내부 구조를 직접 제어하는 방식이 모델의 출력 공간을 제어하는 방식보다 더 효율적임을 시사한다.

본 논문은 주로 선형 레이어(Linear Layers)의 가중치에 집중하여 분석을 진행했으며, 바이어스(bias)나 레이어 정규화(LayerNorm) 파라미터의 영향은 분석 범위에서 제외했다. 또한 TFS가 실제 사전 학습 모델에서 완벽하게 성립하기 어려운 이상적인 조건이라는 점을 인정하며, 이를 보완하기 위해 정규화 기법을 제안했다.