Stitchable Neural Networks

카테고리 없음 2024. 11. 27. 19:35 Posted by 랏츠베리

안녕하세요 KDST 박민철입니다.

 

이번 공유 내용에서는 지난 업로드와 마찬가지로 생소한 주제를 준비해보았습니다. 제가 준비한 내용은 직접 범주화하기에 까다롭지만, Neural Architecture Search (NAS)처럼 보일 수 있는, model stitching에 관한 한 가지 연구, "Stitchable Neural Networks, CVPR 2023 (Spotlight)"를 소개드리고자 합니다. Model stitching은 최근 공개 저장소에 딥러닝 모델들의 급격한 배포가 이루어짐에 따라 (HuggingFace ~81k models, timm ~800 models), 수많은 사전 학습된 모델의 중간 중간 피쳐를 잘 연결하면, 기존 모델보다 더 비용적으로 효율적이고, 우수한 성능을 만들 수 있지 않을까?하는 아이디어에서 출발합니다. 아래의 그림을 통해 이해해봅시다. 우리가 잘 알고 있는 Network compression (pruning, quantization, knowledge distillation)의 기법은 단일 모델에 대해서 sub-optimal 구조를 찾고자 합니다 (one-to-one 관계). 경량화에서 NAS는 단일 모델에서 여러 가능한 sub-network를 찾는 문제로 볼 수 있습니다 (one-to-many). 이 방법들은 결국 하나의 large 모델에서 출발해야 한다는 scalability의 문제가 있습니다.

따라서, Stitching 기법은 우수한 성능을 가진 다수의 배포된 모델을 바탕으로, 레이어별 요소들을 적절하게 연결하여 배포해야할 여러 하드웨어 플랫폼에 적합하게 비용 효율적이고 우수한 성능의 새로운 모델을 디자인하는 것을 목적으로 합니다 (many-to-many).

 

Stitchable Neural Networks (SN-Net)는 사전 학습된 같은 계열의 모델을 연결합니다. 이는, 같은 task를 통해 사전학습된 같은 계열의 딥러닝 모델의 구조 (예, DeiT-Ti/S/B)는 서로 stitching될 수 있다는 과거 연구 결과를 바탕으로 합니다. 이를 위해, SN-Net은 같은 계열의 모델 집합에서 잘 동작하는 모델 한 가지를 anchor로 도입하여, 다른 모델의 특정 레이어로 activation을 전파하기 위해 간단한 stitching layer를 도입합니다. 이 방법을 통해, SN-Net은 자연스럽게 anchor에서 다른 모델로 accuracy-efficiency의 다이나믹한 trade-off를 활용하며 interpolation할 수 있고, 이러한 특성을 바탕으로 하나의 neural network 설계를 배포할 하드웨어의 조건에 맞도록 컨디셔닝할 수 있습니다.

 

SN-Net의 동작은 간단합니다. 예를 들어 아래와 같이 anchor 모델이 L개의 레이어로 구성되어 있다고 합시다.

간단한 아이디어는 아래와 같이 l번째 레이어에서 anchor 모델의 동작을 끊어 stitching layer를 통과하여 다른 모델의 l+1번째 레이어부터 마지막 레이어까지 연결하는 것입니다.

Notation의 혼란을 피하기 위해 m은 l+1이라고 생각하시면 편합니다.

저자들은 위의 방법대로 stitching할 때, 서로 다른 구조의 딥러닝 모델, 예를 들면, ViTs와 CNNs,에 대해서도 적은 성능 저하를 통해 SN-Net 구현이 가능하다고 말하고 있습니다.

저자들은 먼저, 사전 학습된 DeiT (Vision Transformer) 계열 모델을 바탕으로 SN-Net을 구축하는 방법을 소개합니다.

Stitching은 한 개의 anchor 모델이 존재할 때, 또 다른 같은 계열의 모델에 대하여 최적의 위치에서 연결하는 것을 목표로 합니다. 따라서, 2개의 학습된 모델은 같은 도메인에서 사전 학습되어 있어야 하며, 본 연구는 ImageNet과 COCO에 대하여 학습된 모델을 대상으로 하였습니다. Stitching layer는 1x1 컨볼루션 레이어를 활용하였고, least squares 문제를 통해 얻은 해를 바탕으로 layer initialization을 처리하였습니다.

본 논문에서 도입하는 stitching 전략은 크게 두 가지 관점에서 출발하는 데, 하나는 Fast-to-Slow, 다른 하나는 Slow-to-Fast 관점이며, Fast-to-Slow의 경우 작은 모델의 특정 레이어에서 큰 모델의 다음 레이어로 연결하는 방법 그리고 Slow-to-Fast는 이와 반대입니다. Fast-to-Slow가 경험적으로 더 좋은 성능을 보였으며, Anchor와 모델의 구조가 다른 경우, 실험적으로 같은 경우보다 최적점을 찾기 어려운 현상을 언급합니. Stitching 방법은 2가지를 고민할 수 있는데, paired stitching과 unpaired stitching으로 살펴볼 수 있습니다. 설명은 아래의 그림과 같습니다.

 

마지막으로, stitching space는 configuration하기에 따라 다릅니다만, 2개의 모델 사이에는 한개의 최적점만 고려하고, 모델을 여러개 도입하는 것이 가능한데, 이 것만 하더라도 layer와 block 단위로 initial stitching layer를 구성할 때 상당히 많은 stitching point가 존재하게 됩니다. 이를 모두 아울러, 최적의 point를 찾는 문제는 복잡해 보이지만, 서치 공간이 NAS에 비해 매우 작기 때문에 서치에 대한 상대적 효율성을 강조합니다.

아래는 Stitchable Neural Network의  구축 알고리즘입니다.

성능 향상을 위해 추가적으로 Knowledge Distillation을 활용했으며, teacher 모델로 ResNetY-160을 사용했습니다.

한편, 학습을 위해서 ImageNet 기준 50 에폭만 사용하였습니다.

 

Experiment results

먼저, Swin-Ti/S/B를 사용하여 ImageNet-22K에 대한 SN-Net의 성능은 아래와 같습니다. 본 논문은 Top-1의 gap이 Swin-S와 Swin-B 사이에서 83.1%, 83.5%로 minor한 반면, FLOPs 차이는 8.7G와 15.4G로 상당한 데 반해 두 모델을 조합하는 최적화된 SN-Net이 성능과 FLOPs 측면에서 효율성을 가질 수 있음을 보입니다.

 

CNNs들 간의 stitching 및 CNN-ViT 사이의 stitching에 대한 결과를 아래와 같이 보여줍니다. 해당 실험을 위해 저자들은 30 에폭만 사용했습니다.

 

다음은 DeiT-Ti/S/B에 대하여, stitching layer의 initialization에 따른 성능 변화를 보여줍니다. He initialization 대비 least square가 효과적임을 보입니다.

 

다음은 Fast-to-Slow, Slow-to-Fast에 따른 sensitivity 및 anchor 모델에 대해 연결할 모델의 크기에 따른 sensitivity, training 방법에 따른 sensitivity를 보입니다.

 

본 논문은 배포된 모델들을 얼마나 유용하게 조합하여 사용할 수 있을지에 대한 현실적이고, 상상할만한 고민을 풀어보기 위해 stitching이라는 생소한 주제로 접근하여 효과성을 입증한 연구라고 생각합니다.

추가적으로 궁금한 부분은 논문을 참조 부탁드리며, 질문을 댓글로 남겨주시면 함께 고민할 좋은 기회가 될 것으로 생각합니다.

 

감사합니다.