References
https://arxiv.org/abs/2502.20583
LiteASR: Efficient Automatic Speech Recognition with Low-Rank Approximation
Modern automatic speech recognition (ASR) models, such as OpenAI's Whisper, rely on deep encoder-decoder architectures, and their encoders are a critical bottleneck for efficient deployment due to high computational intensity. We introduce LiteASR, a low-r
arxiv.org
Introduction
현 시점 가장 많이 쓰이고 있는 사전학습 음성인식(ASR) 모델을 뽑자면 CTC 기반 Conformer, AED(Attention Encoder-Decoder, seq2seq 스타일)기반 Whisper 두 가지로 추릴 수 있다.
그렇다 보니 ASR의 경량화도 보통 이 두 종류의 모델을 기준으로 수행한다.
본 논문은 그 중 Whisper를 기준으로, Tensor Decomposition이라는 경량화 계열 방법론을 중심으로 인코더 파라미터 수 및 연산 속도를 최적화하는 LITEASR 이라는 기법을 제안한다.
시작하기 전, Tensor Decomposition(= Low-rank approximation)?
한국어로 텐서 분해라고 부르는 이 방법론은, n개의 rank를 갖는 기존 모델 weight matrix를, r개 rank를 갖는 matrices로 분할하는 경량화 방법론이다.
'어떤 기준으로 어떻게 분해하냐' 는 제안하는 세부 방법론마다 조금씩 다른데,
그림과 설명만 보았을 때 matrix 수가 더 늘어났는데 이게 어떻게 경량화 기법인가 의문이 들 수 있다.
일단 이게 왜 경량화 기법이냐 하면 수식으로 이해하는 게 직관적이다.
y=Wx 라는 연산이 존재한다고 가정하자. 이때 x는 input, W는 model weight이다.
W는 M*N, 즉 N의 입력차원 및 M의 출력차원을 갖는다. 그럼 파라미터 수는? M*N이다.
그런데 Tensor Decomposition 이후, n개 tensor를 갖던 W가, (M*R 차원의) A * (R*N차원의) B (= W~= AB)로 표현된다고 가정하자.
그럼 연산은 y=Wx에서 → y=A(Bx)로 바뀐다. 이때 파라미터 수는? M*R + R*N = R(M+N) 으로 바뀐다.
이때 M=N=1024이고 R=128이었다면, 파라미터 수는 1024^2 =1,048,576 에서 128(1024 + 1024) = 262,144 로 확연히 줄어든다. (그리고 동일 계열 연산에서 파라미터 수가 줄어들면 당연히 연산량도 감소한다.)
결과적으로 이전보다 더 적은 파라미터로, 최소한의 성능 손실(rank 소거 시 가장 effective한 rank를 살리는 방향으로 수행하기 때문)을 유도할 수 있다.
LITEASR, Tensor Decomposition을 어떻게 적용했나
앞서 모델 weight를 low-rank로 분해하는 디테일은 방법론마다 다르다고 언급했는데, 이 논문은 PCA(차원 축소)를 기준으로 분해를 수행한다.
일단 파이프라인부터 요약하자면 다음과 같다.
- 소량의 calibration audio를 인코더에 통과시킨다.
- 각 self-attention 및 MLP 선형층의 출력 activation을 수집한다.
- 각 층의 activation에 PCA/SVD를 적용하여 주요 출력 부분공간을 구한다.
- 누적 설명 분산이 임계값 를 넘도록 층별 rank k를 자동 결정한다.
- 기존 선형층 XW+b를 두 개의 저랭크 선형변환으로 치환한다.
- Q/K/V rank가 충분히 작을 경우 attention 계산 순서를 변경한다.
1. Calibration activation 수집
데이터셋의 특성을 충분히 대표한다고 가정할 수 있는 소량의 서브셋을 먼저 준비한다.
이 서브셋을 바탕으로 모델을 forward하여 output Y를 수집한다.
2. Activation PCA
위에서 수집된 activation의 평균을 μ_Y라고 하면 centered activation은 Y_c = Y - μ_Y으로 둘 수 있다.
이 Y_c에 SVD를 수행하면 Y_C=U∑V^T 로 나타내진다.
그리고 PCA에 의해 V의 첫 k개 열을 V_k라고 하면 Y - μ_Y ~= (Y - μ_Y)V_k*V_k^T 로 둘 수 있다.
3. Linear layer 재구성 및 저랭크 변환
원래 식 Y=XW + b를 위의 PCA 근사식에 대입하면Y ~= (XW+b - μ_Y)V_k*V_k^T + μ_Y = X(W*V_k)V_k^T + [μ_Y + (b - μ_Y)V_k*V_k^T]로 나타내진다.따라서 분해된 W1, W2와 bias를 다음과 같이 나타낼 수 있다.
최종 압축 선형층은 Y' = XW1W2 + b~ 로 나타낼 수 있겠다.
PCA, SVD 위주 수식 전개 기반의 방법론이라 설명을 풀어 쓰기 쉽지 않다..
LITEASR, Tensor Decomposition을 어디에 적용했나
Whisper의 모든 layer가 항상 같은 effective rank 수를 갖지는 않는다.
Transformer 초기 레이어는 음향 상관성을 많이 포함해 강하게 압축 가능하고, 뒤로 갈 수록 transcription에 직접적인 정보를 많이 만들어내므로 초기보다 더 많은 rank가 필요할 수 있다.
그렇기 때문에 논문은 각 layer마다 남겨둘 rank k를 다르게 두었다.
이때 원래 성능 보존률 𝛉 값을 정의해 두고, 랭크 수를 줄였을 때 원래 activation variance의 𝛉% 미만으로 성능이 줄어들면 거기까지가 최소 k의 마지노선이다.
다만, Transformer 진입 전 Conv layer의 경우 이 방법론으로 압축하였을 경우 성능 손실이 커 압축을 따로 수행하지 않았다고 한다.
그렇게 해서 줄어든 파라미터 및 성능 변화는 아래와 같다.
과연 속도는 얼마나 빨라졌을까?
가장 많이 압축을 수행한 (C)에선 약 기존대비 1.5배 정도의 속도 향상이 관찰되었다.
Discussion
Tensor Decomposition류 방법론 자체가 수학적 개념에 크게 의존하다 보니 여러모로 글로 쉽게 풀어 쓰긴 어려운 것 같다.
가장 중요했던 포인트는 텐서 분해를 수행하는 직접적인 방법론에 PCA를 활용했고, 레이어별로 성능 보존도를 측정하여 동적으로 경량화를 수행했다는 점으로 볼 수 있겠다.
다만 초기 전제조건인 'Calibration set'에 대한 activation 수집 절차는 재현성도 그렇고 모든 실험에서 일관적인 성능을 내기 어렵게 하는 요소라는 점에서는 한계가 있어 보인다.
재현하는 사람이 사용할 데이터셋에서 '알아서' 대표성을 띄는 작은 데이터셋을 임의로 만들어야 하니 말이다. 이 방법론 자체가 이 부분에 크게 병목을 띄고 있는 건 아닌가 싶다.
다만 ASR 분야에서 크게 시도되지 않던 Tensor Decomposition의 구체적인 적용사례라는 점에서 반가운 논문이었다.
또 일반적으로 Tensor Decomposition은 기존보다 레이어 수가 늘어나는 구조다 보니 wall-clock 측면에서 속도 향상에 제한적일 때가 많은데, 본 논문은 아주 크진 않을지라도 실질적 속도 향상을 이루어냈다.
정리하기 난해한 부분들이 이 글에서 일부 생략되었으니, 관심있다면 논문을 더 자세히 읽어보길 권한다.