기술 요약: BaLoRA (베이지안 저랭크 적응)

1. 문제 제기

저랭크 적응 (LoRA) 은 대규모 사전 훈련 모델의 매개변수 효율적 파인튜닝 (PEFT) 표준으로 자리 잡았으며, 업데이트를 저랭크 구조로 제한함으로써 계산 비용을 줄여줍니다. 그러나 해당 논문은 표준 LoRA 및 관련 결정론적 PEFT 방법에서 세 가지 중요한 한계를 지적합니다:

1. 제한된 표현력: 저랭크 점 추정 (point-estimate) 업데이트는 종종 전체 파인튜닝 (FT) 에 비해 지속적인 정확도 격차를 남깁니다.
2. 불확실성 정량화 부재: 표준 LoRA 는 예측 불확실성을 추정할 수 있는 내장 메커니즘을 제공하지 않으며, 이는 재료 과학, 기상 예보와 같은 안전이 중요하거나 데이터가 부족한 분야의 적용을 방해합니다.
3. 부족한 유연성: 기존 방법들은 훈련 과정에서 불확실성을 자연스럽게 인코딩하지 못하며, 종종 불확실성을 근사하기 위해 모델 앙상블 (메모리 및 계산량 증가) 을 필요로 합니다.

2. 방법론: BaLoRA

BaLoRA 는 고정된 매개변수가 아닌 확률 변수로 축소 행렬을 취급하는 LoRA 의 베이지안 확장을 도입합니다. 핵심 혁신은 LoRA 축소 행렬 ($\theta_A$) 에 대한 새로운 입력 적응형 베이지안 재파라미터화입니다.

2.1 입력 적응형 베이지안 재파라미터화

결정론적 업데이트 $y = \theta_0 x + \theta_B \theta_A x$ 대신, BaLoRA 는 입력 의존적 가우시안 노이즈로 축소 행렬 $\theta_A$의 항목을 교란시킵니다:
$$\omega_{A;ij} = \theta_{A;ij} + \sqrt{\alpha(x)\theta_{A;ij}^2} \epsilon_{ij}, \quad \epsilon_{ij} \sim \mathcal{N}(0, 1)$$
여기서 $\alpha(x)$는 입력 $x$에 기반하여 분산 스케일을 조절하는 경량 추론 네트워크입니다. 이 공식은 크기가 큰 가중치가 비례적으로 더 많은 불확실성을 가지며, 불확실성이 입력 활성화와 함께 스케일링됨을 의미합니다.

2.2 훈련 목적 함수

모델은 데이터 적합도와 모델 복잡성 (사후분포와 약하게 정보화된 사전분포 간의 KL 발산) 을 균형 있게 맞추는 증거 하한 (ELBO) 을 최대화하도록 훈련됩니다. KL 발산은 스칼라 $\alpha(x)$와 드롭아웃 하이퍼파라미터 $p$에만 의존하는 폐쇄형 해를 허용하므로, 계산적으로 평가 비용이 저렴합니다.

2.3 저랭크 로컬 재파라미터화 트릭

베이지안 샘플링의 단순한 적용은 $O(k^2)$ 크기의 전체 출력 공분산 행렬을 물리적으로 생성해야 하므로 계산적으로 불가능합니다. BaLoRA 는 LoRA 업데이트의 저랭크 구조를 활용하여 효율적인 샘플링 절차를 유도합니다.

• 노이즈는 고차원 출력 공간이 아닌 저차원 잠재 공간 ($\mathbb{R}^r$) 에서 샘플링됩니다.
• 정확한 사후 예측 샘플은 다음과 같이 계산됩니다:
  $$y = \theta_0 x + \theta_B \theta_A x + \theta_B (\sqrt{d(x)} \odot \epsilon_d)$$
  여기서 $d(x) = \alpha(x) \theta_A^2 x^2$이고 $\epsilon_d \sim \mathcal{N}(0, I_r)$입니다.
• 이 접근법은 표준 LoRA 의 계산 스케일링 ($O(Bkr)$) 을 유지하면서 전체 공분산 행렬의 물리적 생성을 피합니다.

2.4 추론 모드

BaLoRA 는 두 가지 배포 모드를 지원합니다:

1. 결정론적 모드: 어댑터 가중치가 사전 훈련 가중치와 병합됩니다 ($\theta = \theta_0 + \theta_B \theta_A$). 이는 표준 LoRA 에 비해 추가 지연 시간 없이 수행됩니다.
2. 확률적 (베이지안) 모드: 예측 불확실성 (평균 및 분산) 을 몬테카를로 샘플링을 통해 추정하기 위해 샘플링된 가중치를 사용하여 여러 순전파를 수행합니다.

3. 주요 기여

1. 암시적 정규화를 통한 확률적 PEFT: BaLoRA 는 입력 의존적 노이즈 주입이 암시적 정규화제로 작용하는 방법을 제안합니다. 높은 활성화 방향에 교란을 집중시킴으로써 과도한 확신을 가진 업데이트를 억제하여 최적화 및 작업 성능을 향상시킵니다.
2. 효율적인 저랭크 샘플링: 저차원 잠재 공간에서 사후 예측 분포의 정확한 샘플링을 가능하게 하는 로컬 재파라미터화 트릭을 유도하여, 표준 LoRA 의 계산 효율성과 일치합니다.
3. 통합된 정확도 및 불확실성: 이 방법은 단일 파인튜닝 실행 내에서 (앙상블 방법과 달리) 상식 추론 및 이미지 분류에서 최첨단 (SOTA) 성능을 달성하면서 동시에 보정된 불확실성 추정을 제공합니다.

4. 실험 결과

논문은 BaLoRA 를 세 가지 도메인에서 평가합니다:

4.1 상식 추론 (NLP)

• 모델: Llama-2-7B 및 Llama-3-8B.
• 벤치마크: BoolQ, PIQA, ARC-c, WinoG 등 8 가지 작업.
• 결과: BaLoRA 는 Llama-2 의 경우 평균 약 2.9%, Llama-3 의 경우 약 5% 더 높은 성능으로 표준 LoRA 를 능가합니다. Llama-3-8B 의 경우 8 개 벤치마크 중 6 개에서, Llama-2-7B 의 경우 5 개에서 SOTA 정확도를 달성하여 DoRA, MoRA, HiRA 와 같은 결정론적 경쟁자들을 능가합니다.
• 강건성: 성능은 다양한 사전 확률 하이퍼파라미터 ($p$) 범위에서 안정적으로 유지됩니다.

4.2 시각 지각 (컴퓨터 비전)

• 모델: ImageNet-21K 에서 사전 훈련된 ViT-L/16.
• 벤치마크: CIFAR-10, CIFAR-100, Oxford-IIIT Pets, Oxford Flowers-102.
• 결과: BaLoRA 는 PEFT 방법들 사이에서 일관되게 1 위를 차지합니다. 특히 Oxford Flowers-102 에서 전체 파인튜닝 정확도 (99.61%) 와 일치하며 CIFAR-100 에서는 전체 파인튜닝을 능가합니다.
• 보정: BaLoRA 는 4 개 데이터셋 중 2 개에서 가장 낮은 기대 보정 오차 (ECE) 를 달성하고 나머지는 2 위를 기록하며, 사후 보정 없이도 LoRA 및 DoRA 를 능가합니다.

4.3 재료 물성 예측 (과학적 응용)

• 작업: QMOF 데이터셋과 MOFTransformer 를 사용한 금속 - 유기 골격체 (MOFs) 의 밴드갭 예측.
• 비교: BaLoRA 대 LoRA 모델 앙상블 (여러 훈련 실행 필요).
• 결과: BaLoRA 는 평균 절대 오차 (MAE) 를 0.266 eV 로 달성하여 전체 파인튜닝 (0.224 eV) 과의 격차를 좁히면서 훈련 가능한 매개변수의 6.75% 만 사용합니다.
• 불확실성: BaLoRA 는 제로샷 불확실성 추정을 생성하며, 이는 단일 파인튜닝 실행만 필요로 함에도 불구하고 훈련된 LoRA 모델 앙상블 (0.309) 보다 모델 오차와 더 강하게 상관관계를 가집니다 (스피어만 상관관계 0.346). 불확실성 추정은 테스트 시 몬테카를로 단계 수에 따라 단조롭게 향상됩니다.

5. 중요성 및 주장

이 논문은 BaLoRA 가 PEFT 에서의 표현력과 불확실성 정량화라는 이중 과제를 해결한다고 주장합니다. 베이지안 귀납적 편향을 저랭크 어댑터의 기하학에 통합함으로써, 이 방법은 순수한 결정론적 접근법보다 상호 보완적이고 가산적인 이점을 제공합니다.

주요 주장은 다음과 같습니다:

• 효율성: BaLoRA 는 표준 LoRA 의 계산 스케일링과 일치하며, 결정론적 모드에서는 추론 지연 시간이 발생하지 않습니다.
• 성능: 베이지안 확장은 모델 용량 증가가 아닌 원칙적인 불확실성 모델링에 힘입어 NLP 및 비전 작업 전반에 걸쳐 전체 파인튜닝과의 격차를 줄이며 예측 정확도를 향상시킵니다.
• 과학적 유용성: 재료 과학 분야에서 BaLoRA 는 안전이 중요한 응용 분야에 필수적인 특징인 오차와 상관관계가 있는 신뢰할 수 있는 불확실성 추정을 제공하며, 계산 오버헤드가 훨씬 낮은 앙상블 기준을 능가합니다.
• 일반화 가능성: 이 접근법은 특정 도메인에 국한되지 않으며, 자연어 추론, 시각 지각, 과학적 회귀 작업에서 효과성을 입증합니다.

저자들은 BaLoRA 가 정확한 예측만큼이나 신뢰할 수 있는 예측이 중요한 환경에서 불확실성 인식 PEFT 의 보다 광범위한 채택을 장려한다고 결론지었습니다.