LoRA에 그냥 가중치 감쇠를 걸면 생기는 미묘한 함정

• 이 글의 핵심은 꽤 날카로움. LoRA 미세조정에서 weight_decay를 그냥 켜면, 우리가 풀 파인튜닝에서 기대하던 가중치 감쇠와 다른 정규화가 걸릴 수 있다는 얘기임

  • 풀 파인튜닝에서는 모델의 모든 가중치 θ를 직접 업데이트함
  • 이때 weight decay는 대략 “각 가중치를 0에 가까운 쪽으로 조금씩 밀자”는 의미가 됨
  • 그래서 특정 가중치 하나가 예측에 지나치게 큰 영향을 주는 걸 막는 정규화 역할을 함

• 문제는 LoRA가 애초에 가중치를 직접 학습하지 않는다는 점임

  • LoRA는 원본 가중치 W_init을 얼려두고, 작은 어댑터 행렬 A와 B를 학습함
  • 실제로 모델에 반영되는 변화량은 ΔW = AB이고, 최종 적응 행렬은 W = W_init + AB가 됨
  • 그러니까 학습 파라미터는 A, B지만, 우리가 정규화하고 싶은 대상은 사실 W_init + AB 전체일 수 있음

• 숫자로 보면 LoRA가 왜 매력적인지도 확실히 보임

  • 예시 행렬 크기가 4096 × 16384라면 풀 파인튜닝에서는 이 행렬 하나만으로도 67,108,864개, 약 6710만 개 파라미터를 학습해야 함
  • LoRA rank r=4를 쓰면 필요한 파라미터는 4096×4 + 4×16384 = 81,920개뿐임
  • 원래 대비 0.1%도 안 되는 파라미터만 학습하는 셈이라, 가중치·그래디언트·옵티마이저 상태를 저장하는 비용이 확 줄어듦

• 풀 파인튜닝은 유연하지만 요즘 모델 크기에서는 운영 비용이 바로 튀어나옴

  • GPT-3 175B 같은 모델은 학습할 숫자가 1750억 개라는 뜻임
  • 모델 자체뿐 아니라 그래디언트와 옵티마이저 상태까지 들고 있어야 해서 일반적으로 모델 메모리의 최소 3배가 필요함
  • 앱 안에 태스크가 10개 있으면 풀 파인튜닝 방식은 모델도 10벌 호스팅해야 하는 느낌이 됨. 비용이 바로 현실 문제가 됨

• LoRA는 이 지점을 꽤 깔끔하게 우회함

  • 원본 모델은 한 벌만 로드하고, 태스크별로 작은 어댑터만 갈아끼울 수 있음
  • 그래서 “태스크별 튜닝 모델”을 운영하는 게 훨씬 현실적이 됨
  • 다만 바로 이 구조 때문에 optimizer의 정규화 항도 LoRA 구조를 이해해야 함

• 글에서 제안하는 수정은 정규화 대상을 A나 B가 아니라 전체 적응 행렬 W_init + AB로 잡는 것임

  • 풀 파인튜닝의 weight decay가 전체 가중치 행렬을 0 쪽으로 미는 거라면, LoRA에서도 같은 의미를 유지하려면 W_init + AB 기준으로 regularization term을 둬야 함
  • 그러면 A와 B에 대한 decay term은 각각 (W_init + AB)B^T, A^T(W_init + AB) 형태로 나옴
  • 결론적으로 LoRA용 weight decay는 “파라미터에 weight_decay를 곱해서 더하기”보다 한 단계 더 구조적인 계산이 필요함

sequenceDiagram
    participant 학습데이터 as 학습 데이터
    participant 원본가중치 as 원본 가중치 W_init
    participant 어댑터 as LoRA 어댑터 A,B
    participant 옵티마이저 as 옵티마이저
    participant 적응행렬 as 적응 행렬 W_init + AB
    학습데이터->>어댑터: 손실 그래디언트 계산
    원본가중치->>적응행렬: 얼린 가중치 제공
    어댑터->>적응행렬: AB 변화량 제공
    적응행렬->>옵티마이저: 수정된 decay term 계산
    옵티마이저->>어댑터: A,B 업데이트

• Optax의 일반 weight decay 구현은 아주 단순함

  • 각 업데이트 g에 weight_decay * p를 더하는 방식임
  • 풀 파인튜닝에서는 이게 자연스럽지만, LoRA에서는 p가 전체 행렬이 아니라 kernelA, kernelB 같은 어댑터 파라미터가 됨
  • 글의 구현은 파라미터 경로를 따라가서 같은 레이어의 kernel, kernelA, kernelB를 꺼낸 뒤 LoRA용 decay term을 따로 계산함

• 실무적으로는 “LoRA 쓰면 메모리 싸다”에서 끝낼 게 아니라, optimizer 설정까지 같이 봐야 한다는 메시지임

  • 특히 기존 학습 코드에서 weight decay를 공통 옵션으로 켜둔 팀이라면, LoRA 어댑터에도 같은 의미로 적용되고 있는지 확인할 필요가 있음
  • 논문이나 튜토리얼에서 rank, learning rate, target module만 챙기고 정규화 의미를 놓치면 성능 차이가 애매하게 숨어들 수 있음

기술 맥락

• 여기서 선택의 핵심은 “무엇을 정규화할 것인가”예요. 풀 파인튜닝에서는 학습 파라미터와 실제 모델 가중치가 거의 같은 대상이라 weight decay의 의미가 단순한데, LoRA에서는 학습 파라미터가 A와 B이고 실제로 쓰이는 행렬은 W_init + AB거든요.

• 왜 이게 중요하냐면, A와 B를 작게 만드는 것과 W_init + AB를 0에 가깝게 만드는 건 같은 목표가 아니기 때문이에요. LoRA는 저랭크 분해 구조라서 여러 A, B 조합이 비슷한 AB를 만들 수 있고, 그래서 파라미터 자체에 거는 패널티가 최종 행렬에 대한 패널티와 어긋날 수 있어요.

• 구현 관점에서는 일반 옵티마이저의 g + weight_decay * p 패턴을 그대로 못 쓰는 게 포인트예요. 레이어 안에서 얼려둔 W_init과 어댑터 A, B를 함께 꺼내고, 현재 파라미터가 A인지 B인지에 따라 다른 decay term을 넣어야 해요.

• 이 글이 실무적으로 좋은 이유는 수식에서 코드까지 이어진다는 점이에요. JAX와 Optax처럼 파라미터 트리를 타고 들어가는 구조에서는 tree_map_with_path로 현재 파라미터 이름을 확인한 뒤 LoRA 어댑터에만 별도 로직을 적용할 수 있거든요.