기술 요약: TAS-LoRA

문제 정의

Transformer 아키텍처 검색 (TAS) 은 디자인에 필요한 수동 노력을 줄이기 위해 최적의 Vision Transformer(ViT) 아키텍처를 자동으로 발견하는 것을 목표로 합니다. 현재 최첨단 TAS 방법들은 주로 모든 가능한 서브넷을 포함하는 "슈퍼넷"을 학습하는 원샷 (one-shot) 접근법에 의존합니다. 이러한 가중치 얽힘 (weight-entangled) 슈퍼넷에서는 모든 서브넷이 공통의 가중치 세트를 공유합니다.

이 논문은 이러한 패러다임의 치명적인 한계를 지적합니다: 특징 붕괴 (feature collapse) 문제입니다. 슈퍼넷이 모든 서브넷에 걸친 전역 손실을 최소화하도록 최적화되기 때문에, 공유된 가중치는 다양한 아키텍처에 적절하게 작동하는 일반화된 특징을 학습하는 경향이 있지만, 서브넷 고유의 특성을 포착하지는 못합니다. 결과적으로 개별 서브넷은 고유한 아키텍처적 특성 (예: 특정 수의 어텐션 헤드나 MLP 비율) 을 완전히 활용하지 못하여, 처음부터 독립적으로 학습된 서브넷에 비해 성능이 최적화되지 못합니다. 논문의 실증 분석에 따르면, 가중치 얽힘 슈퍼넷 내의 서브넷들은 거의 동일한 특징 표현 (높은 코사인 유사도) 을 보이는 반면, 독립적으로 학습된 서브넷들은 명확하게 구분된 표현을 개발합니다.

방법론: TAS-LoRA

특징 붕괴를 해결하면서도 계산 효율성을 유지하기 위해, 저자들은 TAS-LoRA(Mixture-of-LoRA Experts 를 활용한 Transformer 아키텍처 검색) 를 제안합니다. 이 방법은 슈퍼넷 학습 과정에 저랭크 적응 (LoRA) 을 통합하여, 모든 가능한 서브넷 구성에 고유한 매개변수를 할당하는 prohibitive 한 비용을 치르지 않고도 각 서브넷을 위한 전문화된 특징 학습을 가능하게 합니다.

1. Mixture-of-LoRA-Experts(MoLE)

매우 광범위한 검색 공간으로 인해 모든 서브넷에 고유한 LoRA 모듈을 할당하는 것은 계산적으로 불가능하므로, TAS-LoRA 는 Mixture-of-LoRA-Experts(MoLE) 전략을 사용합니다.

• 아키텍처: LoRA 전문가 (experts) 는 슈퍼넷의 선형 레이어 (QKV 프로젝션 및 MLP 레이어 포함) 에 적용됩니다.
• 동적 라우팅: 경량 라우터가 샘플링된 서브넷의 특정 아키텍처에 기반하여 제한된 수의 LoRA 전문가들을 동적으로 선택하고 결합합니다.
• 추론 효율성: 추론 시에는 선택된 LoRA 매개변수가 원래 슈퍼넷 가중치에 병합됩니다. 서브넷 구성은 추론 중에 고정되므로 라우터는 한 번만 실행되며, 병합된 가중치는 추가적인 계산 오버헤드를 발생시키지 않습니다.

2. 라우터 설계

라우터는 트랜스포머 블록의 순차적 의존성을 포착하도록 설계되었습니다:

• 입력: 블록 수준 속성(예: 어텐션 헤드 수, MLP 비율)과 서브넷 수준 속성(예: 임베딩 차원, 깊이) 을 모두 입력으로 받습니다.
• 메커니즘: 이러한 속성들은 학습 가능한 임베딩 레이어에 의해 인코딩되고, 블록 간의 순차적 의존성을 포착하기 위해 LSTM에 의해 처리됩니다. LSTM 은 은닉 표현을 출력하며, 이는 완전 연결된 "라우터 헤드"를 통과하여 각 레이어의 사용 가능한 LoRA 전문가들에 대한 확률 분포를 생성합니다.

3. 그룹별 라우터 초기화

저자들은 학습 초기 단계에서 라우터가 모든 전문가에 유사한 가중치를 할당하여 중복을 초래하고 전문화를 방해한다는 중요한 문제를 발견했습니다. 이를 완화하기 위해 저자들은 그룹별 라우터 초기화 기법을 도입합니다:

• 그룹화: 서로 다른 서브넷의 트랜스포머 블록은 아키텍처적 유사성 (특히 어텐션 헤드 수, MLP 비율, 임베딩 차원) 을 기반으로 그룹화됩니다.
• 초기화: LoRA 전문가의 수는 그룹 수와 일치하도록 설정됩니다. 라우터는 각 전문가가 서로 다른 그룹으로 편향되도록 초기화됩니다. 이는 해당 그룹 - 전문가 쌍에 대해 라우터 헤드의 특정 편향 항을 양수 $\beta$로 설정하고 나머지는 0 으로 설정함으로써 달성됩니다.
• 효과: 이는 학습 초기에 라우터가 서로 다른 아키텍처 그룹에 서로 다른 전문가를 할당하도록 강제하여 다양한 특징 학습을 장려합니다. 학습이 진행됨에 따라 라우터는 서브넷의 속성에 기반하여 이러한 할당을 동적으로 업데이트합니다.

주요 기여

1. TAS 내 특징 붕괴의 식별: 이 논문은 가중치 얽힘 슈퍼넷에서의 특징 붕괴 문제를 공식적으로 식별하고 분석하여, 공유된 가중치가 서브넷이 전문화된 표현을 학습하는 것을 방해함을 입증합니다.
2. TAS-LoRA 프레임워크: LoRA 를 동적 라우팅 메커니즘 (MoLE) 과 결합하여 서브넷 고유의 특징을 효율적으로 학습하는 새로운 방법입니다. 이는 전체 슈퍼넷을 처음부터 다시 학습하지 않고도 서브넷이 아키텍처적 속성에 기반하여 공유된 특징을 정제할 수 있게 합니다.
3. 그룹별 라우터 초기화: 학습 초기 단계에서 전문가 중복을 방지하여 LoRA 전문가들이 서로 다른 특징 표현을 학습하도록 보장하는 기법입니다.
4. 최첨단 성능: 광범위한 실험을 통해 TAS-LoRA 가 다양한 검색 공간과 데이터셋에서 기존 TAS 방법들 (예: AutoFormer, DYNAS, FocusFormer) 보다 현저히 우수한 성능을 보임을 입증했습니다.

실험 결과

저자들은 ImageNet 데이터셋과 여러 전이 학습 벤치마크 (CIFAR-10/100, Flowers, CARS, INAT-19) 에서 TAS-LoRA 를 평가했습니다.

• ImageNet 성능: TAS-LoRA 는 Tiny, Small, Base 검색 공간 전반에 걸쳐 AutoFormer 와 DYNAS 로 학습된 슈퍼넷의 top-1 정확도를 일관되게 향상시켰습니다.
  • 예를 들어, AutoFormer-T 공간에서 TAS-LoRA 는 **75.7%**의 정확도를 달성하여 원래 AutoFormer(74.7%) 와 유연성이 낮은 PreNAS-Ti(77.1%) 와 같은 다른 최첨단 방법들을 능가했습니다.
  • AutoFormer-B 공간에서 TAS-LoRA-B 는 54M 개의 매개변수로 **82.6%**의 정확도를 달성하여 AutoFormer-B(82.4%) 를 능가했고, 96M 개의 매개변수를 가진 더 큰 모델인 GLiT-B(82.3%) 와도 맞먹거나 능가했습니다.
• 전이 학습: TAS-LoRA 는 우수한 일반화 능력을 보여주어 CIFAR-10/100, Flowers, CARS, INAT-19 에서 경쟁력 있거나 최첨단 결과를 달성했습니다.
• Ablation 연구:
  • 초기화: 그룹별 초기화가 무작위 초기화보다 일관되게 우월한 성능을 보여, 전문가 중복을 줄이는 데 그 역할이 있음을 확인했습니다.
  • 특징 유사성: 코사인 유사도 분석 결과, 그룹별 초기화를 적용한 TAS-LoRA 는 무작위 초기화나 표준 가중치 얽힘 방법에 비해 전문가 간에 훨씬 더 다양한 특징 표현을 생성함을 보여주었습니다.
  • 파인튜닝과의 비교: TAS-LoRA 는 MoLE 레이어에서 50 에포크만 학습했음에도 불구하고 모든 가중치를 업데이트한 완전 파인튜닝된 슈퍼넷보다 우수한 성능을 보였습니다. 이는 성취된 이득이 단순히 학습 기간 연장에 기인한 것이 아니라 MoLE 의 서브넷 고유의 특징 학습 능력에서 비롯됨을 시사합니다.

중요성과 주장

이 논문은 TAS-LoRA 가 가중치 얽힘 슈퍼넷에 내재된 특징 붕괴 문제를 효과적으로 해결함으로써 Transformer 아키텍처 검색 분야에서 중요한 진전을 이루었다고 주장합니다.

• 효율성 대 성능: 여러 슈퍼넷을 학습해야 하는 방법 (예: MoS) 이나 하드웨어 제약이 변경될 때 처음부터 다시 학습해야 하는 방법 (예: PreNAS) 과는 달리, TAS-LoRA 는 서브넷별 적응 레이어를 갖춘 단일 슈퍼넷을 학습합니다. 이는 전체 모델을 다시 학습하지 않고도 새로운 하드웨어 제약 (예: FLOPs, 지연 시간) 에 유연하게 적응할 수 있게 합니다.
• 확장성: 제한된 수의 전문가와 동적 라우터를 사용함으로써, 이 방법은 매개변수나 계산 비용의 선형적 증가 없이 방대한 수의 가능한 서브넷 구성을 가진 슈퍼넷으로 확장됩니다.
• 일반화 가능성: 이 접근법은 근본적인 학습 전략 (예: AutoFormer 또는 DYNAS) 에 관계없이 성능을 향상시키는, 어떤 가중치 얽힘 슈퍼넷에도 적용 가능한 일반적인 솔루션으로 제시됩니다.

저자들은 TAS-LoRA 가 특징 붕괴를 성공적으로 완화하여 서브넷이 전문화된 표현을 학습하게 하고, 아키텍처 검색 및 전이 학습 작업 모두에서 최첨단 성능을 달성했다고 결론지었습니다.