LiteStage: 다단계 추론을 위한 지연시간 인식 레이어 스킵핑

기본 정보

• 논문 ID: 2510.14211
• 제목: LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning
• 저자: Beomseok Kang, Jiwon Song, Jae-Joon Kim (서울대학교)
• 분류: cs.CL, cs.AI
• 발표 시간: 2025년 10월 16일
• 논문 링크: https://arxiv.org/abs/2510.14211
• 코드 링크: https://github.com/beomseokg/LiteStage

초록

다단계 추론은 복잡한 문제를 순차적 부분 단계로 분해하여 소형 언어 모델의 추론 능력을 향상시키는 효과적인 전략이 되었습니다. 그러나 이는 지연시간 증가를 초래합니다. 저자들은 기존의 적응형 가속 기술(예: 레이어 스킵핑)이 이러한 설정에서 효율성과 정확도의 균형을 맞추기 어렵다는 것을 관찰했으며, 주로 두 가지 핵심 과제에 직면했습니다: (1) 단계 간 스킵핑 민감도의 차이, (2) 중복 출력 토큰의 생성. 이러한 문제를 해결하기 위해 본 논문은 다단계 추론을 위한 지연시간 인식 레이어 스킵핑 프레임워크인 LiteStage를 제안합니다. LiteStage는 최적 레이어 예산을 할당하는 단계별 오프라인 검색과 불필요한 디코딩을 억제하기 위한 신뢰도 기반 온라인 생성 조기 종료 메커니즘을 결합합니다. OBQA, CSQA 및 StrategyQA 세 가지 벤치마크에서의 실험은 LiteStage가 최대 1.70배의 가속을 달성하면서 정확도 손실이 4.0% 미만이며, 이전의 훈련 불필요 레이어 스킵핑 방법을 능가함을 보여줍니다.

연구 배경 및 동기

문제 정의

다단계 추론은 복잡한 문제를 여러 개의 연속적인 부분 문제로 분해하여 소형 언어 모델의 추론 능력을 향상시킵니다. 예를 들어, TinyThinker는 세 단계 추론을 채택합니다: 회상(Recall), 분석(Analysis) 및 요약(Summary). 이러한 방법이 추론 품질을 효과적으로 향상시키지만, 추론 지연시간도 불가피하게 증가시킵니다.

핵심 과제

저자들은 심층 분석을 통해 두 가지 주요 문제를 발견했습니다:

1. 단계 간 스킵핑 민감도 차이: 서로 다른 추론 단계는 레이어 스킵핑에 대해 현저히 다른 민감도를 보입니다. 실험에 따르면 Stage 3(요약 단계)이 레이어 스킵핑에 가장 민감하고, Stage 1(회상 단계)은 상대적으로 견고합니다.
2. 중복 토큰 생성: 레이어 스킵핑은 토큰당 계산 비용을 줄이지만, 종종 더 많은 토큰을 생성하도록 하여 오히려 종단 간 지연시간을 증가시킵니다.

기존 방법의 한계

기존 레이어 스킵핑 방법(예: SkipDecode, UnifiedSkip, AdaSkip)은 일반적으로 통일된 스킵핑 전략을 채택하여 다단계 추론의 서로 다른 단계의 특성에 적응할 수 없으며, 다음을 초래합니다:

• 민감한 단계에서 과도한 압축으로 인한 정확도의 급격한 하락
• 레이어 스킵핑으로 인한 생성 길이 증가 문제 무시
• 지연시간 인식 최적화 메커니즘 부재

핵심 기여

1. LiteStage 프레임워크 제안: 다단계 추론을 위해 특별히 설계된 첫 번째 지연시간 인식 레이어 스킵핑 프레임워크로, 단계 간 민감도 차이와 중복 토큰 생성 문제를 효과적으로 해결합니다.
2. 단계별 레이어 예산 할당 전략: 가장 느린 단계에서 가장 빠른 단계로의 탐욕 검색 알고리즘을 설계하여 각 추론 단계에 최적의 레이어 스킵핑 예산을 할당합니다.
3. 신뢰도 기반 생성 조기 종료 메커니즘: 온라인 신뢰도 모니터링을 도입하여 낮은 신뢰도의 중복 생성을 동적으로 종료하고 추론 효율을 추가로 향상시킵니다.
4. 현저한 성능 향상: 세 가지 벤치마크 데이터셋에서 1.161.70배 가속을 달성하면서 정확도 손실은 0.44.0%에 불과하며, 기존 훈련 불필요 방법을 크게 능가합니다.

방법 상세 설명

작업 정의

테스트 데이터셋 D가 주어졌을 때, 목표는 주어진 정확도 임계값 ε 내에서 추론 지연시간을 최소화하는 단계별 레이어 예산 L을 찾는 것입니다:

argmin_L (1/|D|) ∑_{d∈D} T(M_L(d))
subject to: A(M_L(d)) ≤ A(M(d)) - ε

여기서 T와 A는 각각 추론 지연시간과 정확도를 나타내며, M_L과 M은 각각 레이어 스킵핑을 적용한 모델과 완전한 레이어의 모델을 나타냅니다.

모델 아키텍처

LiteStage는 두 가지 상호 보완적 구성 요소를 포함합니다:

1. 오프라인 구성 (Offline Configuration)

Step 1: 레이어 중요도 추정

• 부분 레이어 수준의 코사인 유사도를 중요도 대리로 채택
• 다중 헤드 자기 주의(MHSA)와 피드포워드 네트워크(FFN)의 중요도를 각각 계산:

I^(j)_MHSA = (1/N) ∑_{n=0}^{N-1} cos(MHSA^(j)(x) + x, x)
I^(j)_FFN = (1/N) ∑_{n=0}^{N-1} cos(FFN^(j)(x) + x, x)

Step 2: 레이어 예산 검색

• 가장 느린 추론 단계부터 탐욕 검색 시작
• 정확도-지연시간 곡선을 구축하고 정확도 제약 조건을 만족하는 최적 지연시간 구성 선택
• 단계별 최적화로 단계 간 상호작용이 정확하게 반영되도록 보장

2. 온라인 조정 (Online Adjustment)

Step 3: 생성 조기 종료

• 최근 n개 토큰의 신뢰도 캐시 유지
• 평균 신뢰도 μ_Conf를 계산하고 임계값 이하일 때 생성 조기 종료
• 신뢰도는 각 토큰의 최대 로짓 값으로 정의

기술 혁신 포인트

1. 비균등 레이어 예산 할당: 각 단계의 민감도 차이에 따라 서로 다른 레이어 스킵핑 예산을 자적응적으로 할당하여 민감한 단계에서의 과도한 압축을 방지합니다.
2. 지연시간 인식 최적화: 정확도뿐만 아니라 실제 추론 지연시간을 종합적으로 고려하여 더 많은 레이어를 스킵하지만 지연시간이 더 높은 구성을 자동으로 제외합니다.
3. 동적 생성 제어: 신뢰도 모니터링을 통해 생성 길이를 능동적으로 제어하고 레이어 스킵핑으로 인한 중복 토큰 문제를 완화합니다.

실험 설정

데이터셋

TinyThinker의 세 단계 추론 프로세스를 채택하여 세 가지 질의응답 벤치마크에서 평가:

• OpenBookQA (OBQA): 개방형 질의응답 작업
• CommonSenseQA (CSQA): 상식 추론 질의응답
• StrategyQA: 전략적 추론 질의응답

평가 지표

• 정확도: 질의응답 정확률
• 가속비: 완전한 레이어 모델 대비 추론 속도 향상
• 지연시간: 종단 간 추론 시간

비교 방법

• SkipDecode: 점진적 심층 레이어 스킵핑
• UnifiedSkip: 주기적 레이어 스킵핑
• AdaSkip: 코사인 유사도 기반 부분 레이어 중요도 추정

구현 세부사항

• 주로 TinyLlama-1.1B-Chat-v1.0 모델 사용
• 10 에포크 훈련, 배치 크기 16(OBQA/CSQA) 또는 24(StrategyQA)
• 학습률 5×10^-5
• 자기 일관성 프로토콜로 10회 반복 평가 수행
• 신뢰도 임계값 0.5, 캐시 크기 n=5

실험 결과

주요 결과

세 가지 벤치마크 데이터셋에서 LiteStage는 기준선 방법을 크게 능가합니다:

주요 발견

1. 단계 민감도 차이: 단일 단계 스킵핑 실험을 통해 Stage 3이 레이어 스킵핑에 가장 민감하며, 그 정확도 곡선이 전체 성능 상한을 거의 결정함을 확인했습니다.
2. 지연시간 역설: 더 많은 레이어 스킵핑이 항상 더 빠른 추론을 가져오지는 않으며, 생성 길이 증가로 인해 일부 구성은 오히려 지연시간을 증가시킵니다.
3. 신뢰도 패턴: 레이어 스킵핑 모델의 토큰 신뢰도는 단조 감소 추세를 보이는 반면, 완전한 레이어 모델은 후기에 신뢰도가 회복될 수 있습니다.

소거 실험

비균등 레이어 예산의 효과:

• 동일한 레이어 스킵핑 수에서 LiteStage의 정확도는 균등 스킵핑 전략보다 현저히 높습니다
• 스킵핑 레이어 수가 증가함에 따라 성능 차이가 더욱 확대됩니다

생성 조기 종료의 기여:

• 경미한 레이어 스킵핑에서 생성 조기 종료의 영향은 미미합니다(-0.5% 디코딩 단계)
• 심각한 레이어 스킵핑에서는 최대 82.5%의 디코딩 단계를 줄일 수 있습니다
• 정확도는 기본적으로 안정적이며, 일부 경우 약간 향상되기도 합니다

사례 분석

CSQA의 구체적인 사례를 통해 생성 조기 종료가 낮은 신뢰도의 중복 텍스트를 효과적으로 잘라내면서 핵심 추론 논리를 완전히 유지하고 최종 답변을 일치시킬 수 있음을 보여줍니다.

관련 연구

다단계 생성

• TinyThinker: 회상-분석-요약의 세 단계 추론 순환 제안
• DeAR: 분해-분석-재고찰 프로세스 채택
• CasCoD: 분해된 사고의 연쇄 증류 방식
• Self-Discover: 추론 구조의 동적 조직화

레이어 스킵핑 기술

훈련 기반 방법:

• LayerSkip, DeeBERT, EE-LLM: 중간 레이어 조기 종료
• Mixture-of-Depth: 모델 및 라우터 훈련 필요

훈련 불필요 방법:

• SkipDecode: 점진적 심층 레이어 스킵핑
• Unified Skipping: 주기적 스킵핑
• ShortGPT: 코사인 유사도 기반
• AdaSkip: 부분 레이어 수준 중요도 추정

생성 조기 종료

기존 방법은 주로 장시간 추론 모델을 대상으로 하며, 모델 압축으로 인한 생성 연장 문제에 대한 관심이 부족합니다.

결론 및 논의

주요 결론

1. 다단계 추론의 비균등 민감도: 서로 다른 추론 단계는 레이어 압축에 대해 현저히 다른 민감도를 보이며, 차별화된 최적화 전략이 필요합니다.
2. 지연시간 인식 최적화의 필요성: 순수한 레이어 스킵핑은 생성 길이 증가로 인해 지연시간이 악화될 수 있으므로 정확도와 지연시간을 종합적으로 고려해야 합니다.
3. 생성 제어의 효과성: 신뢰도 기반 생성 조기 종료는 레이어 스킵핑으로 인한 중복 생성 문제를 효과적으로 완화할 수 있습니다.

한계

1. 오프라인 검색 오버헤드: 다른 훈련 불필요 방법과 비교하여 LiteStage의 오프라인 구성은 더 많은 계산 리소스가 필요합니다(약 1~7.6시간).
2. 모델 아키텍처 의존성: 주로 Llama 계열 모델에서 검증되었으며, Qwen 등 다른 아키텍처에서는 효과가 제한적입니다.
3. 적용 범위 제한: 다단계 추론 시나리오에 특화되어 있으며, 단일 단계 추론에 대한 적용 가능성이 충분히 검증되지 않았습니다.

향후 방향

1. 더 많은 모델 아키텍처로 확장: 서로 다른 아키텍처의 스킵핑 민감도 특성 연구
2. 동적 예산 할당: 런타임 자적응 레이어 예산 조정 메커니즘 개발
3. 다중 모달 추론 최적화: 프레임워크를 시각-언어 등 다중 모달 추론 작업으로 확장

심층 평가

장점

1. 문제 식별의 정확성: 다단계 추론의 핵심 병목 현상(단계 민감도 차이 및 중복 생성 문제)을 정확히 식별했습니다.
2. 합리적인 방법 설계: 오프라인-온라인 결합 프레임워크 설계가 정교하여 최적화 효과를 보장하면서 런타임 오버헤드를 제어합니다.
3. 충분한 실험 설계: 상세한 동기 실험, 소거 연구 및 사례 분석을 통해 방법의 효과성을 충분히 검증했습니다.
4. 높은 실용 가치: 훈련 불필요 방법으로서 우수한 실제 응용 전망을 가집니다.

부족한 점

1. 이론 분석 부족: 단계 민감도 차이에 대한 이론적 설명이 부족하며 주로 경험적 관찰에 의존합니다.
2. 휴리스틱 매개변수 설정: 신뢰도 임계값, 캐시 크기 등 핵심 매개변수가 주로 휴리스틱 설정에 기반하며 체계적 분석이 부족합니다.
3. 제한된 일반화성: 서로 다른 모델 아키텍처에서의 성능 차이가 크며 일반화 능력이 개선 필요합니다.

영향력

1. 학술적 기여: 다단계 추론의 레이어 스킵핑 최적화 문제를 처음으로 체계적으로 연구하여 관련 연구에 새로운 관점을 제공했습니다.
2. 실용적 가치: 소형 언어 모델의 효율적 추론을 위한 실용적 해결책을 제공하여 엣지 배포 추진에 도움이 됩니다.
3. 재현성: 완전한 코드 구현을 제공하여 후속 연구 및 응용을 용이하게 합니다.

적용 시나리오

LiteStage는 특히 다음 시나리오에 적합합니다:

• 리소스 제약이 있는 엣지 디바이스 배포
• 다단계 추론이 필요한 복잡한 작업
• 지연시간에 민감한 실시간 응용
• 소형 언어 모델의 추론 가속

참고문헌

논문은 다음을 포함한 여러 중요한 관련 연구를 인용했습니다:

• TinyThinker (Piao and Park, 2024): 다단계 추론의 대표적 연구
• AdaSkip (He et al., 2025): 부분 레이어 수준 레이어 스킵핑의 최신 방법
• Mixture-of-Depths (Raposo et al., 2024): 동적 계산 할당의 개척적 연구

전체 평가: 본 논문은 다단계 추론의 레이어 스킵핑 최적화 문제에 대해 혁신적인 해결책을 제시하며, 이론적 통찰력과 실제 효과 모두에서 현저한 기여를 합니다. 일부 한계가 있지만, 소형 언어 모델의 효율적 추론을 위한 새로운 연구 방향을 개척했으며, 중요한 학술적 가치와 실용적 의미를 가집니다.