NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces

이 논문은 그래프 하이퍼네트워크를 통해 가중치 공간을 구조적으로 정렬하고 패치 기반 처리를 가능하게 하여, 훈련 시 보지 못한 다양한 아키텍처에서도 성공적으로 작동하는 너비 무관형 신경망 생성 모델인 NNiT 를 제안합니다.

핵심

1. 문제: "레고 블록"과 "무작위 나열"의 딜레마

기존에 AI 의 두뇌 (신경망) 를 생성하는 방식에는 두 가지 큰 문제가 있었습니다.

1. 고정된 크기 (Width) 의 문제:

   • 비유: 마치 "3 단짜리 레고 성"을 만드는 방법을 배웠는데, 갑자기 "5 단짜리 성"을 만들어달라고 하면 전혀 못 만드는 상황입니다.
   • 현실: 기존 기술은 AI 의 층 (Layer) 이 몇 개인지, 각 층에 뉴런이 몇 개 있는지 미리 정해져 있어야만 작동했습니다. 훈련할 때 본 적 없는 크기의 AI 는 만들 수 없었습니다.

2. 순서 섞임 (Permutation Symmetry) 의 문제:

   • 비유: 같은 레고 성을 만들더라도, 블록을 쌓는 순서만 바꾼다면 완성된 성의 모양은 똑같지만, 블록들의 위치는 완전히 달라집니다. AI 는 이 '순서'가 다르면 같은 기능을 하는 두뇌도 서로 다른 데이터로 인식해 혼란을 겪습니다.
   • 현실: AI 의 뉴런 순서를 바꿔도 기능은 똑같지만, 컴퓨터는 이를 전혀 다른 데이터로 보기 때문에 패턴을 배우기 매우 어렵습니다.

2. 해결책: "NNiT"라는 새로운 건축가

이 논문은 이 두 문제를 해결하기 위해 **NNiT (Neural Network Diffusion Transformers)**라는 새로운 시스템을 제안합니다.

① 구조 정렬: "나침반이 있는 지도" (Graph HyperNetworks)

먼저, AI 의 두뇌를 만들 때 뉴런의 순서가 뒤죽박죽이 되지 않도록 **GHN(그래프 하이퍼네트워크)**이라는 도구를 사용합니다.

• 비유: 기존 방식은 "무작위로 흩어진 레고 조각"을 줬다면, GHN 은 **"어떤 조각이 어디에 붙어야 하는지 정해진 나침반이 달린 지도"**를 줍니다.
• 효과: 이렇게 하면 AI 의 두뇌가 항상 일정한 규칙 (국소적 상관관계) 을 따르게 되어, 컴퓨터가 패턴을 쉽게 배우게 됩니다.

② 패치 토크나이제이션: "이미지 조각"처럼 자르기

가장 혁신적인 부분은 AI 의 두뇌를 '하나의 긴 줄'로 보지 않고, '작은 정사각형 조각 (패치)'들의 집합으로 본다는 점입니다.

• 비유:
  • 기존 방식: AI 의 두뇌를 "길이가 100 인치인 긴 줄"로 보았습니다. 길이가 100 인치가 아닌 120 인치면 줄이 끊어지고 아무것도 못 했습니다.
  • NNiT 방식: AI 의 두뇌를 **"작은 스티커 (패치)"**로 봅니다.
    • "3 단짜리 성"을 만들려면 스티커 3 개를 붙이면 됩니다.
    • "5 단짜리 성"을 만들려면 스티커 5 개를 붙이면 됩니다.
    • 핵심: 스티커의 크기나 모양은 변하지 않고, 붙이는 개수만 늘리면 됩니다.
• 결과: 훈련할 때 본 적 없는 크기의 AI 도, 필요한 만큼 스티커를 더 붙여주면 바로 작동하는 두뇌가 됩니다. 이를 'Zero-Shot (제로 샷)' 생성이라고 합니다.

3. 어떻게 작동하나요? (멀티모달 시퀀스)

NNiT 는 **AI 의 구조 (어떤 모양인지)**와 **AI 의 두뇌 (어떤 데이터인지)**를 동시에 학습합니다.

• 비유: 건축가가 "집의 설계도 (구조)"와 "집을 짓는 벽돌 (가중치)"을 한 번에 보고, "이런 모양의 집에는 이런 벽돌이 필요해"라고 학습하는 것입니다.
• 효과: 사용자가 "이런 크기로 만들어줘"라고 하면 바로 만들어주고, "아무거나 좋은 거 만들어줘"라고 해도 작동하는 AI 를 만들어냅니다.

4. 실제 성과: 로봇이 잘하는가?

연구진은 이 기술을 로봇 팔을 조종하는 AI 에 적용해 보았습니다.

• 상황: 로봇이 큐브를 잡고, 밀고, 쌓는 작업을 시켰습니다.
• 결과:
  • 기존 기술들은 훈련할 때 본 적 없는 크기의 로봇 두뇌를 만들면 대부분 실패했습니다 (성공률 0~60%).
  • 하지만 NNiT는 훈련할 때 본 적 없는 크기의 두뇌를 만들어도 85% 이상 성공했습니다.
  • 특히 로봇은 아주 작은 오차도 치명적이기 때문에, 이 결과가 얼마나 놀라운지 알 수 있습니다.

5. 요약: 왜 이것이 중요한가?

이 기술은 **"AI 의 크기에 구애받지 않는 유연한 AI 생성"**을 가능하게 합니다.

• 과거: "이 크기의 AI 만 만들 수 있어. 다른 건 못 해."
• NNiT: "네가 원하는 크기의 AI 라도, 내가 알아서 맞춰서 만들어 줄게. 처음 보는 모양도 문제없어!"

마치 레고를 다룰 때, 특정 모양만 만들 수 있는 게 아니라, 어떤 모양이든 필요한 블록 수만큼만 더 붙이면 바로 완성되는 마법 같은 기술을 개발한 것과 같습니다. 이는 앞으로 로봇, 자율주행, 그리고 다양한 하드웨어에 맞춰 AI 를 즉시 배포하는 시대를 열 수 있는 중요한 발걸음입니다.

기술 요약

NNiT: 구조적으로 정렬된 가중치 공간을 활용한 너비 무관 (Width-Agnostic) 신경망 생성

1. 문제 정의 (Problem Statement)

기존의 신경망 파라미터 생성 (Generative Modeling) 은 두 가지 근본적인 한계에 직면해 있습니다.

• 고정된 차원 의존성: 표준 파라미터 표현은 알려진 가중치 행렬의 차원에 의존합니다. 기존 방법들은 가중치를 고정된 크기의 1 차원 벡터로 평탄화 (Flattening) 하는 경향이 있어, 학습 시 보지 못한 너비 (Width) 를 가진 새로운 아키텍처로 일반화하기 어렵습니다.
• 치환 대칭성 (Permutation Symmetry): 신경망의 입력 - 출력 함수는 동일하게 유지되면서도 뉴런의 순서를 임의로 바꿀 수 있습니다 (치환 대칭성). 이로 인해 인접한 가중치들 간의 공간적 상관관계가 깨지고, 가중치 공간이 정렬되지 않아 패치 기반 (Patch-based) 생성 모델링이 불가능해집니다.

2. 방법론 (Methodology)

이 논문은 **Neural Network Diffusion Transformers (NNiT)**를 제안하며, 다음과 같은 핵심 기법들을 통해 위 문제를 해결합니다.

가. 구조적 정렬을 위한 그래프 하이퍼네트워크 (GHN) 활용

• CNN 디코더가 탑재된 GHN: 기존 SGD(확률적 경사 하강법) 로 학습된 모델은 가중치 공간이 정렬되지 않은 반면, 저자들은 CNN 디코더를 가진 그래프 하이퍼네트워크 (GHN) 를 데이터 생성기이자 정렬 메커니즘으로 사용합니다.
• 국소적 상관관계 유도: GHN 은 아키텍처 그래프를 통해 정보를 전파하고, CNN 디코더를 통해 가중치 텐서를 생성합니다. 이 과정에서 **명시적인 국소성 편향 (Locality Bias)**이 가중치 공간에 부여되어, 서로 다른 시드 (Seed) 에서 생성된 모델들 간에도 일관된 국소적 공간 상관관계 (예: 수직 밴딩 구조) 가 형성됩니다. 이는 패치 기반 토큰화 (Tokenization) 에 필수적인 전제 조건을 충족시킵니다.

나. 너비 무관 (Width-Agnostic) 패치 토큰화

• 패치 기반 표현: NNiT 는 가중치 행렬을 고정된 벡터가 아닌 $p \times p$ 크기의 **패치 (Patch)**로 분할하여 토큰화합니다.
• 다중 모드 시퀀스 모델링: 이 접근법은 아키텍처 토큰 (이산적, Discrete) 과 가중치 패치 (연속적, Continuous) 를 단일 시퀀스로 통합합니다.
  • 아키텍처 토큰: 레이어의 너비를 나타내는 이산적 토큰 시퀀스.
  • 가중치 패치: 가중치 텐서를 공간적으로 상관관계가 있는 패치로 변환한 연속적 임베딩.
• 너비 확장성: 레이어의 너비가 변하더라도 토큰화 방식은 변하지 않습니다. 단순히 더 많은 패치를 생성하면 되므로, 학습 중 보지 못한 너비를 가진 아키텍처에도 제로샷 (Zero-shot) 으로 적용 가능합니다.

다. 혼합 노이즈 레벨 (MoNL) 을 통한 학습

• Joint Generation Mode: 아키텍처와 가중치를 동시에 확산 (Diffuse) 하여 $p(a, w)$ 결합 분포를 학습합니다.
• Conditional Synthesis Mode: 아키텍처 토큰은 노이즈가 없는 상태로 유지하고 가중치만 확산시켜, 주어진 아키텍처에 대한 가중치 합성 $p(w|a)$을 학습합니다.

3. 주요 기여 (Key Contributions)

1. 가중치 공간의 구조적 정렬 증명: CNN 디코더가 탑재된 GHN 이 치환 대칭성을 줄이고 일관된 국소적 공간 상관관계를 가진 가중치 분포를 생성함을 실험적으로 입증했습니다.
2. 너비 무관 생성을 위한 패치 토큰화: 가중치를 패치 단위로 토큰화하여 생성 과정이 고정된 행렬 차원에 종속되지 않도록 하였습니다. 이를 통해 학습 중 보지 못한 아키텍처 토폴로지에 대한 제로샷 합성이 가능해졌습니다.
3. NNiT 프레임워크 제안: 아키텍처와 가중치를 단일 다중 모드 시퀀스로 모델링하는 확산 트랜스포머 (Diffusion Transformer) 를 개발하여, 결합 생성 (Joint Generation) 과 조건부 합성 (Conditional Synthesis) 을 모두 지원합니다.

4. 실험 결과 (Results)

실험은 로봇 제어 태스크 (ManiSkill3 환경) 에서 MLP 정책을 생성하는 데 적용되었습니다.

• 구조적 정렬 검증 (Table 1 & Figure 3): GHN 으로 생성된 가중치는 SGD 모델과 유사한 성능 (>99% 성공률) 을 내면서도, 시드 간 일관된 공간적 구조 (수직 밴딩) 를 보여줍니다. 이는 모드 붕괴 (Mode Collapse) 없이 가중치 다양성을 유지하면서 정렬이 이루어졌음을 의미합니다.
• 제로샷 너비 전이 (Zero-Shot Width Transferability, Table 2):
  • 학습된 아키텍처: NNiT 와 기존 방법 (D2NWG, SANE) 모두 높은 성능을 보였습니다.
  • 미학습 아키텍처 (Unseen Topologies): 기존 방법들은 성능이 급격히 저하되었습니다 (성공률 0~59%). 반면, NNiT 는 모든 태스크에서 85% 이상의 성공률을 유지하며 강력한 일반화 능력을 입증했습니다.
  • 특히 SANE 은 다양한 너비의 혼합 학습 시 위치 인코딩의 불안정성으로 인해 실패했습니다.
• 다중 모드 결합 생성 (Table 3 & 4): NNiT 는 고정된 아키텍처 프롬프트 없이도 아키텍처와 가중치를 동시에 생성하여 높은 성공률 (99~100%) 을 달성했습니다. 학습 데이터에 존재하지 않는 완전히 새로운 아키텍처 (예: [32, 16, 16, 16]) 에 대해서도 98% 성공률을 보이며 구조적 논리를 내재화했음을 확인했습니다.

5. 의의 및 의의 (Significance)

• 신경망 생성의 패러다임 전환: 신경망 생성을 고정된 벡터 공간의 문제가 아닌, 이미지 생성과 유사한 공간적 필드 (Spatial Field) 문제로 재정의했습니다.
• 실용적 적용 가능성: 로봇 공학 및 임베디드 AI 분야에서 하드웨어 제약 (연산량, 메모리 등) 에 맞춰 유연하게 네트워크를 생성할 수 있는 기반을 마련했습니다.
• 확장성: 패치 기반 토큰화 방식은 비디오 생성 모델의 효율성 최적화 (선형 어텐션 등) 를 차용하여, 향후 수십억 파라미터 규모의 파운데이션 모델 생성으로 확장 가능한 잠재력을 가지고 있습니다.

결론적으로, NNiT 는 신경망 생성의 가장 큰 장벽이었던 '치환 대칭성'과 '고정된 차원 의존성'을 동시에 해결하여, 학습 데이터에 존재하지 않는 다양한 구조와 너비를 가진 고성능 신경망을 실시간으로 합성할 수 있는 새로운 가능성을 제시했습니다.