On the Expressive Power of Permutation-Equivariant Weight-Space Networks

문제 정의
가중치 공간 학습(Weight-space learning)은 다른 신경망의 파라미터를 직접 다루는 신경 구조를 연구하며, 가중치를 정확도 예측, 메타 최적화, 모델 편집과 같은 작업을 위한 구조화된 데이터로 취급합니다. 다층 퍼셉트론(MLP) 가중치 공간의 핵심적인 특성은 내재된 대칭성입니다. 즉, 은닉층 내의 뉴런을 치환(permutation)하면 원시 파라미터는 변하지만 실현된 함수는 변하지 않습니다. 결과적으로, 최첨단(SOTA) 가중치 공간 네트워크들은 이러한 대칭성을 존중하는 치환 불변(permutation-equivariant) 설계를 채택합니다. 그러나 이러한 대칭성 제약은 가설 공간을 제한하여 네트워크의 표현력을 잠재적으로 저하시킬 수 있습니다. 선행 연구들이 특정 능력(예: 순전파 시뮬레이션)을 입증해 왔으나, 치환 불변 가중치 공간 네트워크의 표현력에 대한 포괄적인 이론적 규명은 여전히 미비한 상태입니다. 또한, 다른 기하학적 딥러닝 영역과 달리, 가중치 공간 학습은 가중치 공간과 함수 공간 모두에서 작동하는 사상(map)을 포함하므로 표현력 분석이 특히 까다롭습니다.

방법론
저자들은 MLP 가중치에 작용하는 치환 불변 가중치 공간 네트워크의 표현력을 분석하기 위한 체계적인 이론적 프레임워크를 개발합니다. 방법론은 다음 네 가지 주요 단계로 진행됩니다:

1. 구조적 동등성(Architectural Equivalence): 본 논문은 Deep Weight Space (DWS), Neural Functional Networks (NFN), Graph Meta-Networks (GMN), Neural Graph GNNs (NG-GNN)를 포함한 주요 치환 불변 아키텍처들이 표현력 측면에서 동등함을 먼저 입증합니다. 이는 각 아키텍처가 다른 아키텍처의 기본 레이어를 근사할 수 있음을 보여줌으로써 표현력의 포함 관계(cycle of expressive containment)를 형성함으로써 증명됩니다. Neural Functional Transformers (NFT)는 비표준 어텐션 메커즘으로 인해 초기에는 예외로 식별되었으나, "일반 위치(general position, GP)" 가정 하에서는 다른 아키텍처들과 동일한 표현력을 갖는 것으로 나타났습니다.
2. 근사 프레임워크(Approximation Framework): 저자들은 대상 사상을 분류하기 위해 네 가지 근본적인 근사 설정을 정의합니다:
   • 함수 공간 범함수(Function-space functionals): 함수에서 벡터로 가는 사상 (예: 정확도 예측).
   • 치환 불변 범함수(Permutation-invariant functionals): 가중치에서 벡터로 가는 사상으로서, 뉴런 치환에는 불변하지만 특정 파라미터화에는 의존할 수 있음 (예: 가중치 $\ell_2$-노름).
   • 함수 공간 연산자(Function-space operators): 함수에서 함수로 가는 사상 (예: 모델 편집).
   • 치환 불변 연산자(Permutation-equivariant operators): 가중치에서 가중치로 가는 사상으로서, 대칭성을 존중하지만 실현된 함수에 의해서만 결정되지는 않음 (예: 프루닝 마스크).
3. 보편성 분석(Universality Analysis): 각 설정에 대해, 보편적 근사가 달성 가능한 조건을 분석합니다. 이 분석의 핵심 요소는 "일반 위치(GP)" 가정이며, 이는 각 은닉층 내의 편향(bias) 항들이 서로 쌍마다 구별되어야 함을 요구합니다. 이 가정은 측도가 0인 퇴화된 구성(degenerate configurations) 집합을 제외합니다.
4. 이론적 증명(Theoretical Proofs): 저자들은 분리 성질(separation properties), 정형화 맵(canonization maps), Stone-Weierstrass 정리 등 기하학적 딥러닝 도구를 활용하여 보편성 결과를 증명합니다. 이들은 특정 설정(특히 치환 불변 범함수 및 고정된 아키텍처 상의 연산자)에서 보편성이 일반론적으로는 실패하지만, GP 가정 하에 혹은 출력 아키텍처를 충분히 크게 허용할 경우 회복됨을 입증합니다.

주요 기여

1. 표현력 동등성: 본 논문은 MLP를 위한 모든 주요 치환 불변 가중치 공간 네트워크(DWS, NFN, GMN, NG-GNN)가 동일한 표현력을 가짐을 증명합니다. NFT는 GP 가정 하에서 이 클래스와 동등함이 밝혀졌습니다.
2. 통합 근사 프레임워크: 저자들은 네 가지 자연스러운 근사 설정을 식별하고 공식화하여, 가중치 공간 네트워크를 평가할 수 있는 구조화된 관점을 제공합니다.
3. 보편성 규명:
   • 함수 공간 범함수: 임의의 컴팩트 집합 위에서 불변 가중치 공간 네트워크에 의해 보편적으로 근사 가능함이 증명되었습니다.
   • 치환 불변 범함수: 전체 범위에서는 보편적으로 근사되지 못함이 밝혀졌습니다(동일한 계산 그래프를 가진 비동등한 가중치 구성을 구별하는 데 한계가 있기 때문). 그러나 GP 가정 하에서는 보편성이 달성됩니다.
   • 함수 공간 연산자: 고정된 아키텍처의 출력 네트워크로 제한될 경우 비보편적임이 증명되었습니다(표현 용량의 한계 때문). 출력 아키텍처가 대상 함수의 복잡성을 수용할 만큼 충분히 크다면 보편성이 달성됩니다.
   • 치환 불변 연산자: 치환 불변 범함수와 유사하게, 전체 범위에서는 보편성이 실패하지만 GP 가정 하에서는 성립합니다.
4. 출력 용량 확장(Output Capacity Expansion, OCE)을 통한 실증적 검증: 함수 공간 연산자가 더 큰 출력 용량을 필요로 한다는 이론적 발견에 따라, 저자들은 $k$개의 출력 네트워크 앙상블을 예측하고 이를 평균내는 간단한 수정을 제안합니다. 이는 모델 자체에 파라미터를 추가하지 않고도 유효 출력 용량을 증가시킵니다.

결과

• 이론적 결과: 본 논문은 표현력의 지형도를 완벽하게 제공하여, 가중치 공간 네트워크가 특정 유형의 사상을 근사할 때 언제, 왜 성공하거나 실패하는지를 명확히 합니다. 저자들은 "제외 집합(exclusion set)"(퇴화된 편향 구성)이 불변 설정에서의 보편성을 가로막는 주요 장벽임을 식별했습니다.
• 실증적 결과: 함수 공간 연산자 작업인 MNIST Implicit Neural Representation (INR) dilation 벤치마크에서, 제안된 OCE 수정 방식은 상당한 개선을 보였습니다. 구체적으로, DWS 및 GMN 아키텍처에 OCE를 적용했을 때 이전 SOTA(ScaleGMN + GradMetaNet++) 대비 34%의 향상을 달est했습니다. 이 결과는 유효 출력 용량을 늘리는 것이 함수 공간 연산자 설정에서 식별된 이론적 한계를 직접적으로 해결함을 보여줍니다.

의의
본 논문은 치환 불변 가중치 공간 네트워크의 표현력에 대한 최초의 포괄적인 이론적 토대를 제공한다고 주장합니다. 다양한 아키텍처 패밀리를 단일 표현력 클래스로 통합하고 보편성을 위한 조건을 정밀하게 규명함으로써, 이 연구는 가중치 공간 모델 설계에 대한 원칙적인 가이드를 제공합니다. 출력 용량이 함수 공간 연산자의 병목임을 식별한 것은, 이론에 기반한 실용적인 개선책인 OCE로 이어져 SOTA 성능을 달성했습니다. 저자들은 이 작업이 단순한 순전파 시뮬레이션 능력을 넘어, 가중치 공간에서의 학습에 대한 깊은 수학적 이해로 나아가기 위한 필수적인 단계라고 위치 지으며, 기존의 작업들과 차별화합니다.