ArXiv	https://arxiv.org/abs/2501.02336
Github Code	https://github.com/ASISys/AdaSkip
Authors	Zhuomin He1*†, Yizhen Yao1*†, Pengfei Zuo2*, Bin Gao3†, Qinya Li1‡, Zhenzhe Zheng1, Fan Wu1
Affiliation	1Shanghai Jiao Tong University 2Huawei Cloud 3National University of Singapore

💡

 Key Differentiator

 Sublayer-wise skipping

기존은 layer 전체 단위로 스킵했지만,

Attention과 FFN을 분리해서 선택적으로 건너뜀

 Auto-adaptive

기존에는 Decoding단계에만 스킵을 했지만,

IO similarity 기반으로  prefilling(offline) +  decoding(online) 단계를 다르게 다룸

 Applicable to both prefilling & decoding

Prefilling(처음 입력 처리 단계)에서도 스킵 전략을 쓰고,

Decoding(토큰별 생성 단계)에서도 online learning으로 동적 보정

🤷

Why I chose this paper?

VLM Agent 경량화해서 모바일에서 잘 돌아가게 하는 연구 진행중

Early Exit 방식을 통해 더 정확도가 좋아지는 경우가 있어서 아이디어 얻기 위해서

어떻게 실험했는지 참고하기 위해서

Background and Motivation

IO Similarity and Transformer Module Importance

 IO Similarity? (Input-Output)

Transformer의 한 모듈(한 블록의 Attention 서브레이어 or FFN 서브레이어)에서

입력 벡터와 출력 벡터가 얼마나 비슷한지 코사인 유사도(cosine similarity)로 측정

$$\begin{equation} Similarity(\vec{a}, \vec{b}) = \frac{\vec{a} \cdot \vec{b}}{\|\vec{a}\| \|\vec{b}\|} = \frac{\sum_{i=1}^n a_i b_i}{\sqrt{\sum_{i=1}^n a_i^2} \, \sqrt{\sum_{i=1}^n b_i^2}} \end{equation}$$

IO similarity가 높다 → 입력과 출력이 거의 같음 → 모듈이 큰 변화 없이 데이터를 전달 → 중요도가 낮음

IO similarity가 낮다 → 입력과 출력이 많이 다름 → 모듈이 적극적으로 변환 → 중요도가 높음

 
IO Similarity와 중요도의 관계 검증

1차 실행: 각 레이어의 IO similarity를 측정하여 프로파일링.

2차 실행: 특정 similarity 값을 기준으로 레이어를 선택적으로 스킵(skip)하고 성능(GPT score) 확인.

결과:

 LeastSkip: IO similarity가 낮은 레이어를 스킵 → 성능 급격히 하락 (1개만 스킵해도 점수 < 1.0).

 MostSkip: IO similarity가 높은 레이어를 스킵 → 1, 3, 5개 스킵 시에도 점수 8.9, 6.1, 4.2 유지.

 → IO similarity가 높을수록 스킵해도 성능 손상이 적다.

Existing Layer-wise Skipping Strategies

1.  Early Skipping

   항상 앞쪽 몇 개 레이어를 고정적으로 스킵.

2. 장점: 배치(batch) 연산 호환성 좋음.
3. 단점: 앞부분이 중요한 경우 성능 손실 가능.

2.  Periodic Skipping

   일정 간격으로 중간 레이어를 스킵 (예: 4개마다 1개 스킵).

3. 장점: 배치 연산 가능.
4. 단점: 레이어 중요도의 변동성을 반영하지 못함.

3.  Early Exit

   각 레이어 계산 후 조건(예: confidence)이 충족되면 뒤 계산 생략.

4. 장점: 불필요한 연산 절약.
5. 단점: 중요한 뒷부분 레이어를 건너뛸 위험, classifier 학습이나 모델 파인튜닝 필요.

Motivation

기존 레이어 스키핑 전략들이 장문 맥락 추론에서 한계가 있는 이유

 Observation 1:
The layer importance distribution exhibits significant variation across diverse models.
모델에 따라 레이어 중요도 분포가 크게 다르다.

기존 layer-wise skipping 기법은 “항상 같은 위치의 레이어”를 건너뛰는데,

이렇게 하면 모델별 중요도 패턴 차이를 무시하게 돼서  적응성이 떨어짐.

따라서  모델별로 맞춤형(Adaptive) 스킵 전략이 필요.

 Observation 2:
The importance distributions of attention and FFN modules are different.
어텐션(Attention)과 FFN(Feed-Forward Network) 모듈의 중요도 분포가 서로 다르다.

Transformer 레이어는 보통  Attention sublayer와  FFN sublayer로 구성됨.

Attention sublayer와 FFN sublayer가 뭐야?

 Attention sublayer:

문장 속 단어들이 서로를 얼마나 참고해야 하는지 계산하는 부분

→ "이 단어가 저 단어와 얼마나 관련 있는지" 점수를 매기고, 중요한 정보에 더 집중하도록

 FFN (Feed-Forward Network) sublayer:

Attention에서 모아온 정보를 각 단어별로 따로  가공하는 작은 신경망

→ "집중해서 모은 정보"를 더 복잡하게 변형하고, 다음 단계로 넘겨줌

이 둘의 IO similarity를 따로 분석해 보니 분포가 다르게 나옴.

 Attention: 최고 IO similarity가 거의  0.97으로 높고, 값이 서로 비슷하게 모여 있음.

 FFN: 최고 IO similarity가 0.95 정도이고, 값이 퍼져 있음.

→ Attention은 FFN보다 스킵할 후보가 많음

→ 기존 방법처럼 레이어 전체를 한 번에 스킵하는 건  비효율적이고,

 Observation 3:
The importance distribution of sublayers in the prefilling and decoding phases have similar trends but different fluctuation degrees.
프리필링과 디코딩 단계에서 서브레이어 중요도 분포의 경향이 유사하지만, 변동 정도는 다르다.

Prefilling: 처음 문맥을 읽어들이는 단계.

Decoding: 토큰을 하나씩 생성하는 단계.

두 단계에서 Attention과 FFN의 IO similarity 변화를 비교:

 전체적인 추세는 비슷 → 두 단계에서 비슷한 스킵 전략을 공유할 수 있음.

하지만 FFN sublayer는 IO similarity가 Decoding 단계에서 Prefilling 때보다  높음

→ Decoding 단계에서  FFN sublayer를 더 많이 스킵해도 성능 영향이 적을 가능성이 큼.

---

Methodology

Sublayer Skipping during Prefilling with Offline Importance Learning

 왜 Prefilling 단계에서 스킵이 중요한가?

Prefilling은 긴 입력을 처음 읽는 단계라서:

TTFT (Time To First Token)가 길어짐

 KV 캐시 사용량이 많음

모델마다 IO similarity 분포가 달라서,  고정된 레이어 스킵은 최적이 아님.

문제: Prefilling 시작 전에는  중요도에 대한 사전 정보가 없음 → adaptive skipping이 어려움.

 Insight

hit rate = unimportant sublayer(= skip target)를 얼마나 잘 맞췄는지 정확도

skip한 sublayer 개수 : 4, 6, 10

그 중에서 얼마나 일치하게 선택되었는지(hit)를 평균 : 3.76, 4.86, 9.31

여러 데이터셋과 LLaMA3.1-8B-128k 모델로 Prefilling IO similarity를 측정

→ 한 데이터셋에서 측정한 평균 IO similarity로 다른 데이터셋에서 스킵 대상을 예측해도  Hit Rate가 높음

Prefilling에서  과거 IO similarity → 현재 스킵 대상 예측이 가능하고, 데이터셋 간 공유도 가능.

 Offline Importance Learning Workflow

1.  데이터 수집
2. 여러 개의 inference task(프롬프트)를 준비
3. 각 task는 길이가 다를 수 있고, 모델에는 M개의 transformer 레이어(각각 attention + FFN sublayer)가 있음
4. 각 sublayer(예: attention 1, FFN 1, attention 2, FFN 2, …)별로 입력 벡터 vs 출력 벡터의 유사도(IO similarity)를 측정
5. 모든 task의 토큰마다 값을 모아  평균 IO similarity를 구함 → 이게 “얼마나 덜 중요한지”를 알려줌

   $$\begin{equation}\overline{Similarity}_{j} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{|T_i|} Similarity(\vec{a}^{\,j}_{it}, \vec{b}^{\,j}_{it})}{\sum_{i=1}^{N} |T_i|}\end{equation}$$

   a_{jit} = 입력 벡터, b_{jit} = 출력 벡터.

2.  Deviation 보정 (Scale factor)
3. 단순 similarity만 쓰면 입력과 출력 벡터 크기 차이 때문에 약간의 오차가 생김
4. 그래서 각 sublayer에 대해  입력과 출력 벡터 크기 비율을 평균 내서 보정 계수(scale factor)를 계산

   $$\begin{equation}\overline{Scale}_{j} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{|T_i|} \frac{\lVert \vec{b}^{\,j}_{it} \rVert}{\lVert \vec{a}^{\,j}_{it} \rVert}}{\sum_{i=1}^{N} |T_i|}\end{equation}$$

   보정된 출력 벡터:

   $$\begin{equation}\hat{\vec{b}}^{\,j}_{it} = \overline{Scale}_{j} \cdot \vec{a}^{\,j}_{it}\end{equation}$$

3.  Sublayer 중요도 정렬
4. 모든 sublayer(attention+FFN, 총 2M개)에 대해 평균 similarity를 구했으니
5. Similarity‾ 값이 높은 순서(= 덜 중요한 순서)로 정렬

4.  가속 비율(Acceleration ratio) α로 스킵량 조정
5. 전체 속도 ↔ 성능 trade-off를 제어하기 위해  acceleration ratio α를 설정.
6. α가 크면 더 많이 스킵, 작으면 덜 스킵.

$$\begin{equation}m = M - \frac{M}{\alpha}\end{equation}$$

7. 최종적으로 상위(덜 중요한) 2m개의 sublayer를 스킵 대상으로 선택.

 Prefilling 전: 과거 데이터로 미리 sublayer 중요도를 학습.

 Prefilling 중: IO similarity를 새로 측정할 필요 없이, 과거 평균값 + 보정치로 스킵 대상 결정.

---

 Extra FFN Sublayer Skipping during Decoding with Online Importance Learning

Prefilling 단계에서 이미 스킵할 sublayer를 정했지만, Decoding 중에 실시간 데이터(online learning)를 활용해 더 많은 FFN을 건너뛰어 속도를 높이는 전략

 Decoding 단계에서 FFN을 더 스킵할 수 있나?

Observation 3에 따르면:

1. Prefilling과 Decoding의  IO similarity 추세는 비슷함 → Prefilling에서 골라낸 스킵 후보를 Decoding에서도 재사용 가능.

2. 하지만  FFN sublayer는 Decoding에서 IO similarity가 더 높음 → 즉, 덜 중요한 경우가 많으니  더 많이 스킵 가능.

 Insight

Decoding이 시작될 때,  처음 몇 개 토큰의 IO similarity만 측정해도 이후 생성에서 중요한 sublayer를 잘 예측 가능

hit rate = unimportant sublayer(= skip target)를 얼마나 잘 맞췄는지 정확도

초기 window size를 늘릴수록 hit rate가 올라가다 어느 지점 이후 거의 일정.

→ 굳이 모든 토큰을 분석할 필요 없이,  초반 P개의 토큰만 보고 결정해도 충분.

 Online Importance Learning Workflow

1.  Input
2. Prefilling에서 이미 정한 “스킵 후보 집합(skipped)”
3. Decoding 초반  P개 토큰

2.  초기 윈도우로 FFN 중요도 측정
3. 각 FFN 서브레이어 j에 대해, 처음 P개의 디코딩 토큰에 대한  IO similarity 평균을 계산

   $$\begin{equation}\overline{Similarity}^{P}_{j} = \frac{\sum_{t=1}^{P} Similarity(\vec{a}^{\,j}_{t}, \vec{b}^{\,j}_{t})}{P}\end{equation}$$

4. ajt = 입력 벡터, bjt = 출력 벡터

3.  Prefilling 결과와 비교
4. Prefilling에서 스킵된 sublayer들의 similarity 중 최소값을  β로 설정
5. “이 값 이상이면 ‘덜 중요한 축’에 속하므로, 디코딩에서도 추가로 스킵해도 괜찮다”

   $$\begin{equation}\beta = \min \{ \overline{Similarity}_{j} \mid j \in skipped \}\end{equation}$$

6.  index: 모든 FFN sublayer 인덱스 집합.
7.  skipped: Prefilling에서 이미 스킵하던 sublayer 집합.

4.  추가 FFN 스킵 선별
5. 모든 FFN 서브레이어 인덱스 집합을 훑으면서 $\overline{Similarity}^{P}_{j}$ ≥ β 인 FFN sublayer를 찾기

5.  최종 스킵 집합 생성
6. Prefilling에서 스킵한 sublayer + 추가로 찾은 FFN sublayer를 합쳐 최종 스킵 집합 완성하기
7. skipped^(P) = skipped ∪ EXTRA_SKIP

6.  보정(Scale Compensation)
7. Prefilling에서 미리 구해둔 평균 스케일 값 Scale‾j 를 이용해 입력 벡터를 보정

   $$\begin{equation}\hat{\vec{b}}^{\,j}_{it} = \overline{Scale}_{j} \cdot \vec{a}^{\,j}_{it}\end{equation}$$

 추가 가속: Prefilling만 썼을 때보다 더 많은 FFN을 스킵 → 더 큰 속도 향상

 성능 유지: 초기 토큰으로 추정하기 때문에 성능 저하를 최소화

 실시간 적응: 현재 문맥(context)에 맞춰 동적으로 결정

→ Decoding 초기에  한 번 더 업데이트해서, 덜 중요한 FFN을 건너뛰는 방식

Experiments

Setting

 MFieldQA, TREC, TriviaQA → 작은/중간 크기 데이터셋 →  세부 분석용 (ATTN vs FFN 별 hit rate 등)

 GovReport, MultiNews → 긴 문맥 요약 데이터셋 →  End-to-End 가속 실험용 (실제 성능 + 속도)

Hardware: NVIDIA A100 GPUs (80GB)

KV cache enabled

 Prefilling phase에서  acceleration ratio α는 얼마나 많은 sublayer가 skip되는지 보고 결정

 Decoding phase online learning에서 첫 토큰 개수는 중요도 보고 추정함 (P = 20, 50, 100)

Results of Prefilling Tasks, Decoding Tasks

F1 : QA같은거에서 정밀도, 재현율 조화평균

ACC : 분류 문제에서 정확도

Rouge-L : 생성 태스크에서 참조 요약과 모델 출력 간의 longest common subsequence 기반 유사도

SU (Speedup) : 기존 대비 속도 향상비율

실제 Speedup는 왜 AdaSkip이 다른 skip 보다 빠를까??

→  진짜 쓸모없는 sublayer만 골라내 건너뜀 →  더 많은 연산 절약 + 성능 유지

“AdaSkip achieves superior speedup compared with the state-of-the-art skipping strategies, as it adaptively skips the most unimportant sublayers in both prefilling and decoding phases without requiring additional training.”

Results of End-to-End Testing

실제로 전체를 실행해도 AdaSkip이 정확도가 높음

---

Limitation & Future Work

 Prefilling 단계에서 사전 데이터가 필요함

처음 실행할 때는 어떤 레이어를 건너뛸지 알 수 없어서 과거 데이터를 기반으로 예측해야 하는 구조.

→ Prefilling에서 더 정교한  zero-shot layer importance predictor 설계

 모델·데이터셋 특화 적응성 문제

완전히 새로운 모델이나 도메인에 적용할 때는 아직 한계가 있다

→ Dataset 간 transfer를 넘어  cross-model generalization까지 확장

 Not adaptive parameter α (acceleration ratio),  P (online learning window size)

실험자가 α와 P를 고정해서 설정해야 함 → 데이터셋이나 태스크가 바뀌면 최적값이 달라짐

→ 모델이 스스로 α와 P를 상황에 맞게 조정할 수 있는 메커니즘 개발

→ 태스크 난이도나 문맥 길이에 따라 dynamic하게 skip ratio를 조절 → 더 범용적이고 안정적인 성능 확보.

---

Q&A

논문 Presentation 발표 중 제대로 답변 못한 Q&A

GPT Score

강력한 평가자 모델(GPT-4, Claude, Gemini 등)을 불러와서

문제와 후보 답변을 같이 제시

평가 모델이 어느 쪽이 더 낫다고 판단하거나 점수 매기기

→ 사람이 직접 평가하는 것보다 빠르고 일관성 있음

Skip 방법은 이전 결과가 필요할텐데 어떻게 하는거지?

Early Exit은 종료하는거니까 상관 없는데, Skip 방법은 이전 결과가 필요할텐데…

이 논문에서 Scaling 한다는 것을 알고 있었어서, 그게 일반적인건지 이 논문에서 도입한건지 확실하지 않았었는데, 이 논문에서 도입한 거였다

 Identity Mapping (입력 그대로 전달)

스킵된 레이어는 아무 계산도 하지 않고, 입력을 그대로 다음 레이어로 넘김

Transformer에서 output = input + f(input) 에서 f(input) = 0

기존 방법들은 이걸 썼음. (SkipDecode, Unified Layer Skipping)

 Scaling (입력에 보정 계수 곱하기) - AdaSkip

단순히 입력 aj를 출력 대신 쓰면 크기 차이 때문에 deviation(편차)이 생김

그래서 각 sublayer에 대해  입력과 출력 벡터 크기 비율을 평균 내서 보정 계수(scale factor)를 계산

$$\begin{equation}\overline{Scale}_{j} = \frac{\sum_{i=1}^{N} \sum_{t=1}^{|T_i|} \frac{\lVert \vec{b}^{\,j}_{it} \rVert}{\lVert \vec{a}^{\,j}_{it} \rVert}}{\sum_{i=1}^{N} |T_i|}\end{equation}$$

보정된 출력 벡터:

$$\begin{equation}\hat{\vec{b}}^{\,j}_{it} = \overline{Scale}_{j} \cdot \vec{a}^{\,j}_{it}\end{equation}$$

Prefilling은 병렬이 가능해서 Decoding보다 금방 하지 않나?

근데 Prefilling은 병렬이 가능해서 Decoding이 더 오래걸리지 않아? 왜 Prefilling에서 굳이?

기존 기법들은 Decoding에서 skip하는 방식을 채택했음

 Prefilling: 병렬이라 속도 자체는 빠르지만, 계산량과 메모리 소비가 매우 큼

→ TTFT(첫 토큰 지연)가 길어짐

→ 사용자 입장에서는 "첫 응답까지 기다리는 시간"이 길어짐

 Decoding: 토큰이 하나씩 나오면서 체감 속도는 빠르게 보일 수 있음 (streaming)

 Prefill에서 Skip

→ 많은 입력 토큰 × 많은 레이어라면, sublayer 하나만 스킵해도 전체 계산량을 크게 줄일 수 있음

→ 즉, 폭발적인 연산량이 들어가는 지점이라 skip 효율이 큼

Skip Decode, Unified Skipping은 이미 있는건가?

이미 있는거 쓴건가? 아니면 뭔가 그 skip에 맞게 만든건가?
막 정확도 0.0도 있던데 제대로 측정한게 맞아?

early skipping으로 SkipDecode, periodic skipping으로 Unified Skipping 을 baseline으로 사용

 근데, 이 기법들은 원래 Decoding단계용임

 디코딩 태스크는 원래 설계대로, Prefill이랑 end-to-end 태스크는 같은 스킵 규칙을 확장 적용

“Its Vicuna Rouge-L scores for the two summarization tasks fall below 4.0, possibly due to the accumulation of errors in the autoregressive process.”

그리고 Decoding에서도 꽤 차이가 큰데, 오토리그레시브 생성 과정에서의 오차 누적(accumulation of errors) 때문일 수 있다고 언급함

→ Decoding용 방식을 억지로 규칙을 확장해서 오차가 누적되어 성능이 0.0처럼 낮게 나온 것이다!

왜 Offline Learning할때 다른 Dataset으로 해도 비슷해? Observation에는 아니라며

오해했는데 Observation1에서는  Model별로 차이 크다고 했음!

Offline Importance Learning during Prefilling 에서

다른  Dataset으로 측정해도 괜찮다고 했음!

 → Limitation : Model별로 차이가 크니까 결국 모델을 바꾼다면 Offline Learning으로 또 학습해야함