Relaxed Recursive Transformer

조다영·2026년 4월 24일
post-thumbnail

깃허브에 로그인 했다가 현재 trending이라는 https://github.com/kyegomez/OpenMythos.git 을 보게 되었다.

요즘 많이 언급되는 클로드의 Mythos 모델을 구현한 리포지토리 같다. 학습데이터랑 방법론이랑 뭐랑 뭐랑 다 엄청 자세하게 나와있는데
Recursive Transformer 구조랑 PEFT 방법론 좀 알아가면 좋을 것 같아서 찾아봄.


PEFT(Parameter Efficient Fine Tuning)

  • 핵심 개념: Fine-Tuning에 필요한 자원(시간 및 컴퓨팅 자원)의 양을 줄인다

LoRA(Low-Rank Adaptation)

  • 핵심 개념: PEFT 방법론 중 하나로, 모델 전체를 다시 학습시키는 대신 소규모의 학습 가능한 '저차원 행렬(어댑터)'만 추가해 파라미터 수를 1% 미만으로 줄이면서도 특정 작업이나 화풍에 맞춰 빠르게 파인튜닝하는 효율적인 기법
  • 원문 링크: https://arxiv.org/abs/2106.09685

lora

  • 파인튜닝 전: h(출력값)=x(입력값)W(가중치매트릭스)h(출력값) = x(입력값) * W(가중치 매트릭스)
  • LoRA 적용 후: h=(xW)+(xAB)h = (x * W) + (x * A * B )

만약 x가 14096 차원이라고 하고, W의 차원이 4096 4096 이라고 하면 W는 16,777,216 (16백만) 패러미터를 가지게 된다. 
LoRA에서 사용되는 Weight인 A는, 예를 들어서  차원을 맞춰야 하기 때문에, 일단 입력 차원은 4096 16, B는 164096이라고 하면, 전체 패러미터수는 2164096=131072 으로 16777216 패러미터 대비 0.8%의 패러미터 사이즈가 된다.
즉 0.8%의 패러미터만 학습을 하면되기 때문에 학습 효율이 뛰어나다. 
출처: https://bcho.tistory.com/1452 [조대협의 블로그:티스토리]

라고 한다....아주 효율적인듯함ㅎㅎ


Recursive Transformer

  • 핵심 개념: Transformer의 layer를 반복적으로 재사용(parameter sharing)하여 파라미터 수를 줄인다.
  • looped layer tying: 모델의 파라미터 수를 줄이면서도 깊은 연산 능력을 유지하기 위한 가중치 공유 방식.
  • 원문 링크: https://arxiv.org/abs/2410.20672

젬미니한테 논문 파일 던져주고 모델 구조 좀 요약해달라고 함.


1. 기본 개념: CYCLE 전략

  • 고유한 파라미터를 가진 단일 레이어 블록을 정의한 뒤, 이를 정해진 횟수만큼 반복해서 통과시키는 구조입니다.
  • 이는 레이어가 순차적으로 쌓인 일반적인 트랜스포머와 달리, 출력이 다시 입력 블록으로 들어가는 루프(Loop) 형태를 띱니다.

(1) 수식과 작동 원리

모델의 숨겨진 상태(Hidden State)는 다음과 같은 수식으로 계산됩니다:

htl=f(htl1;Φ((l1)mod L/B)+1)h_{t}^{l}=f(h_{t}^{l-1};\Phi_{((l-1)mod~L/B)+1}^{\prime})

  • htlh_{t}^{l} (Hidden State): 시간 단계 tt에서의 ll번째 레이어 출력값. 가장 처음 값인 ht0h_{t}^{0}는 이전 단어의 임베딩(Embedding) 벡터가 됨.
  • f()f(\cdot) (Layer Function): 트랜스포머의 한 층이 수행하는 연산(멀티 헤드 어텐션 + 피드 포워드 네트워크).
  • Φ\Phi^{\prime} (Shared Parameters): 모델이 실제로 메모리에 들고 있는 유일한 가중치 블록. 여기에는 선형 레이어 가중치와 RMSNorm 파라미터 등이 포함됨
  • LL (Total Layers): 원래 모델이 목표로 하는 전체 깊이(논리적 레이어 수)
  • BB (Number of Looping Blocks): 전체 레이어를 몇 개의 묶음(블록)으로 나눌 것인지를 정하는 수치. 반드시 LL의 약수여야 함.

(2) 인덱스 매핑 로직의 핵심

((l-1) mod L/B) + 1 이 부분이 바로 "어떤 공유 가중치를 가져올 것인가"를 결정하는 '인덱스 저글링'의 핵심입니다.

  • 블록 크기 계산 (L/BL/B): 만약 L=18L=18이고 B=2B=2라면, 모델은 물리적으로 9개의 레이어(18/218/2)만 가집니다.
  • 순환(Modulo) 연산: (l-1) mod 9를 하면, 레이어 번호(ll)가 9를 넘어가는 순간 다시 0부터 시작하게 됩니다.
  • 결과 예시 (L=18,B=2L=18, B=2 일 때):
    • $l=1$ (첫 번째 층): (11)mod9+1=1(1-1) mod 9 + 1 = 1 \rightarrow 공유 블록의 1번 가중치 사용.
    • $l=9$ (아홉 번째 층): (91)mod9+1=9(9-1) mod 9 + 1 = 9 \rightarrow 공유 블록의 9번 가중치 사용.
    • $l=10$ (열 번째 층): (101)mod9+1=1(10-1) mod 9 + 1 = 1 \rightarrow 다시 공유 블록의 1번 가중치를 재사용!

(3) 구체적인 사례 (Gemma 2B 기준)

  • 구조 변화: 18개의 레이어를 가진 Gemma 2B 모델을 루핑 블록 B=2B=2로 설정하면, 실제로는 9개의 레이어 가중치만 저장합니다.
  • 연산 과정: 데이터를 처리할 때 이 9개의 레이어 블록을 두 번 반복해서 통과(Loop twice)시켜 총 18층의 연산을 수행합니다.
  • 공유 범위: 선형 레이어(Linear layers)의 가중치뿐만 아니라 RMSNorm과 같은 정규화 레이어의 파라미터까지 모두 공유(Tie)합니다.

(4) 이 방식의 장점

  • 메모리 효율성: 가중치를 재사용하기 때문에 모델의 메모리 점유율(Footprint)을 크게 낮출 수 있습니다.
  • 추론 최적화와의 결합: 레이어가 반복되는 특성 덕분에, 특정 단계에서 연산을 멈추는 조기 종료(Early-exiting) 메커니즘을 적용하기에 매우 유리한 구조가 됩니다.

2. 릴랙스드 재귀 트랜스포머(Relaxed Recursive Transformer)

재귀 레이어를 사용하게 되면,

  • (1) the set of possible model sizes is limited to scaling the number of layers -> 모델 크기를 조절할 수 있는 옵션이 레이어 수를 늘리거나 줄이는 것밖에 없다
  • (2) each model layer ends up having to serve multiple roles associated with different depths of the model -> 하나의 레이어가 각 반복에서 다른 역할을 수행해야 한다

과 같은 문제가 있다고 한다......

\rightarrow relaxed recursive transformer는 각 반복 회차마다 약간의 변주를 주는 것이 핵심!

(1) 레이어별 LoRA (Multi-LoRA)

relaxed recursive transformer

  • 구조: 한 블록의 공유 레이어(WW')가 있고, 여기에 루프 회차마다 서로 다른 작은 가중치(ΔW\Delta W')가 더해집니다.
  • 수식: h=Wx+ΔWx=Wx+BAxh = W'x + \Delta W'x = W'x + BAx
    • AABB는 랭크(rank) rr을 가진 아주 작은 행렬들
  • 반복별 차별화: 예를 들어 2번 루프를 돈다면, 첫 번째 바퀴에서는 LoRA 모듈 1번을, 두 번째 바퀴에서는 LoRA 모듈 2번을 사용하여 같은 레이어라도 깊이에 따라 미세하게 다른 연산을 수행하게 합니다.

(2) 수식의 변화

기존 수식에 LoRA 파라미터인 ΔΦl\Delta\Phi_{l}^{\prime}이 추가됩니다:

htl=f(htl1;Φ((l1)mod L/B)+1,ΔΦl),l[1,L]h_{t}^{l}=f(h_{t}^{l-1};\Phi_{((l-1)mod~L/B)+1},\Delta\Phi_{l}^{\prime}), \quad l\in[1,L]

(3) SVD 초기화 (Truncated SVD)

릴랙스드 모델의 성능이 좋은 이유는 초기화 방식에 있습니다. 무작위로 시작하는 대신, 원래 거대 모델의 가중치(WW)와 공유 가중치(WW') 사이의 차이(잔차) 를 추출하여 LoRA를 초기화합니다.

  • 원리: Truncated SVD를 통해 원래 모델의 정보를 가장 잘 보존하는 하위 행렬을 찾아냅니다=.
  • 효과: 이 방식을 쓰면 LoRA의 랭크(rr)가 0이면 완전 재귀 모델이 되고, 랭크가 충분히 크면 원래의 거대 모델과 거의 똑같아집니다. 즉, 사용자가 성능과 크기 사이에서 아주 정밀하게 모델을 튜닝할 수 있게 됩니다.

0개의 댓글