OLoRA 리뷰

PEFT 라이브러리를 보다가 LoRA의 학습을 발전시킨 OLoRA를 발견하여 리뷰해본다.

기본 정보

• Title: OLoRA: Orthonormal Low-Rank Adaptation of LLMs
• 저자: Kerim Büyükakyüz(Sicim AI)
• 링크: https://arxiv.org/abs/2406.01775

내용 정리

요약

LoRA를 구성하는 AB matrix를 pretrained 파라미터인 W matrix를 QR decomposition해서 얻은 QR matrix로 간주하였다. 여기서 Q matrix가 orthonormal 한데 이런 특성 때문에 학습의 안정성을 높이고 모델의 최종 성능을 높일 수 있었다고 한다. 제목의 Orthonormal은 이렇게 해서 얻은 matrix Q가 orthonormal하기 때문에 들어간 모양이다. QLoRA는 있으니 OLoRA로 했나보다.

방법론

1. Orthonormality in Neural Network
   • Neural network 파라미터 행렬이 orthonormal하면 optimization, generalization에 아래와 같은 긍정적인 영향을 준다.
   • Gradient flow 향상:
     • Orthonormal 행렬은 backpropagation 과정 중에 gradient norm을 유지하여 vanishing, exploding gradients 문제를 줄여준다.
   • Optimization landscape 향상:
     • Orthogonal group은 optimization을 향상시키는 geometric 성질을 갖고 있다. 따라서 수렴 속도가 빨라지고 더 넓은 범위의 parameter를 탐색해볼 수 있어 generalization이 잘 된다.
2. OLoRA
   • W를 QR decomposition하여 Q,R matrix를 얻는다.(line 3)
   • 이렇게 얻은 Q에서는 첫 r개의 column을 선택하여 Q_r, R에서는 첫 r개의 row를 선택하여 R_r을 얻는다. 이 r이 기본 LoRA의 r과 같은 역할을 한다(A, B 행렬이 d x r, r x d의 꼴을 갖는 것과 같은 방식).(line 4)
   • 원래 파라미터 W를 W - sQrRr로 초기화한다(s는 scaling coefficient).(line 5)
   • 학습 step이 시작할 때 W를 W + QrRr로 업데이트 한다. 이렇게 하여 매 forward step마다 W의 값이 달라진다. 이렇게 변하는 W는 학습을 위해서만 사용되는거 같다.(어짜피 저장할 떄 LoRA만 저장할 것이다.) (line 12)
   • Qr, Rr을 gradient를 이용해 학습한다.(line 16 ~ 19)
   • 이렇게 W의 left singular vector들로 Qr을 만들어서 학습하면 well-conditioned subspace 상에서 학습이 이루어져 수렴이 빠르고 안정성이 높아진다.
   • (여기서부터는 좀 이해가 잘 안되긴 하는데)또한 원래 weight의 spectral properties를 보존한다. 이를 통해 안정성과 pretrained 모델의 능력을 유지할 수 있다.
   • Orthonormal한 기저들로 이루어진 low-rank subspace로 학습을 제한하면 새로운 structural inductive bias를 도입할 수 있다. 이 bias는 모델이 fine tuning 과정에서 가장 중요한 방향으로 변화할 수 있게 해준다. Low-rank로 학습하는 것은 또한 training data에 overfitting되는 걸 막아준다(이건 모든 LoRA에 적용되는 거 아닌가)
     (Inductive bias for generalization에 적힌 설명들이 학습 과정 중 생긴 W의 변화와 OLoRA만의 학습 방법에 대한 특별한 내용을 설명하는 것인지 잘 이해가 가지 않는다...)
    

 

후기

단순히 QR matrix 꼴로 LoRA 파라미터를 초기화하고 학습하는 걸 넘어서 조금 재밌는 시도들을 했는데, 그런 아이디어들은 어디에서 근거를 찾아서 시도했는지 궁금하다. 그리고 이 학습 방법에 대해 의미를 부여하기 위해서는 W + QrRr 에 대한 설명, 근거가 필요할 거 같은데 왜 그렇게 했는지, 왜 좋은 방법인지는 잘 이해가 가지 않는다.