기술 요약: 대리 모델 지원 신경 아키텍처 탐색을 위한 코드 지향적 언어 모델 임베딩

문제 제기

신경 아키텍처 탐색 (NAS) 은 고성능 신경 아키텍처의 발견을 자동화하지만, 후보 아키텍처의 적합성을 평가하는 비용이 매우 높다는 심각한 계산 병목 현상에 직면해 있습니다. 이를 완화하기 위해 연구자들은 완전한 학습 없이 적합성을 근사화하는 **대리 모델 (성능 예측기)**을 활용합니다. 그러나 기존 대리 모델 접근법은 두 가지 주요 한계를 겪고 있습니다:

1. 표현의 경직성: 구조 인코딩 (예: 인접 행렬, 경로 인코딩) 은 종종 특정 탐색 공간 (셀 기반 DAG 등) 에 맞춰져 있어 새로운 아키텍처로 일반화하기 어렵거나 커널 크기 (kernel sizes) 와 같은 복잡한 하이퍼파라미터를 포착하는 데 커스텀 엔지니어링이 필요합니다.
2. 파인튜닝 오버헤드: 언어 모델 (LM) 을 범용 회귀기로 사용하려는 최근 시도는 아키텍처를 비표준 텍스트 형식 (예: 파생 트리 문자열 또는 ONNX-to-text) 으로 변환한 후, 이러한 외래 표현에 대해 LM 을 파인튜닝해야 합니다. 이 "워밍업" 단계는 탐색 시작 전 고비용의 사전 학습을 요구하므로, 대리 모델 사용으로 인한 효율성 향상을 무효화합니다.

저자들은 이러한 비효율성은 신경 아키텍처가 본질적으로 실행 가능한 코드로 존재한다는 사실을 현재 방법들이 간과하기 때문에 발생한다고 주장합니다. LM 은 방대한 코드 코퍼스로 사전 학습되었으므로, 인공적인 텍스트 인코딩을 입력으로 제공할 때 발생하는 구조적 논리 이해에 대한 선재적 편향 (inductive bias) 불일치를 해결할 수 있는 고유한 능력을 갖추고 있습니다.

방법론: COLE

저자들은 NAS 를 위한 제로샷 특징 추출기로 고정된 사전 학습 언어 모델을 활용하는 코드 지향적 언어 모델 임베딩 (COLE) 전략을 제안합니다.

1. 네이티브 코드 문자열화

아키텍처를 추상적인 문자열로 변환하는 대신, COLE 는 탐색 공간을 직접 PyTorch 클래스 정의(torch.nn.Module) 로 매핑합니다.

• NAS-Bench-201: 저자들은 표준 그래프 토폴로지 문자열을 구문 분석하여 유효한 Python 클래스로 변환합니다. 세 가지 상세도 (verbosity) 수준을 조사했습니다:
  • 헬퍼 메서드: 복잡한 연산을 별도의 클래스로 정의.
  • 인라인: nn.Sequential 을 사용하여 클래스 내부에 논리 정의.
  • 헬퍼 제외: 구현을 정의하지 않고 헬퍼 클래스를 참조 (가장 토큰 효율적).
• einspace: 결정론적 트랜스파일러가 문맥 자유 문법 파생 트리를 직접 실행 가능한 PyTorch 코드로 변환하며, 비단말 노드 (예: Branching, Routing) 를 사용자 정의 컨테이너 클래스로, 말단 노드를 원시 연산자로 매핑합니다.

2. 고정 임베딩 추출

시스템은 원시 코드 텍스트를 처리하기 위해 고정된 사전 학습 LM(특히 CodeLlama Python 7B) 을 활용합니다.

• 파인튜닝 부재: LM 파라미터는 완전히 고정되어 도메인별 사전 학습의 계산 비용을 제거합니다.
• 풀링: 토큰 시퀀스 전반에 걸쳐 **평균 풀링 (Mean Pooling)**을 사용하여 마지막 레이어의 숨겨진 상태 활성화 값을 집계하여 고정 차원의 임베딩 벡터를 생성합니다.

3. 대리 헤드

고차원 임베딩 (예: 4096 차원) 은 경량 회귀 파이프라인을 통해 처리됩니다:

• 차원 축소: 과적합을 방지하기 위해 주성분 분석 (PCA) 을 사용하여 임베딩을 128 차원으로 축소합니다.
• 회귀: 3 층 다층 퍼셉트론 (MLP) 이 아키텍처 성능을 예측합니다.
• 손실 함수: 모델은 **쌍대 힌지 손실 (Pairwise Hinge Loss)**을 사용하여 켄달의 타우 (Kendall's Tau) 순위 상관관계에 최적화됩니다. 이는 진화적 탐색 알고리즘에 필수적인 후보의 정확한 순서를 예측하는 데 중점을 두며, 정확한 절대 정확도 값 예측보다는 순위 예측을 우선시합니다.

주요 기여

1. 네이티브 코드의 우수성: 저자들은 고정된 LM 을 사용할 때 아키텍처를 원시 PyTorch 코드로 표현하는 것이 다른 텍스트 기반 인코딩 (ONNX-to-text, 파생 트리) 보다 더 높은 예측 성능 (켄달의 타우) 을 제공함을 입증했습니다.
2. 파인튜닝 제거: COLE 는 효과적인 NAS 임베딩을 위해 계산 비용이 많이 드는 파인튜닝이 반드시 필요하지 않음을 증명하며, 오프더셸 (off-the-shelf) LM 이 "아웃오브더박스"로 경쟁력 있는 특징 추출기로 작용함을 보여줍니다.
3. SA-NAS 통합: COLE 를 BANANAS 대리 모델 지원 탐색 알고리즘에 통합하면 표준 구조 경로 인코딩을 사용하는 경우보다 탐색 효율성이 크게 향상됩니다.
4. 실용적 튜닝 가이드라인: 이 연구는 토큰 효율성을 위한 "헬퍼 제외" 모드 사용과 차원 축소를 위한 PCA 와 같은 효과적인 구성을 문서화하여, 입력 문자열에 전체 구현 세부 사항이 없어도 LM 이 연산 의미론을 추론할 수 있음을 보여줍니다.

실험 결과

이 프레임워크는 NAS-Bench-201과 einspace 탐색 공간에서 평가되었습니다.

• 임베딩 품질:

  • NAS-Bench-201 에서 COLE(CodeLlama 7B 사용) 은 879 개의 학습 샘플로 0.703의 켄달의 타우를 달성하여 ONNX-to-text 인코딩 (0.678) 보다 우수했으며, einspace 의 파생 트리 문자열 (0.763) 과 비교하여 PyTorch 코드 (0.778) 가 더 높은 성능을 보였습니다.
  • 모델 선택: 코드 중심 모델 (CodeLlama, Yi Coder) 은 일관되게 일반 인코더 전용 모델 (ModernBERT, CodeBERT) 보다 우수한 성능을 보였습니다. 흥미롭게도 작은 모델 (1.5B) 이 더 큰 7B/34B 모델과 유사한 성능을 보여 거대한 규모가 반드시 필요하지 않음을 시사합니다.
  • 손실 함수: 평균 제곱 오차 (MSE) 가 절대 오차는 낮았음에도 불구하고, 쌍대 힌지 손실이 표준 MSE 보다 더 나은 순위 상관관계를 제공했습니다.

• 탐색 효율성 (BANANAS):

  • BANANAS 알고리즘에서 경로 인코딩을 COLE 로 대체하면 최적 성능에 도달하는 데 필요한 평가 예산이 감소했습니다.
  • CIFAR-100: COLE 는 132 번의 평가로 가장 적합한 아키텍처의 1% 이내 성능에 도달하여, 경로 인코딩이 요구하는 200 번 평가 대비 34% 감소를 달성했습니다.
  • CIFAR-10: 13% 감소 (52 대 60 번 평가).
  • ImageNet16-120: 18% 감소 (164 대 200 번 평가).
  • 이러한 개선은 대부분의 탐색 궤적에서 통계적으로 유의미했습니다 ($p < 0.05$).

중요성과 주장

이 논문은 COLE 가 사전 학습 LM 의 기존 코드 지식을 활용하여 NAS 를 발전시키기 위한 다용도적이고 효율적인 기반을 확립한다고 주장합니다.

• 탐색 중립성: 모든 신경 아키텍처는 코드로 표현될 수 있으므로, COLE 는 구조적 인코딩과 달리 새로운 탐색 공간을 위해 인코딩 체계를 재설계할 필요가 없습니다.
• 콜드스타트 능력: 파인튜닝 필요성을 제거함으로써 COLE 는 대리 모델이 즉시 효과적인 "콜드스타트" 탐색을 가능하게 하여, 이전 LM 기반 NAS 방법들의 효율성 역설을 해결합니다.
• 일반화 능력: 이 접근법은 토폴로지와 하이퍼파라미터 표현을 단일 형식 (코드) 으로 통합하여 LM 이 구조적 흐름과 매개변수 값 (예: kernel_size=3) 을 모두 자연스럽게 해석할 수 있게 합니다.

저자들은 현재 평가가 특정 벤치마크 (NAS-Bench-201, einspace) 로 제한되지만, 이 방법론이 특히 리소스 제약이 있는 NAS 시나리오에서 복잡한 GNN 또는 파인튜닝된 트랜스포머에 대한 실용적이고 저비용의 대안을 제공한다고 결론지었습니다.