BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization

기술 요약: BROS: 메모리 효율적인 단일 루프 이층 최적화를 위한 편향 보정 무작위 부분공간

문제 정의
확률적 이층 최적화 (SBO) 는 초매개변수 최적화, 데이터 재가중치, 표현 학습, 그리고 초데이터 정제와 같은 딥러닝 응용을 위한 근본적인 프레임워크입니다. 표준 형식은 상위 변수 $x$와 하위 변수 $Y$(일반적으로 신경망의 매개변수) 를 포함하며, 여기서 상위 목적 함수는 하위 최적화 문제의 해에 의존합니다.

기존의 정확한 단일 루프 SBO 방법들 (예: MA-SOBA) 은 강력한 수렴 보장 ($O(\epsilon^{-2})$ 샘플 복잡도) 을 제공하지만, 하위 네트워크가 대규모일 경우 심각한 메모리 병목 현상을 겪습니다. 이러한 방법들은 보조 변수를 유지하기 위해 전체 공간 헤시안 - 벡터 곱 (HVP) 과 야코비안 - 벡터 곱 (JVP) 에 의존하여 전체 매개변수 수에 비례하는 메모리를 요구합니다. 반면, 메모리 효율적인 대리 방법들 (예: 페널티 방법, 0 차 접근법) 은 메모리 사용량을 줄이지만 종종 정확한 수렴 보장을 희생하거나, 더 강한 가정을 요구하거나, 원래의 이층 목적 함수가 아닌 대리 문제를 최적화합니다.

방법론: BROS
저자들은 다층 행렬 매개변수화 문제를 위해 설계된 메모리 효율적인 단일 루프 SBO 방법인 BROS(Bilevel optimization via RandOmized Subspaces) 를 제안합니다. BROS 는 무작위 저차원 부분공간 내에서 하위 및 보조 수준 업데이트를 수행하면서도 정확한 단일 루프 하이퍼그래디언트 구조를 유지함으로써 메모리 병목 현상을 해결합니다.

BROS 의 핵심 구성 요소는 다음과 같습니다:

1. 무작위 부분공간 분해:
   전체 하위 매개변수 $Y \in \mathbb{R}^{m \times n}$을 업데이트하는 대신, BROS 는 각 반복에서 스케일링된 Haar 프로젝터 $P \in \mathbb{R}^{m \times r}$ (여기서 $r \ll m$) 를 샘플링합니다. 하위 업데이트는 투영된 목적 함수 $\tilde{g}(x, B) = g(x, Y + PB)$를 통해 축소된 부분공간 $B \in \mathbb{R}^{r \times n}$에서 계산됩니다. 전체 공간 변수는 $Y^{k+1} = Y^k + P B^k$로 업데이트됩니다. 이는 기울기, 활성화, 그리고 보조 방향의 메모리 발자국을 $O(mn)$에서$O(rn)$으로 줄입니다.

2. 편향 보정 보조 업데이트:
   보조 변수 $Z$(헤시안을 포함하는 선형 시스템을 해결함) 에 대한 단순한 "스케치 후 리프트" 접근법은 리프트된 헤시안 작용 $P \tilde{H} [P^\top Z]$이 전체 공간 헤시안 작용 $H[Z]$의 편향 없는 추정치가 아니기 때문에 편향을 도입합니다. 이 편향은 하이퍼그래디언트를 왜곡하고 비정상적인 극한으로 이어질 수 있습니다.
   이를 해결하기 위해 BROS 는 Rademacher 양측 프로브 보정을 도입합니다. 알고리즘은 투영으로 인해 발생하는 트레이스 유사 및 전치 유사 잡음 항을 추정하기 위해 독립적인 Rademacher 벡터 (부호 벡터) $u$와 $\xi$를 샘플링합니다. 이러한 프로브는 전체 공간 행렬 형성이 필요 없이 조건부 편향 없는 전체 공간 헤시안 작용을 보장하는 보정된 HVP 추정치 $\widehat{H}[Z]$의 구성을 가능하게 합니다.

3. 단일 루프 재귀:
   BROS 는 상위 변수 $x$, 하위 변수 $Y$, 보조 변수 $Z$, 그리고 이동 평균 하이퍼그래디언트 추정치 $h$에 대한 결합된 단일 루프 재귀를 유지합니다. 업데이트는 투영된 하위 기울기와 편향 보정된 보조 방향을 활용합니다.

주요 기여

• 메모리 효율적 재귀: BROS 는 하위 수준 및 보조 재귀를 무작위 부분공간으로 분해하여 저장된 기울기, 활성화, 그리고 HVP/JVP 계산의 차원을 크게 줄입니다.
• 편향 없는 하이퍼그래디언트 복구: 논문은 리프트된 보조 헤시안 작용에서 발생하는 특정 투영 유도 편향을 식별하고, 계층별 Rademacher 양측 프로브 보정을 제안합니다. 이는 추가적인 구조적 가정 (예: 대각 블록 헤시안) 을 부과하지 않고도 추정치가 전체 다층 헤시안 작용에 대해 편향 없음을 보장합니다.
• 이론적 보장: 표준 가정 (하위 수준의 매끄러움, 강한 볼록성) 하에서 저자들은 BROS 가 $\epsilon$-정상점을 찾기 위한 정확한 단일 루프 방법들 (특히 MA-SOBA) 의 $O(\epsilon^{-2})$ 샘플 복잡도를 유지함을 증명합니다. 수렴 속도는 투영 비율 $\omega_P = \max(m_\ell/r_\ell)$에 의존하며, $r=m$일 때 전체 공간 속도를 회복합니다.

실험 결과
저자들은 네 가지 이층 학습 작업에서 BROS 를 평가했습니다:

1. MNIST 초데이터 정제: BROS 는 84.30% 의 테스트 정확도를 달성하여 전체 공간 MA-SOBA 기준 (84.46%) 과 밀접하게 일치하고 페널티 기반 및 0 차 방법들을 능가했습니다.
2. 데이터 혼합 학습 (280M GPT 스타일 모델): BROS 는 최종 로그 퍼플렉시티 2.8109 대 2.8098 로 MA-SOBA 의 성능을 일치시키면서, MA-SOBA 대비 약 44.9%, 페널티 방법 대비 27% 의 피크 프록시 학습 메모리를 줄였습니다.
3. CIFAR-10 초표현 학습: BROS 는 78.03% 의 최종 테스트 정확도를 달성하여 MA-SOBA(78.57%) 와 비교 가능한 성능을 보였습니다.
4. CIFAR-100 ViT 샘플 재가중치: BROS 는 비교된 방법들 중 가장 높은 검증 (49.88%) 및 테스트 (49.05%) 정확도를 달성했으며, 전체 공간 기준 (23 GB 이상) 에 비해 6.42 GB 의 피크 메모리만 사용했습니다.

의의 및 주장
이 논문은 BROS 가 확률적 이층 최적화에서 메모리 효율성과 이론적 수렴 사이의 간극을 성공적으로 연결했다고 주장합니다. 정확성과 효율성 사이의 트레이드오프를 하는 대리 방법들과 달리, BROS 는 정확한 단일 루프 하이퍼그래디언트 업데이트 구조와 이에 수반되는 $O(\epsilon^{-2})$ 수렴 차수를 유지합니다. 무작위 부분공간과 새로운 편향 보정 메커니즘을 활용함으로써 BROS 는 메모리 제약으로 인해 본래 불가능했을 대규모 신경망에서의 이층 최적화 작업 (데이터 재가중치 및 초매개변수 조정 등) 을 가능하게 하며, 전체 공간 성능을 밀접하게 유지하면서 피크 메모리 사용량을 최대 44.9% 까지 줄입니다.