[논문 리뷰] Equivariant Architectures for Learning in Deep Weight Spaces

Aviv Navon, Aviv Shamsian외 4명|arXiv (Cornell University)|2023. 01. 30.

Machine Learning in Materials Science인용 수 5

한 줄 요약

이 논문은 사전 훈련된 MLP의 가중치 및 편향 행렬에서 직접 작동하도록 설계된 새로운 신경망 아키텍처인 딥 웨이트 스페이스 네트워크(DWSNets)를 소개한다. 이 방법은 이러한 가중치의 자연스러운 순열 대칭성에 대해 동치성을 유지함으로써 작동한다. 이론적으로 아핀 동치성 보장된 레이어는 풀링, 브로드캐스팅, 완전 연결 연산을 통해 특징화되며, 이는 매개변수 효율적인 학습을 가능하게 하고 일반화 오차 예측 및 INR 분류와 같은 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Designing machine learning architectures for processing neural networks in their raw weight matrix form is a newly introduced research direction. Unfortunately, the unique symmetry structure of deep weight spaces makes this design very challenging. If successful, such architectures would be capable of performing a wide range of intriguing tasks, from adapting a pre-trained network to a new domain to editing objects represented as functions (INRs or NeRFs). As a first step towards this goal, we present here a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trained MLP and processes it using a composition of layers that are equivariant to the natural permutation symmetry of the MLP's weights: Changing the order of neurons in intermediate layers of the MLP does not affect the function it represents. We provide a full characterization of all affine equivariant and invariant layers for these symmetries and show how these layers can be implemented using three basic operations: pooling, broadcasting, and fully connected layers applied to the input in an appropriate manner. We demonstrate the effectiveness of our architecture and its advantages over natural baselines in a variety of learning tasks.

연구 동기 및 목표

• 사전 훈련된 신경망의 가중치를 원본 행렬 형태로 직접 처리할 수 있는 딥 러닝 아키텍처를 설계하는 것.
• 복잡한 대칭성 구조를 지닌 웨이트 스페이스에서의 학습에 도전하는 것.
• MLP 가중치 행렬에 내재된 순열 대칭성을 존중하는 원칙적인 프레임워크를 개발하는 것.
• 웨이트 스페이스에서 직접 모델 적응, 프루닝, 표현 학습 등의 후행 작업을 가능하게 하는 것.

제안 방법

• 아키텍처는 DWS-레이어—MLP 가중치의 순열 대칭성을 존중하는 아핀 동치성 보장 선형 레이어—기반으로 구성된다.
• 이 레이어는 입력 네트워크의 특정 가중치 및 편향 하위공간 간의 맵핑을 이루는 블록 행렬로 특징지어진다.
• 이 레이어는 세 가지 핵심 연산을 사용해 구현된다: 풀링, 브로드캐스팅, 일반적인 완전 연결 레이어.
• 점별 비선형성과 함께 스택된 DWS-레이어를 사용하여 웨이트 스페이스 처리를 위한 딥 아키텍처를 형성한다.
• 이 방법은 동치성을 보장한다: 입력 MLP의 뉴런 순서를 바꾸면 출력도 해당하는, 예측 가능한 변환을 겪는다.
• 구조화된 가중치 공유 덕분에 완전 연결 기반 대비 훨씬 적은 매개변수를 사용하는 매개변수 효율적인 아키텍처이다.

실험 결과

연구 질문

• RQ1 MLP 가중치 행렬의 순열 대칭성에 대해 동치성이 보장되는 선형 레이어에 대한 수학적 제약 조건은 무엇인가?
• RQ2 기본 딥 러닝 연산을 사용해 이러한 동치성 레이어를 효율적으로 구현할 수 있는 방법은 무엇인가?
• RQ3 DWSNet은 일반화 오차 예측이나 조밀한 표현 학습과 같은, 사전 훈련된 MLP의 공간에서 함수를 효과적으로 학습할 수 있는가?
• RQ4 원본 가중치 행렬을 다루는 작업에서 표준 기반 대비 아키텍처는 어떻게 비교되는가?
• RQ5 DWSNets는 INRs나 NeRFs와 같은 다른 아키텍처로 일반화될 수 있는가?

주요 결과

• DWSNets는 CIFAR-10 및 MNIST 데이터셋에서 가중치로부터 모델 성능을 예측할 때 가장 낮은 일반화 오차를 기록했으며, 모든 기반 모델을 압도했다.
• 모델은 주파수와 진폭에 따라 모델을 그룹화하는 2차원 조밀한 표현을 학습하여 해석 가능성을 입증했다.
• 가중치-가중치 및 편향-편향 블록을 모두 추가하면 하나만 사용할 때보다 성능 향상이 있었으며, 대각선 버전은 거의 최적의 성능을 달성했다.
• 데이터 증강과 배치 정규화가 성능 향상에 크게 기여했으며, MNIST INRs 분류 작업에서 테스트 정확도를 77.20%에서 85.71%로 상승시켰다.
• 아키텍처는 뛰어난 표현 능력을 보였으며, 학습된 레이어를 통해 입력 MLP의 전방향 전파를 근사할 수 있었다.
• 프루닝 및 INR 처리의 도전 과제가 있었음에도 불구하고, 적절한 초기화와 도메인 인식 정규화를 통해 방법의 잠재력이 확인되었다.