[AI System] Adaptive Inference

"Adaptive Inference"

[Recap]

· Pruning, Quantization, Knowledge Distillation 모두 최종적으로 하나의 고정된 Model 생성

    → Model이 고정되어 있으므로 Inference 때 발생하는 계산량도 항상 고정(같은 Model을 끝까지 실행)

· 큰 Model : 정확도↑ & 계산량, Delay, 저장 공간↑ / 작은 Model : 계산량, Delay, 저장 공간↓ & 정확도↓

    → Trade-Off : Accuracy ↔ Computation, Delay, Storage

** Pruning, Quantization 같은 방법은 고정된 하나의 Model만 존재 → Inference 도중에 이 Trade-Off를 유연하게 바꾸기 어려움

[Multi-Exit Neural Networks]

▣ Motivation

· 자율 주행 차량 같은 실무 시나리오에서는 환경에 따라 요구 사항이나 제약 조건이 실시간으로 변화

  - 고속 주행 : 추론 연산과 Delay 감소 필요 ∵ 정확도가 낮더라도 안전을 위해 빠르게 반응하는 것이 중요

  - 저속 주행 : Delay가 여유 있으므로, 연산량을 더 많이 투입해서 정확하고 안전한 예측 수행

** Pruning, Quantization 된 고정 Model은 동적 환경에 빠르게 적응 불가

▣ Multi-Exit Neural Network

· 중간 Layer/Block마다 여러 개의 Exit(조기 예측 지점)을 도입하여 적응형 추론 가능

  - 앞쪽 Exit 예측 : 속도↑, 정확도↓ / 뒤쪽 Exit 예측 : 속도↓, 정확도↑

· 기존 Model은 Block 1, 2, 3 전체를 통과한 후 Prediction 3만 생성

· Multi-Exit 구조는 중간에도 Prediction 1, 2 생성 → Full Model을 끝까지 실행하지 않아도 예측 가능

** 중간 Block의 출력은 바로 최종 Prediction 형태가 아닐 수 있으므로 각 Exit에 Auxiliary Network 추가

    → Auxiliary Network : 중간 Feature를 받아 최종 출력 차원과 같은 Prediction을 만들도록 해줌

▣ Training & Inference

<Original Model>

· θ₁ → θ₂ → … → θ₆을 차례로 통과한 후 마지막에서만 Loss(L) 계산

→ 최종 출력 하나(L)를 줄이기 위해 Forward, Backward를 통해 학습

<Multi-Exit Model>

· 중간에 Auxiliary Model(w₁, w₂, …, w₅) 추가하고, 각 Exit마다 Loss(L₁, L₂, …, L₆) 계산

→ 특정 Exit 하나가 아니라 모든 Exit의 평균 Loss(∑L_k/6)를 줄이기 위해 앞쪽 Exit도 어느 정도 좋은 예측을 하도록 학습

> Gradient Flow

                               ▼

· 각 Block의 Gradient는 해당 Block 이후에 있는 Loss들의 영향을 모두 받음 (θ₆ : L₆ / θ₁ : L₁ ~ L₆) 

→ 앞쪽 Layer일수록 더 많은 Loss로부터 Gradient를 받음

→ 각 Block은 자신 이후에 존재하는 Exit들의 성능을 좋게 만드는 방향으로 학습

> Backpropagation Path

· 여러 Exit의 Loss가 동시에 앞쪽 Block들을 학습시키는 구조 (L₆ : θ₆ ~ θ₁ / L₄ : θ₄ ~ θ₁)

> Memory

  - Activation Memory = (|a₁| + |a₁| + … + |a₅|) x n

· 각 Block 사이 Activation을 저장

· Mini-Batch 크기 n에 비례

  - Parameter Memory N = ∑|θ_k| + ∑|w_k| (θ_k : k = 1 ~ 6 / w_k : k = 1 ~ 5)

· SGD : 2N / SGD with Momentum : 3N / Adam : 4N

· 원래 Block Parameter "θ"뿐만 아니라 Auxiliary Network Parameter "w"도 포함 → 기존 Model보다 약간 증가

∴ Multi-Exit Model : Inference Latency/Computation을 줄일 수 있지만, Parameter Storage 자체를 크게 줄이진 못함

> Inference

· 추론 시에는 환경에 따라 어느 Exit에서 예측할 지 선택 가능

    (급하면 앞쪽 Exit에서 빠르게 예측 / 시간이 충분하거나 배터리가 여유로우면 뒤쪽 Exit에서 예측)

→ Accuracy와 Delay/Computation 사이의 Trade-Off를 동적으로 조절 가능

> Performance

<정확도>

· 앞쪽 Exit : 계산량, Delay↓ & 정확도↓ / 뒤쪽 Exit : 계산량, Delay↑ & 정확도↑

> Local-Loss-Based Training과 비교

  - Local-Loss-Based Training

· Training Latency를 줄이기 위한 방법

· 각 Local Loss가 대응되는 일부 Layer만 Update → Backward-Locking 문제↓

  - Multi-Exit Training

· Inference 시 여러 Exit을 활용하기 위한 방법

· 각 Loss가 이전 Layer들까지 Backpropagation되어, 여러 Exit에서 좋은 성능을 내도록 학습

▣ Options To Improve The Performance

> Inference 중 Model Ensemble 사용

· 특정 Latency 또는 Computation 제한이 있을 때, 해당 제한 안에서 도달 가능한 Exit들의 예측을 활용

· 3번째 Exit 하나의 출력만 사용하지 않고, 1~3번째 Exit의 출력들을 평균 내거나 결합하면 더 안정적인 예측 가능

> Training 중 Knowledge Distillation 사용

· 마지막 Exit : 가장 깊은 Model을 사용하므로 일반적으로 가장 신뢰도↑

    → 마지막 Exit의 출력을 Teacher처럼 사용하고, 앞쪽 Exit들이 이를 따라가도록 학습 진행

· 전체 Full Model(= 마지막 Exit까지 간 Model)을 Teacher Model, 앞쪽 Exit에서 끝나는 Sub-Model을 Student Model로 간주

    → 앞쪽 Exit들은 정답 Label뿐만 아니라 Teacher의 Soft Prediction도 학습

<Performance>

· Distillation 적용 시 각 Exit의 정확도 상승(특히 초기 Exit에서 성능 향상↑)

∵ 앞쪽 Exit들이 제한된 계산량만 사용하더라도, 마지막 Exit의 지식을 받아 더 좋은 예측을 하도록 학습

[Extension To Split Learning & Inference]

▣ Motivation

· AI를 사용할 때는 저장 공간, 연산량, Delay가 작아야 함

> 전체 Model을 Client에 올리는 방식

· Server와 통신하지 않아도 되므로 Privacy 측면에서 좋음

· Clinet가 큰 Model을 저장하고 계산해야 해서 배터리가 부족하거나 기기 성능이 낮으면 매우 느려질 수 있음

> 전체 Model을 Server에 올리는 방식

· Client는 데이터를 Server로 보내고, Server가 예측을 수행하여 Client의 부담 감소

· 통신이 필요하기 때문에 네트워크 상태가 나쁘면 느려지고, Privacy 문제 발생

▣ Adaptive Inference

· 작은 앞부분 Model은 Client에, 나머지 큰 Model은 Server에 배치

· Client 쪽에는 작은 Model과 Auxiliary Network 존재

※ Auxiliary Network : Client-Side Model의 중간 출력을 바로 예측 결과로 바꿔주는 작은 Network

→ 긴급하거나 쉬운 입력이면 Client 쪽에서 바로 예측 진행

→ 덜 긴급하거나 어려운 입력이면 중간 Activation을 Server로 보내고, Server의 큰 Model이 이어서 계산한 후 예측 진행

** Exit이 2개인 Multi-Exit Neural Network로 볼 수 있음

> 1st Exit

· Client 쪽에 존재

· 작은 Model만 사용하므로 빠르지만 정확도는 상대적으로 낮음

> 2nd Exit

· Server 쪽에 존재

· 원래 Full Model에 가까운 계산을 사용하므로 정확도는 높지만, 통신과 추가 계산 때문에 Delay 증가

△ Model Splitting을 통해 Client-Side 저장 공간과 연산량 감소 (∵ 전체 Model을 사용자 기기에 올리지 않아도 됨)

△ Adaptive Inference를 통해 Delay 감소 (∵ 급한 상황에서는 Client의 1st Exit에서 바로 예측)

▣ Federated Learning

> Option 1 - 각 Client가 전체 Multi-Exit Neural Network를 학습

· 각 Client는 자신의 Local Dataset으로 Loss 계산 (Loss : 1st Exit Loss와 2nd Exit Loss를 섞은 형태)

· 각 Client가 Model을 학습한 후, Server는 Client들이 보낸 Model들의 평균 계산(FedAvg)

▼ 각 Client가 전체 Multi-Exit Model을 가지고 학습 → Client의 저장 공간과 연산 부담↑

> Option 2 - Model을 실제 배치 방식처럼 나누어 학습

· Client : 앞부분 Model과 1st Exit 학습 / Server : 뒷부분 Model과 2nd Exit 학습

· Forward 단계 : Client가 Activation과 Label을 Server로 전송

    → Server는 Server-Side Model을 통해 2nd Exit Loss 계산

· Backward 단계 : Server가 Gradient 또는 Backpropagation Signal을 Client로 다시 전송

    → Client는 Client-Side Model과 Auxiliary Network를 Update

△ Client가 전체 Model을 들고 있지 않아도 되므로 저장 공간과 연산량 감소

▼ Client가 Server의 Backpropagation Signal을 기다려야 하므로 학습 Delay 문제 발생

▣ Summary

> Training - Federated Learning Option 2

· 두 Exit의 Loss를 모두 사용 → Client-Side Exit과 Server-Side Exit이 모두 좋은 예측을 하도록 학습

· Federated Learning을 사용하면 여러 Client의 데이터를 활용 가능

> Inference - Adaptive Inference

· 상황에 따라 Exit 선택

    → 긴급하면 Client-Side의 1st Exit에서 바로 예측

    → 덜 긴급 또는 높은 정확도가 필요하면 Activation을 Server로 보내고 Server-Side의 2nd Exit에서 예측

→ 정확도와 Latency 사이의 Trade-Off를 유연하게 조절하면서, 동시에 Client의 저장 공간과 연산량 감소

▣ Enhance Performance

> Knowledge Distillation

· Server-Side의 2nd Exit은 더 큰 Model을 사용하므로 일반적으로 더 정확

→ 2nd Exit을 Teacher처럼 사용하고, Client-Side의 1st Exit을 Student처럼 학습

> Ensemble

· 추론할 때 1st Exit과 2nd Exit의 출력을 함께 사용해 최종 예측 생성

→ 단일 Exit만 사용하는 것보다 더 안정적인 예측 가능

[Reference]

· Adaptive Inference (AI System : Week 13) - Dongjun Han