KDST

이번에는 NIPS 2020 Poster session에 발표된 논문인 AdaShare: Learning What To Share For Efficient Deep Multi-Task Learning 을 리뷰하려고 합니다. 논문은 링크를 참조해주세요

Background and Introduction

우선 Mutli-task learning이라는 게 어떤 것일까요?
Multi-task learning은 task가 여러 개라는 것인데 큰 dataset $D$ 아래에 다른 dataset $D_{i}$로 구분된 $D=\{D_1, ... , D_i \}$ 를 사용하는 learning method를 말합니다.

Multi-task learning의 종류에는 크게 두 가지 기준인 feature와 method으로 분류합니다.
feature로 구분하는 경우
- Homogeneous-feature MTL: 다른 dataset에서 동일 feature space를 갖는 것(동물만 분류)
- Heterogeneous-feature MTL: 다른 dataset에서 다른 feature space를 갖는 것(동물, 기계 분류)

Method로 구분하는 경우
- Homogeneous MTL: 다른 dataset이지만 동일 method를 다루는 경우(only classification, only regression)
- Heterogeneous MTL: 다른 dataset에서 다른 method를 다루는 경우(Depth Prediction과 Semantic Segmentation을 다루는 경우)

보통의 Multi-task learning(MTL)에서는 Homogeneous MTL을 MTL이라고 부르지만, 이 논문에서는 Heterogeneous MTL을 다루고 있습니다.

다른 Learning Method와 비교할 때 MTL과 다른 점은 다음과 같습니다.

Transfer Learning은 source task를 통해 학습한 후, Knowledge transfer를 통해 target task를 학습시키지만, MTL은 source task와 target task가 동일하게 knowledge를 공유하여 성능을 높이는데 그 목적이 있습니다.

Multi-label learning과 같은 경우, 노트북을 관찰할 때, 키보드와 화면이라는 label을 동시에 분류해내는 것을 의미하지만 MTL은 task에 맞게 노트북만을 혹은 키보드만을 분류해냅니다.

Multi-view learning과 같은 경우는 각각의 특징 앞서 말한 키보드와 화면이라는 관점을 가지고 이를 조합하여 data로 인지하여 노트북으로 분류해내는 것을 말하지만, MTL은 노트북을 보면 각각의 관점을 조합하는 것이 아니라 그냥 단순히 노트북 하나만을 분류해내는 것을 말합니다.

Single-task learning과 비교할 때, 장점은 dataset이 늘고 related data를 사용하는 경우, feature extraction이 잘되기 때문에 training time을 줄일 수도 있고, overfitting을 피하기 용이하여 일반화 성능이 높게 나타납니다.

하지만, Multi-task learning에서도 아직 해결해야하는 난제가 있는데요. 바로 Negative Transfer 현상입니다. 사람이 보기에는 related data일지라도 neural network입장에서는 아닌 것입니다. 그래서 각 task가 서로의 성능향상에 악영향을 끼치는 경우도 상당수 존재합니다. 이러한 점들을 해결하기 위해서 knowledge sharing을 적재적소에 하는 것이 중요한데, 대부분의 논문들은 다음의 3가지 관점에서 이를 관찰합니다.

-When: 학습 이전, 몇 iteration, 몇 epoch부터 sharing할지
-What: parameter(weight, bias), feature, and instance
-How: 어떻게 (selection of what to share, related feature analysis)

이 논문에서는 어떤 parameter를 어디서 share할지에 관한 논문입니다.

Parameter Sharing Method
 Sharing Method는 크게 weight를 모두 공유하는 hard-parameter sharing과 feature나 weight의 정보를 공유(그대로 복사x)하는 soft-parameter sharing이 존재합니다.
- Hard-parameter sharing: Shared feature network + task-specific network로 구성이 되며, task가 서로 많이 다를 경우, shared feature network의 깊이가 얕아져서 다른 task-specific하게 동작하게 합니다. Negative Transfer 현상이 빈번히 발생하며, 해당 network는 그 task끼리의 조합에만 유효한 경우가 대부분입니다. (Adhoc network)

- Soft-parameter sharing: 주로 feature sharing을 하며, task 별로 network를 각각 가지기 때문에 메모리를 많이 먹고, 적절한 network를 구성하였을 때, task별로 capacity가 달라 overfitting이 특정 task에만 일어나는 단점이 존재하기도 한다. 다음 사진은 CVPR에 2016년에 게재된 Cross Stitch Network로 activation map을 task-specific network에서 공유하여 학습을 진행하는 방식이다.

이를 넘어서 parameter sharing하는 node를 선택하여 학습을 진행하는 논문이 2019년 ICCV에 게재된 Deep Elastic Network입니다. Deep Elastic Network에서는 별도의 task-specific한 노드를 고르는 selector network를 RL로서 가장 좋은 성능을 내는 node를 골라 동작할 수 있도록 하였습니다. Model distribution에서 network를 골라 행동하게 하는 것을 의미합니다. Pipeline은 다음 사진과 같습니다.

자, 이런 Network search하는 과정을 Neural Architecture Search(NAS)라고 부릅니다. 이는 AutoML 즉 자동으로 적합하게 학습하는 방법론중의 하나입니다. NAS를 하는 방법에는 두 가지가 존재하는데 첫번째는 Reinforcement Learning을 통해서 구조를 search하는 방식입니다. Action으로 적합한 architecture를 생성하고, network performance를 reward로 부여하여 finetuning을 진행합니다. 두번째로는 Evolutionary Computation을 활용합니다. Genetic Algorithm을 통해서 network들의 집합을 골라 거기서 진화하는 방식(layer생성이나 hyperparameter tuning)을 택하고, skip connection등을 생성하는 것이 있습니다. 이 논문에서는 Gradient-based optimization을 통해 최적의 architecture를 찾을 수 있다고 말을 하고 있습니다. (Genetic Algorithm에서도 optimization을 사용하지만, 이 논문에서는 해당 사항에 대해 언급하고 있지는 않습니다)

기존까지는 Task-specific과 Knowledge sharing하는 part를 manually하게 구성해야 했고, 그랬기에 항상 task 종류에 맞는 network를 주로 제작하여 다른 task 조합에서는 해당 network가 제대로 동작하지 못하는 현상이 있었습니다. 그래서 이 논문에서는 다음과 같이 말하고 있습니다.

Multi-task learning을 위해서 task를 share할지 task-specific layer를 결정할지 자동으로 선택하고 싶고, 이를 통해서 Negative Transfer할 때는 해당 layer를 skip하고, Positive Transfer일 때는 layer를 share하여 성능을 높일 수 있지 않을까라고 제시하고 있습니다.

Proposed Method

논문에서는 Residual Network(ResNet, ResNeXt, DenseNet)등의 shortcut connection의 성공에 주목하여, feature extraction할 때, 특정 task는 skip하여 negative transfer를 막을 수 있지않을까 라는 idea를 주목하였습니다. Shortcut connection여부를 residual block마다 결정하면 task-specific한 block과 knowledge share block을 나눌 수 있다고 생각하였습니다.

해당 block $l$과 task $k$에서 select-or-skip random variable를  $u_{l,k}=[0,1]$라고 하고, policy를 ${U}$라고 할 때, ${n(U)}$가 $2^{N\times K}$로 기하급수적으로 늘어납니다. 그렇기 때문에 Reinforcement Learning으로는 해결하기 어려울 것이라 여겨졌고, 또한 discrete하기에 non-differentiable하여 gradient-based optimization하기 어렵다고 여겨졌습니다. 하지만 binary r.v.인 $u_{l,k}$의 distribution vector $\pi_{l,k}$를 Gumbel Softmax Sampling을 통해서 differentiable하게 만들었습니다.

Gumbel Softmax Sampling은 VAE(Variational AutoEncoder)에서 generative model의 output $x$를 생성할 때의 sampling 또한 미분이 불가능하여 sampling대신에 sampling function을 구현하여 back propagation을 하게한 Reparameterization Trick중 한 방법입니다.
Reparameterization Trick은 Stochastic 한 Node를 Stochastic한 input과 Deterministic한 input으로 구별하여 이 조합을 function화 하여 deterministic한 node로 흐르게 하여 back propagation을 가능하게 하는 것을 말합니다.

Gumbel-Softmax Sampling은 Gumbel-Max Sampling에서 argmax를 통해 sampling하는 것을 softmax로 바꾸어 differentiable하게 만든 것을 의미하며, Gumbel-Max Sampling에서 Gumbel distribution을 $\pi_{l,k}$ term에 더해서 Sampling하게되면 gradient가 흐를 수 있게됩니다. 이를 통해서 를 결정하는 logit $a_{l,k}$를 update할 수 있게되는 것입니다.

이러한 update를 위해 사용하는 Loss term은 다음과 같습니다.

$l_k$는 각 task가 갖는 고유의 loss term입니다. (classification의 경우, cross entropy)

Sparsity Regularization term인 $L_{sparsity}$는 logit $a_{l,k}$를 minimize하여 model을 compact하게 해주는 효과를 일으킵니다.

Sharing loss term은 $a_{l,k}$에서 같은 $l$간의 격차를 줄여 모두 task-specific한 block만 남지 않게함이고, 맨 뒤 block부터 가중치를 낮게 부여하여 앞으로 갈수록 높게 부여하는 형식으로 앞 block은 최대한 sharing을 하고, 뒤 block에서 task-specific하게 구성되도록 하였습니다.

이러한 Model의 전체적인 pipeline은 다음과 같습니다.

1.     맨처음부터 distribution parameter($\pi_{l,k}$)를 학습시키면 학습이 잘 되지 않으므로 curriculum learning을 도입하여 hard-parameter sharing을 통해서 network를 먼저 학습 시킵니다.

2.     이후 매 $l$ epoch에서 끝에서부터 $l$ block까지의 distribution parameter를 학습시킵니다. (Curriculum Learning)

3.     모든 policy distribution parameter가 train되면, select-or-skip decision을 하여 finetuning합니다.

Experiments

기존까지의 state-of-art method와 basline과 비교한 실험입니다. Single-task는 동일 network에 한개의 task만 learning을 각각 진행하여 표에 기재하였고, multi-task의 경우 마지막 fully-connected layer만 분리시켜서 사용하였습니다. ResNet18을 사용하여 NYU를 학습한 결과이며, Segmentation과 Surface Normal Prediction의 두 가지 task에서 사용한 결과입니다.

현재 sota인 MTAN method와 비교할 때 엎치락 뒤치락하는 성능을 보여줍니다.

위 사진은 Task별 logit $a_{l,k}$과 graph를 통해서 sampling하였을 때 실제 select-of-skip된 connection들입니다. 대체로 logit이 동일 image를 사용하기에 낮지는 않지만, surface normal prediction과 segmentation은 물체 구분과 표면각을 구분하는 다른 task여서 비교적 낮게 나오는 것을 뒤에 나올 실험을 통해 확인할 수 있습니다.

다른 dataset인 Tiny-Taskomony에서 ResNet34를 사용하여 실험한 내용입니다. 다음은 dataset과 ground truth의 예시입니다. 각 task의 서로 다른 feature 추출을 통해서 대략적인 task가 related가 되어있는지 여부를 추측해볼 수 있습니다.

5가지 task인 Segmentation, Surface Normal Prediction, Depth Prediction, Key-point Detection, Edge Detection에 대해서 비교 실험을 하였고 결과는 다음과 같습니다.

아래의 policy visualization을 통해서 depth prediction이 다른 task와 달리 상대적으로 unrelated되어있음을 추측해볼 수 있고, 위의 사진결과를 봐도 depth prediction 성능이 다른 task와 비교할 때, 상대적으로 낮게 나오는 것을 확인할 수 있습니다.

다른 실험은 논문에서 제시한 policy와 strategy가 있을 때와 없을 때 사이의 성능 차를 비교한 실험입니다.
1. Stochastic Depth(ECCV,2016) 방식과의 비교 (선택적 layer drop방식이므로 이와 비교)
2. Random Policy (AdaShare논문에서 제시한 Sampling이 합리적인지 비교)
3. 각각의 sparsity와 sharing loss term의 중요성 검증 실험
4. NAS기반 MTL의 sota논문인 Deep Elastic Network(ICCV,2019)방식을 도입하여 policy를 결정하는 것과의 비교와 RL(REINFORCE algorithm)을 사용하여 search했을 때와의 비교

Conclusion

MTL보다는 Neural Architecture Search(NAS) 논문이지만, negative transfer 현상을 회피하고 positive transfer를 보강하여 MTL에 있어서 network의 parameter 수를 유지하면서 성능을 끌어올리고자 노력한 논문입니다. RL을 적용할 경우, search space가 커지는 현상을 Gumbel Softmax Sampling을 통해서 optimization problem으로 변환하였고, 이를 통해 computational resource를 절약할 수 있는 방법을 제시한 연구라고 할 수 있겠습니다.

자세한 내용은 full paper를 참고하거나 댓글로 남겨주세요.

감사합니다.

Minsoo Kang