LoRA 원리와 LoRALinear 직접 구현: 큰 가중치는 얼리고 작은 두 행렬만 학습하는 이유
LoRA는 거대한 가중치 행렬 W는 그대로 두고, 작은 행렬 두 개(B, A)의 곱 BA만 학습해서 파라미터 수를 크게 줄이는 방법이다. 오늘은 왜 "low-rank"로 충분한지 직관을 잡고, nn.Linear를 감싸는 LoRALinear를 직접 구현한다. rank를 4·8·16·32로 바꿔가며 학습 파라미터 수와 성능이 어떻게 달라지는지 확인한다.
By 송민성
오늘의 목표
LoRA가 왜 "low-rank(저차원)" 구조인지 이해한다. 큰 가중치 W를 고정(frozen)한 채 변화량 ΔW만 BA라는 작은 두 행렬의 곱으로 표현해 학습하는 원리를 설명할 수 있고, 이를 nn.Linear 위에 직접 구현할 수 있게 된다.
개념과 직관
왜 필요한가. 모델 하나를 통째로 파인튜닝(fine-tuning, 이미 학습된 모델을 내 데이터로 추가 학습하는 것)하면 모든 가중치를 갱신해야 한다. 파라미터가 수십억 개면 가중치 외에 학습에 필요한 부가 정보까지 메모리에 올라가서 GPU가 버티기 어렵다. 그래서 "원래 가중치는 건드리지 말고, 바뀌어야 할 부분만 아주 작게 따로 배우자"는 발상이 나왔다.
핵심 아이디어. 선형층(linear layer)은 y = W x 를 계산한다. 파인튜닝 후의 가중치를 W' = W + ΔW 라고 하자. LoRA(Low-Rank Adaptation, 저랭크 적응)는 ΔW를 통째로 학습하지 않고 두 작은 행렬의 곱으로 놓는다.
ΔW = B · A
W : (d_out × d_in) ← 고정(frozen)
B : (d_out × r) ← 학습
A : (r × d_in) ← 학습
r : rank, d_in·d_out 보다 훨씬 작음 (예: 4~32)랭크(rank, 행렬이 담고 있는 독립적인 방향의 개수)란? 비유하자면, 색을 만드는 데 물감이 수백 종 필요한 것이 아니라 빨강·파랑·노랑 몇 가지 기본색을 섞으면 대부분 표현되는 것과 같다. ΔW가 d_out × d_in 크기의 큰 표지만, "모델을 내 과제에 맞게 살짝 바꾸는 방향"은 몇 개의 기본 방향으로 충분히 표현된다는 가정이 LoRA의 바탕이다. 이것은 원 논문이 제시한 가설이고, 실제로 많은 과제에서 잘 통했지만 모든 경우에 보장되는 것은 아니다.
파라미터 수 계산. d_in = d_out = 4096, r = 8이라면:
- 전체 ΔW: 4096 × 4096 = 16,777,216개
- LoRA: B(4096×8) + A(8×4096) = 65,536개
약 1/256 수준이다. 식으로 쓰면 LoRA 파라미터 = r × (d_in + d_out).
scaling과 alpha. 실제 forward는 이렇게 계산한다.
y = W x + (alpha / r) · B (A x)alpha는 우리가 정하는 상수이고, alpha / r이 scaling(저차원 항에 곱하는 배율)이다. r을 바꿔도 출력 크기가 크게 흔들리지 않도록 맞춰주는 역할이라고 이해하면 된다. 흔한 설정은 alpha = r 또는 alpha = 2r 정도이지만, 정답이 있는 값은 아니어서 실험으로 확인한다.
초기화. A는 작은 랜덤 값, B는 0으로 시작한다. 그러면 처음에는 BA = 0이라 모델 출력이 원래 모델과 정확히 같다. "기존 능력에서 출발해 조금씩 벗어난다"는 안전한 시작점이다.
용어 정리.
- 어댑터(adapter): 원 모델 위에 덧붙이는 작은 학습 모듈. 여기서는 A, B 쌍.
- 타깃 모듈(target module): LoRA를 붙일 층. 보통 Transformer 안의 선형층(예: 쿼리·밸류 투영 등)을 고른다. 아직 Transformer 내부 구조를 자세히 배우지 않았다면, 오늘은 "모델 안의 nn.Linear 중 일부를 골라 바꿔 끼운다" 정도로만 이해하면 충분하다.
- 병합(merge): 학습이 끝난 뒤 W ← W + (alpha/r)·BA 로 더해 하나의 행렬로 합치는 것. 합치면 추론 시 추가 연산이 없다.
코드로 직접 해보기
1단계: LoRALinear 구현.
import math
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r: int = 8, alpha: int = 16):
super().__init__()
self.base = base
for p in self.base.parameters():
p.requires_grad = False # 원래 가중치 고정
self.r = r
self.scaling = alpha / r
self.A = nn.Parameter(torch.empty(r, base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, r))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5)) # A: 랜덤, B: 0
def forward(self, x):
low_rank = (x @ self.A.T) @ self.B.T # x → r차원 → d_out
return self.base(x) + self.scaling * low_rank
@torch.no_grad()
def merge(self):
self.base.weight += self.scaling * (self.B @ self.A)2단계: 처음엔 원본과 같은지 확인.
torch.manual_seed(0)
base = nn.Linear(16, 32)
lora = LoRALinear(base, r=4, alpha=8)
x = torch.randn(5, 16)
print(torch.allclose(base(x), lora(x))) # True 여야 함 (B=0)3단계: 학습 가능한 파라미터 수 세기.
def count_trainable(m):
return sum(p.numel() for p in m.parameters() if p.requires_grad)
print(count_trainable(lora)) # 4*(16+32) = 1924단계: 학습 후 병합이 같은 결과를 내는지 확인. B를 임의 값으로 바꾼 뒤 비교해본다.
with torch.no_grad():
lora.B.normal_()
before = lora(x)
lora.merge()
lora.B.zero_() # 병합 후 어댑터 항을 끄고 base만 사용
after = lora.base(x)
print(torch.allclose(before, after, atol=1e-5)) # True여기서 x @ self.A.T 를 먼저 계산하는 점이 중요하다. BA를 먼저 곱해 큰 행렬을 만들면 메모리 이점이 사라진다. 항상 작은 차원으로 먼저 줄인 뒤 다시 키운다.
오늘의 실험
작은 장난감 문제로 rank 4·8·16·32를 비교한다. 예를 들어 간단한 MLP(층을 쌓은 작은 신경망)를 먼저 일부 데이터로 학습한 뒤, 약간 다른 분포의 데이터에 LoRA로 적응시킨다. 모델 규모는 CPU에서 몇 분 안에 끝나는 정도로 잡는다.
바꾸는 것: r ∈ {4, 8, 16, 32}. 나머지(alpha/r 규칙, 학습률, 스텝 수, 시드)는 고정한다. alpha는 alpha = 2r처럼 규칙으로 정해 모든 r에서 같은 규칙을 쓴다.
관찰하는 것:
- 학습 가능한 파라미터 수 (위의 count_trainable)
- 검증 손실(validation loss) 또는 정확도
- 전체 파인튜닝(모든 가중치 학습)과의 차이
다음과 같은 표로 기록한다.
| r | 학습 파라미터 수 | 전체 대비 비율 | 검증 손실 | 비고 |
|---|---|---|---|---|
| 4 | ||||
| 8 | ||||
| 16 | ||||
| 32 |
파라미터 수는 r에 정확히 비례해서 늘어나는지 확인하고, 성능은 r을 키울수록 계속 좋아지는지 아니면 어느 지점부터 거의 안 변하는지 본다. 시드를 2~3개 바꿔 돌려서 차이가 시드 흔들림보다 큰지도 확인한다. 작은 실험에서는 결과가 노이즈에 묻힐 수 있으니, 결론은 "이 설정에서는 이랬다"까지만 쓴다.
왜 이걸 배우나 (LLM / Post-training 연결)
LLM 전체를 파인튜닝하기엔 메모리와 비용이 크기 때문에, SFT(Supervised Fine-Tuning, 정답 예시로 지도 학습하는 후속 학습)에서 LoRA는 가장 널리 쓰이는 선택지다. 이후 QLoRA(가중치를 낮은 정밀도로 압축해 올린 뒤 LoRA를 붙이는 방식)와 Project 2에서 직접 쓰게 되며, Post-training 현업에서는 rank·target module 선택과 어댑터 병합·배포를 판단하는 일이 실제로 따라온다.
초보자가 막히는 지점
- base 가중치가 같이 학습됨.
requires_grad = False를 빼먹으면 LoRA가 아니라 전체 파인튜닝이 된다. count_trainable로 항상 확인한다. - optimizer에 frozen 파라미터를 넘김. 보통
filter(lambda p: p.requires_grad, model.parameters())로 학습 대상만 넘긴다. - A와 B를 둘 다 0으로 초기화. 둘 다 0이면 기울기가 모두 0이라 학습이 시작되지 않는다. 하나(B)만 0이고 다른 하나(A)는 랜덤이어야 한다.
- 행렬 모양(shape) 혼동. PyTorch의 nn.Linear는 weight가 (out, in) 모양이고 y = x Wᵀ를 계산한다. 그래서 코드에서
.T가 붙는다. 모양이 안 맞으면 각 단계에서print(t.shape)를 찍어본다. - alpha와 학습률을 함께 바꿈. 둘 다 업데이트 크기에 영향을 주므로, 한 번에 하나만 바꿔야 원인을 알 수 있다.
- 병합을 두 번 호출. merge()를 두 번 하면 변화량이 두 번 더해진다. 병합 전후 출력 비교로 검증하고, 병합은 한 번만 한다.
- "r이 클수록 무조건 좋다"는 오해. 파라미터는 늘지만 성능 향상은 어느 순간 멈출 수 있다. 실험 표로 직접 확인한다.
알고리즘 (병행 1시간)
LeetCode Medium 2문제를 DP(동적 계획법, 큰 문제를 작은 부분 문제로 나눠 결과를 저장해 재사용하는 방식)로 푼다. 문제는 본인이 고르되 1차원 DP(예: 계단 오르기 계열, 집 털기 계열)나 2차원 DP(예: 격자 경로 계열) 중에서 선택하면 부담이 적다.
풀이 후 문제마다 다음을 기록한다.
- 접근: dp[i]가 무엇을 뜻하는지 한 문장, 점화식, 초기값
- 시간·공간복잡도: 예) O(n) / O(n), 공간을 줄일 수 있었는지
- 틀린 이유: 초기값 실수, 인덱스 범위, 점화식 누락 등 어디서 막혔는지