[논문 리뷰] Neural Execution Engines: Learning to Execute Subroutines

Yujun Yan, Kevin Swersky외 3명|arXiv (Cornell University)|2020. 06. 15.

Topic Modeling참고 문헌 6인용 수 7

한 줄 요약

이 논문은 강력한 일반화 능력을 갖춘 트랜스포머 기반 아키텍처인 신경 실행 엔진(NEE)을 제안한다. NEE는 정렬 및 최단 경로와 같은 알고리즘 하위 서브루틴을 학습하여 훈련 분포를 초월한 입력, 특히 더 긴 시퀀스와 알려지지 않은 데이터 범위에서도 거의 완벽한 정확도를 달성한다. 학습된 조건부 마스킹 메커니즘과 이진수 표현 방식을 도입함으로써, NEE는 훈련 시퀀스와는 다를지라도 매우 긴 길이의 입력에 대해서도 높은 성능을 유지한다.

ABSTRACT

A significant effort has been made to train neural networks that replicate algorithmic reasoning, but they often fail to learn the abstract concepts underlying these algorithms. This is evidenced by their inability to generalize to data distributions that are outside of their restricted training sets, namely larger inputs and unseen data. We study these generalization issues at the level of numerical subroutines that comprise common algorithms like sorting, shortest paths, and minimum spanning trees. First, we observe that transformer-based sequence-to-sequence models can learn subroutines like sorting a list of numbers, but their performance rapidly degrades as the length of lists grows beyond those found in the training set. We demonstrate that this is due to attention weights that lose fidelity with longer sequences, particularly when the input numbers are numerically similar. To address the issue, we propose a learned conditional masking mechanism, which enables the model to strongly generalize far outside of its training range with near-perfect accuracy on a variety of algorithms. Second, to generalize to unseen data, we show that encoding numbers with a binary representation leads to embeddings with rich structure once trained on downstream tasks like addition or multiplication. This allows the embedding to handle missing data by faithfully interpolating numbers not seen during training.

연구 동기 및 목표

• 알고리즘 추론에 대한 신경망의 강력한 일반화 부족 문제를 해결하며, 특히 훈련 분포를 초월한 입력 범위에서의 일반화를 향상시키기 위해.
• 수치적으로 유사한 값이 많은 장수열에서 트랜스포머의 어텐션 정밀도 저하 문제를 해결하여 오류 누적이 발생하지 않도록 하기 위해.
• 이진 표현을 통한 구조적이고 내삽 가능한 수치 임베딩을 학습함으로써 알려지지 않은 데이터 분포로의 일반화를 가능하게 하기 위해.
• 단일 훈련된 서브루틴(예: 비교)이 재사용 가능하도록 모듈러한 프레임워크를 설계하여, Dijkstra나 Prim 알고리즘과 같은 다양한 알고리즘에 대해 재훈련 없이도 적용 가능하도록 하기 위해.
• 다양한 그래프 유형과 입력 분포에서의 강건성 확보를 위해, 분포 이동에 대해 성능 저하가 최소화되는 성능을 보여주기 위해.

제안 방법

• 하나의 서브루틴 단계에서 값과 포인터를 동시에 예측할 수 있도록 시퀀스-투-시퀀스 트랜스포머를 사용하여 동적 상태 추적을 가능하게 한다.
• 예측된 포인터에 기반해 인코더 마스크를 업데이트하는 학습된 조건부 마스킹 메커니즘을 도입하여, 장수열에서 어텐션 정밀도를 향상시킨다.
• 수치를 이진수 형태로 표현하여 계층적이고 지수적으로 확장 가능한 임베딩을 생성함으로써, 알려지지 않은 값에 대한 내삽을 지원한다.
• 덧셈과 곱셈과 같은 후행 작업에 대해 모델을 훈련시켜, 체계적이고 해석 가능한 수치 표현을 학습한다.
• 전체 알고리즘 대신 서브루틴에 대해 훈련함으로써 알고리즘 논리의 학습과 데이터 분포를 분리함으로써 모듈러한 재사용을 가능하게 한다.
• 동일한 학습된 서브루틴을 Dijkstra나 Prim 알고리즘 등 다양한 알고리즘에 동일하게 적용하여 재훈련 없이도 전이 가능함을 보여주었다.

실험 결과

연구 질문

• RQ1 트랜스포머 기반 모델이 훈련 분포를 초월한 입력에 대해 정렬 및 최단 경로와 같은 서브루틴을 강력한 일반화 능력으로 실행할 수 있는가?
• RQ2 표준 트랜스포머가 더 긴 시퀀스나 수치적으로 유사한 입력에서 왜 일반화에 실패하는가? 그리고 어떻게 어텐션 정밀도를 유지할 수 있는가?
• RQ3 이진 수치 표현이 체계적이고 내삽 가능한 임베딩을 만들어내어 알려지지 않은 수치로의 일반화를 가능하게 할 수 있는가?
• RQ4 한 번 학습된 서브루틴(예: 비교)이 재훈련 없이 여러 알고리즘에 얼마나 널리 적용 가능한가?
• RQ5 입력 데이터의 분포 이동(예: 다양한 그래프 유형 또는 입력 범위)에 대해 신경 실행 엔진은 얼마나 강건한가?

주요 결과

• Dijkstra의 최단 경로 알고리즘과 Prim의 최소 스패닝 트리 알고리즘에 대해, 훈련 시 8개 노드의 그래프에서 학습했음에도 테스트 그래프 크기가 100까지 확장되도 100.00%의 정확도를 기록했다.
• 훈련 그래프와 테스트 그래프의 유형이 다를 경우에도 그래프 알고리즘 성능은 거의 완벽하게 유지되었으며(최단 경로 99.91%, MST 93.00%), 분포 이동에 대한 강건성을 입증했다.
• 이진 수치 표현을 통해 모델은 체계적이고 해석 가능한 임베딩을 학습하였으며, 이는 내삽을 가능하게 한다. 예를 들어, 훈련 중에 제외된 65%의 수치가 정확하게 내삽되었다.
• 학습된 조건부 마스킹 메커니즋试로, 길이 8인 짧은 시퀀스에서 훈련된 모델이 길이 100까지의 시퀀스로도 정확하게 일반화됨을 입증했다.
• 한 서브루틴(예: 비교)에 대해 훈련된 NEE는 재훈련 없이도 Dijkstra, Prim 등 다양한 알고리즘에 직접 적용 가능하며, 이는 높은 모듈성과 전이 가능성의 증거이다.
• 이진 임베딩을 사용한 모델는 원-핫 출력 대비 파라미터 효율성과 확장성에서 뛰어나며, 모든 작업과 입력 범위에서 높은 성능 유지를 유지한다.