기술 요약: PreLort – 랭크 이질성 하에서의 연합 미세 조정을 위한 접두사 중첩형 LoRA (Prefix-Nested LoRA)

1. 문제 정의

연합 학습(Federated Learning, FL)은 원본 데이터를 공유하지 않고 대규모 언어 모델(LLM)의 협력적 학습을 가능하게 하지만, 저계수 적응(Low-Rank Adaptation, LoRA)과 같은 매개변수 효율적 미세 조정(PEFT) 방식을 통해 파운데이션 모델을 적응시킬 때 상당한 어려움에 직면합니다. 주요 장애물은 **랭크 이질성(rank heterogeneity)**입니다. 현실적인 크로스 디바이스 환경에서 클라이언트는 서로 다른 계산 및 메모리 제약을 가지며, 이에 따라 서로 다른 랭크($r$)를 가진 LoRA 어댑터를 초기화하게 됩니다.

기존의 이질적 연합 LoRA 솔루션은 주로 **집계 단계(aggregation phase)**에 집중하며, 서로 다른 랭크의 어댑터를 정렬하기 위해 제로 패딩(zero-padding)을 사용하는 경우가 많습니다. 그러나 저자들은 이러한 접근 방식의 근본적인 한계인 **표현 불일치(representation misalignment)**를 지적합니다. 집계가 수학적으로 정의되어 있더라도, 서로 다른 랭크로 독립적으로 훈련된 모델들을 결합하면 불일치가 발생합니다. 구체적으로, 저계수 클라이언트는 고계수 클라이언트의 업데이트로부터 혜택을 받지 못하는데, 이는 정보가 랭크 차원 전반에 걸쳐 계층적으로 분포되지 않기 때문입니다. 또한 제로 패딩은 고계수 클라이언트의 기여도를 희석시켜, 공유된 저계수 표현의 최적 활용을 방해합니다.

2. 방법론: PreLort

저자들은 PreLort(Prefix-Nested LoRA)를 제안합니다. 이 프레임워크는 **훈련 역학(training dynamics)**과 **집계 규칙(aggregation rules)**을 공동으로 구조화함으로써 랭크 이질성 문제를 해결합니다. 핵심 철학은 어댑터 차원을 **접두사 계층 구조(prefix hierarchy)**로 구성하여, 저계수 차원이 모든 클라이언트 간에 공유되는 작업 관련 정보를 인코딩하는 동시에, 고계수 차원이 점진적인 용량을 포착하도록 보장하는 것입니다.

2.1 중첩된 저계수 구조 (Nested Low-Rank Structure)

PreLort는 어댑터 차원에 접두사 구조화된 구성을 부과합니다. 랭크 $r_k$를 가진 클라이언트의 경우, 처음 $r_k$개의 차원만 활성화됩니다. 결정적으로, 로컬 훈련 중에 어댑터는 단순히 전체 랭크에서만 최적화되는 것이 아니라, 다중 절단 구성(multiple truncated configurations, 접두사 하위 랭크) 하에서도 최적화됩니다.

• 메커니즘: 랭크 $r_k$를 가진 클라이언트는 $r_k$뿐만 아니라, 연합 내에 존재하는 모든 낮은 랭크 수준($r \in \{r \mid r \le r_k\}$)에서도 잘 작동하도록 훈련됩니다.
• 목적: 이는 모델이 작업 관련 신호를 선행(leading)하는 저계수 차원에 집중하도록 강제하여, 해당 차원들이 절단되더라도 유용하고 정보가 풍부한 상태를 유지하도록 합니다.

2.2 로컬 훈련 전략

클라이언트 $k$의 로컬 목적 함수는 모든 유효한 접두사 하위 랭크 $R_k$에 대해 계산된 교차 엔트로피 손실의 평균으로 정의됩니다.
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
실제로 각 훈련 단계는 하나의 풀 랭크 순전파/역전파와 더불어 모든 하위 랭크에 대한 추가적인 패스를 포함합니다. 이 전략은 접두사 세그먼트가 집계 시 사용되는 것과 동일한 입도(granularity)에서 명시적으로 최적화되도록 보장하며, 저계수 차원이 가장 중요한 정보를 포착하는 일관된 계층 구조를 장려합니다.

2.3 세그먼트별 집계 (Segment-wise Aggregation)

제로 패딩으로 인한 신호 희석을 방기하기 위해, PreLort는 그래디언트(업데이트) 공간에서 작동하는 세그먼트별 집계 규칙을 도입합니다.

• 과정: 모든 클라이언트에 대해 모든 차원을 평균 내는 대신, 특정 랭크 세그먼트 $[r_{l-1}, r_l)$에 대해 해당 세그민트에 기여할 수 있는 충분한 랭크 $r_k$를 가진 클라이언트들의 집합에 대해서만 그래디언트를 집계합니다 (즉, $r_k \ge r_l$).
• 공식: $l$번째 세그먼트에 대해 집계된 그래디언트는 다음과 같습니다:
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• 이점: 이는 저계수 클라이언트의 제로 패딩으로부터 발생하는 신호 희석을 방지하며, 고계수 차원이 실제로 이를 학습할 능력을 갖춘 클라이언트에 의해서만 업데이트되도록 보장하는 동시에, 저계수 차원은 전체 인구(population)로부터 혜택을 받도록 합니다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 설명합니다:

1. 표현 불정합의 식별: 저자들은 랭크 차원 간의 구조화된 정렬 부족이 현재의 이질적 연합 LoRA 방법론의 핵심적인 한계이며, 이것이 일관되지 않은 표현을 초래한다는 점을 정확히 짚어냈습니다.
2. 중첩 공식화 및 통합 전략: 공통된 접두사 구조를 부과하는 중첩된 저계수 공식화를 제안합니다. 이는 제로 패딩 없이도 이질적인 랭크 간에 일관된 집계를 가능하게 하는 통합 훈련 및 집계 전략과 결합됩니다.
3. 실증적 검증: 이 방법은 지시어 튜닝(instruction-tuning) 벤치마크에서 기존의 이질적 연합 LoRA 방법들보다 모델 품질과 안정성을 일관되게 향상시킨다는 것을 입증했습니다.

4. 실험 결과

저자들은 두 가지 베이스 모델(TinyLlama-1.1B 및 Qwen2.5-0.5B)을 세 가지 데이터셋(Alpaca, Databricks-dolly-15, 20 Newsgroups)에 대해 평가했습니다.

• 성능 지표: PreLort는 정확도(MMLU, 20 Newsgroups) 및 ROUGE-L 점수에서 베이스라인(ZeroPad, HetLoRA, FLoRA)을 일관되게 능가했습니다.
  • TinyLlama-1.1B에서 PreLort는 MMLU (Dolly)에서 33.86%, Alpaca에서 30.76%의 정확도를 달성하여 차순위 베이스라인들을 앞질렀습니다.
  • Qwen2.5-0.5B에서는 MMLU (Dolly)에서 39.85%의 정확도를 달성했습니다.
• Perplexity: PreLort는 두 베이스 모델 모두에서 Dolly와 NewsGroup에 대해 가장 낮은 퍼플렉시티를 달성하여, 우수한 토큰 레벨 가능도(likelihood)를 보여주었습니다.
• 랭크 강건성: 분석 결과, PreLort는 다양한 접두사 절단 수준에서도 안정적인 성능을 유지했습니다. 반면, "Nested Aggregation Only"와 같은 베이스라인은 특히 높은 랭크에서 불안정성을 보였으며, 이는 중첩된 훈련 목적 함수 없이는 집계만으로는 불충분하다는 것을 확인시켜 줍니다.
• 차원별 중요도: 차원 중요도 시각화를 통해 PreLort가 작업 관련 신호를 선행 차원에 집중시키는 것(구조화된 분포)을 확인했습니다. 반면 표준 훈련은 균등한 분포를 보였습니다. 이러한 계층적 우선순위 지정은 이질적인 랭크 간의 안정적인 업데이트를 용이하게 합니다.

5. 의의 및 주장

본 논문은 PreLort가 랭크 이질성을 다루는 방식의 전환을 나타낸다고 주장합니다. 즉, 이를 사후적인 집계 문제로 취급하는 것이 아니라, 훈련 과정 자체를 구조화함으로써 해결합니다.

• 핵심 논거: 저자들은 랭크 이질성을 해결하기 위해서는 개선된 집계뿐만 아니라, 랭크 간의 구조화된 표현 정렬이 필요하다고 주장합니다.
• 효율성: 이 방법은 접두사 차원이 일관되게 학습되고 집계되기 때문에, 저계수 클라이언트가 고계수 클라이언트가 기여한 더 풍부한 정보로부터 혜택을 받을 수 있게 합니다.
• 한계에 대한 겸허한 인정: 저자들은 중첩된 훈련 전략으로 인해 배치당 다중 순전파/역전파가 발생하여 약 40%의 추가적인 계산 오버헤드가 발생함을 인정합니다. 또한, 매우 제한적인 로컬 에포크(epoch)를 사용할 경우 모델이 잘 정렬된 저계수 접두사 표현을 완전히 학습하지 못할 수 있다는 점도 언급합니다.

결론적으로, PreLort는 랭크 이질성이 (중첩된 접두사 최적화를 통한) 어댑터 훈련 방식과 (세그먼트별 규칙을 통한) 업데이트 집계 방식을 공동으로 구조화할 때 가장 잘 처리될 수 있음을 보여주며, 이를 통해 기존의 이질적 집계 방법들보다 더 정렬되고 강건한 접두사 표현을 생성합니다.