[논문 리뷰] Discrete Key-Value Bottleneck

Frederik Träuble, Anirudh Goyal외 5명|arXiv (Cornell University)|2022. 07. 22.

Domain Adaptation and Few-Shot Learning인용 수 5

한 줄 요약

이 논문은 사전 훈련된 인코더를 동결하고 추론 중에만 희소한 학습 가능한 값 코드 집합을 선택적으로 업데이트함으로써 지속적 학습을 가능하게 하는 이산 키-밸류 버티브 메커니즘을 제안한다. 작업별 지식을 저장하고 검색하기 위해 이산 키-밸류 쌍을 사용함으로써, 작업 경계가 없는 클래스 증분 학습 벤치마크에서 치명적인 잊음 현상을 줄이며, 다양한 사전 훈련된 모델에서 기존 방법들을 능가한다.

ABSTRACT

Deep neural networks perform well on classification tasks where data streams are i.i.d. and labeled data is abundant. Challenges emerge with non-stationary training data streams such as continual learning. One powerful approach that has addressed this challenge involves pre-training of large encoders on volumes of readily available data, followed by task-specific tuning. Given a new task, however, updating the weights of these encoders is challenging as a large number of weights needs to be fine-tuned, and as a result, they forget information about the previous tasks. In the present work, we propose a model architecture to address this issue, building upon a discrete bottleneck containing pairs of separate and learnable key-value codes. Our paradigm will be to encode; process the representation via a discrete bottleneck; and decode. Here, the input is fed to the pre-trained encoder, the output of the encoder is used to select the nearest keys, and the corresponding values are fed to the decoder to solve the current task. The model can only fetch and re-use a sparse number of these key-value pairs during inference, enabling localized and context-dependent model updates. We theoretically investigate the ability of the discrete key-value bottleneck to minimize the effect of learning under distribution shifts and show that it reduces the complexity of the hypothesis class. We empirically verify the proposed method under challenging class-incremental learning scenarios and show that the proposed model - without any task boundaries - reduces catastrophic forgetting across a wide variety of pre-trained models, outperforming relevant baselines on this task.

연구 동기 및 목표

• 시간이 지남에 따라 훈련 데이터 분포가 변화하는 지속적 학습 환경에서 치명적인 잊음 현상을 해결하기 위해.
• 사전 훈련된 인코더 가중치를 미세조정하지 않고도 효율적이고 국소화된 모델 업데이트를 가능하게 하기 위해.
• 이산 버티브를 통해 가설 클래스 복잡도를 감소시키고 공변량 이동 상황에서 일반화 성능을 향상시키기 위해.
• 사전 훈련된 인코더에서의 지식을 재사용하면서도, 희소하고 맥락 의존적인 값 코드 업데이트를 통해 새로운 작업에 적응할 수 있는 모델을 설계하기 위해.
• 작업 경계 주석 없이 다양한 사전 훈련된 모델에서 클래스 증분 학습 환경에서 방법의 실증적 타당성을 검증하기 위해.

제안 방법

• 모델은 인코딩, 이산 키-밸류 버티브를 통한 처리, 디코딩의 세 단계 파이프라인을 사용한다.
• 사전 훈련된 인코더가 연속적 표현을 생성하며, 이를 C개의 헤드로 투영하여 키-밸류 검색을 수행한다.
• 각 헤드는 학습 가능한 키 코드북에서 가장 가까운 키를 검색하고 해당하는 연속적 값 코드를 검색한다.
• 훈련 중에는 오직 값 코드만 업데이트되며, 키와 인코더는 동결되어 국소화된 적응이 가능하다.
• 키-밸류 코드북은 다양한 데이터 분포에 대해 EMA를 통해 초기화되어 특징 다양체의 광범위한 커버리지가 보장된다.
• 키(인코딩)와 값(디코딩) 코드를 분리함으로써 가설 클래스 복잡도를 감소시키고 모델 용량을 제한함으로써 일반화 성능을 향상시킨다.

실험 결과

연구 질문

• RQ1 작업 경계가 없는 클래스 증분 학습 환경에서 이산 키-밸류 버티브가 치명적인 잊음 현상을 줄일 수 있는가?
• RQ2 이산 버티브 아키텍처는 입력 분포 이동 상황에서 일반화에 어떤 영향을 미치는가?
• RQ3 오직 값 코드에만 매개변수 업데이트를 국소화함으로써 이전 작업의 지식을 얼마나 잘 유지할 수 있는가?
• RQ4 비i.i.d. 데이터 스트림 상황에서 다양한 사전 훈련된 인코더에서 모델의 성능은 어떠한가?
• RQ5 키-밸류 쌍의 활용 패턴은 어떠한가? 미사용된 용량은 활용 부족을 의미하는가, 아니면 강건성을 의미하는가?

주요 결과

• 제안된 모델은 클래스 증분 학습 환경에서 테스트된 모든 사전 훈련된 모델(ResNet50, ViT, DINO, SwAV, ConvMixer)에서 치명적인 잊음 현상을 크게 감소시켰다.
• 작업 경계 주석이나 작업별 헤드 재초기화 없이도 관련 베이스라인을 능가하는 성능을 보였다.
• 훈련 중에 70% 이상의 키-밸류 쌍이 최소 한 번 이상 사용되었으며, 28%의 키는 여전히 미사용 상태로 남아 있어 여전히 보존된 용량과 희소한 활용도를 나타냈다.
• 초기화 단계에서 키 코드는 데이터 다양체 전반에 걸쳐 넓게 분포되어 있었으며, 목표 데이터셋에 도달함에 따라 더 뭉쳐지는 경향을 보이며 도메인 특화된 특징에 적응하는 모습을 보였다.
• 가설 클래스 복잡도 감소 덕분에 분포 이동 상황에서 표준 모델보다 더 나은 일반화 경계를 달성했다.
• 값 코드 업데이트가 매우 국소화되어 있었으며, 대부분의 키는 여러 입력에 의해 공유되어 효율적이고 맥락 의존적인 적응이 가능했다.