Federated Learning & Continual Learning in Multi-modality Edge AI Research Overview
본 포스트는 이화여자대학교 2024 캡스톤과디자인창업프로젝트 스타트 프로젝트를 위해 작성되었습니다.
Overview
여기 미노타우르스를 죽인 영웅 테세우스가 그의 위대한 여정에서 탔던 배가 있습니다.
후세의 사람들은 이 배를 오랫동안 보존했는데, 문제가 생겼습니다. 배의 판자가 썩기 시작한거죠.
이에 사람들은 썩은 판자를 떼고 새 판자로 교체했습니다. 그러자 어느 시점부터 이 배에는 테세우스가 있었던 원래 배의 조각이 하나도 남지 않고 말았습니다.
이 때 이 교체된 배를 테세우스의 배라고 할 수 있을까요?
딥러닝에서도 이와 비슷한 문제가 발생합니다.
Background
인공지능 모델에 꾸준히 새로운 데이터로 반복하여 학습시키는 것은 모델이 좀 더 general 해지도록 도움을 줍니다.
이를 Continuous learning이라 부릅니다.
하지만, 이렇게 계속해서 새로운 데이터를 투입시키다 보면 먼 미래에 모델은 새로 학습한 task는 잘 수행하지만, 원래의 학습된 weight를 모두 잃어버려 결국 기존의 task를 망각하는 문제가 발생합니다. 이 현상을 Catastrophic Forgetting Problem (재앙적 망각)이라고 부릅니다.
이러한 문제를 해결하는 가장 좋은 방법은 예전에 학습한 데이터셋을 학습 중간중간 다시 투입하는 것입니다. 하지만, 그러기 위해선 데이터셋을 어딘가에 저장해두어야 하며, 한 번 했던 학습을 두세 번 수행해야 하는 문제가 있죠. 만약 기존 데이터셋 투입 시기가 너무 빨라지면 Overfitting(과적합)되는 문제도 발생할 것입니다.
특히, Decentralized Edge AI와 같은 computation cost에 민감한 환경에서 이는 매우 크리티컬 한 문제입니다.
이를 해결하기 위한 다양한 시도가 있었는데요,
그중에서 저희 팀은 서로 다른 Edge 간에 weight를 공유함으로써 catastrophic forgetting problem을 해결하면서도, 더욱 generalize 된 모델을 얻을 수 있는 방법인 Pick-a-back 알고리즘을 개선하는 방법을 모색하고자 합니다.
Main Idea
Pick-a-back 알고리즘에선 서로 다른 Edge AI가 서로 selective knowledge transfer을 수행하여 학습시키는 방법을 선택하고 있습니다.
이 방법은 실제로 Federated Learning 환경에서 좋은 결과를 보였는데요,
저희는 기존 연구를 Multi-modality에 적용시킴과 동시에, 기존 연구에서 model similarity를 구할 때 raw data를 전송하며 발생하는 통신보안 문제를 해결할 방법을 모색하고 있습니다.
Pick-a-Back 예비실험
연구주제의 보다 깊은 이해를 위해 선행연구를 직접 실험해보았습니다.
먼저, 선행연구를 직접 실행하기 위해 팀원들과 함께 사용할 수 있는 GPU 실험서버를 구축했습니다.
학교에서 별도로 지원하는 데이터센터가 없어, 학교로부터 GPU가 부착된 랩탑을 대여하였습니다.
처음 대여했을 때 Window OS가 깔려있어, OS를 날리고 Ubuntu 22.04를 설치하였습니다.
이후 CUDA, cuDNN을 설치하고 포트포워딩을 수행하여 WAN 환경에서도 SSH를 통해 GPU서버에 접속할 수 있도록 만들었습니다.
이제 예비실험을 수행합니다!
Pick-a-back 논문에서 사용한 CIFAR100 Dataset을 다운받고, dataloader에서 쉽게 불러올 수 있게 저장하였습니다.
다음으로 테스트를 자동화하는 스크립트를 별도로 작성해 모든 task, model을 실험한 뒤, 이를 분석하는 Interactive Python notebook을 작성해 성능을 비교하였습니다.
자세한 결과는 아래 팀 깃허브에서 확인하실 수 있습니다.
https://github.com/EWHA-Tespa/Pick-a-back
GitHub - EWHA-Tespa/Pick-a-back: Pre-experiment
Pre-experiment . Contribute to EWHA-Tespa/Pick-a-back development by creating an account on GitHub.
github.com
연구설계
Data Preprocessing
Pick-a-back을 multi-modal 환경에 적용하기 위해, 현재로선 몇 가지 조건이 선행됩니다.
- 모델의 구조가 동일할 것
- 서로 다른 multi-modal model간의 task가 유사할 것
- Model이 너무 무겁지 않을 것 (parameter등 computation cost가 적을 것)
이를 위해 저희는 아래와 같이 실험을 설계했습니다.
1. multi-modal 간에 같은 class를 가진 dataset을 구한다.
2. 동일한 구조의 서로 다른 모델을 single-modality로 학습시킨다.
3. 이들 간의 selective knowledge transfer을 수행하며 성능변화를 확인한다.
이를 수행하기 위해 찾은 적절한 dataset은 N24 Dataset인데요, 이 데이터는 [(이미지,텍스트), 클래스]으로 이루어져있으며, 총 61개의 뉴스 토픽으로 된 class를 갖고 있습니다.
실험과정을 단순화하기 위해 class를 12개의 superclass로 나눠 각 superclass당 5개의 class를 배정하였습니다.
이후, 데이터셋을 모달리티별로 쪼개( [image, class] [text, class] ) 두 개의 trainset을 만들어 저장했습니다.
Model Architecture
저희가 선택한 모델은 Perceiver 입니다.
Perceiver는 Google Deepmind에서 발표한 multi-modality를 학습할 수 있는 유연한 모델로, 입력값을 latent vector로 표현하는 autoencoder의 성격을 가짐과 동시에 Self-attention layer을 활용하여 기존 autoencoder보다 개선된 computation cost, 성능을 보입니다. (자세한 것은 HuggingFace의 Perceiver IO 블로그 참고)
우리 dataset에 맞게 Perceiver을 구현하였습니다. 코드는 아래 깃허브에서 볼 수 있습니다.
https://github.com/EWHA-Tespa/Perceiver
GitHub - EWHA-Tespa/Perceiver
Contribute to EWHA-Tespa/Perceiver development by creating an account on GitHub.
github.com
향후 계획
이제 이 모델을 MNIST, Spam 데이터셋에 입력하여 성능을 검증한 후, 모델이 정상적으로 작동한다면 준비된 N24 Dataset을 넣어 학습시키고 각각 모델에 selective knowledge transfer을 수행할 예정입니다.