8 분 소요

0. Introduction

Paper link

FlexRank는 model compression을 단순히 작은 모델 하나를 만드는 문제가 아니라, 하나의 pretrained model에서 여러 deployment budget에 맞는 nested submodel들을 뽑아내는 문제로 본다. 이 관점은 실제 서비스에서 중요하다. 같은 model을 항상 같은 latency와 같은 cost로 serving할 필요는 없다. 요청 난이도, device, traffic, cost budget에 따라 더 작은 path를 쓰거나 더 큰 path를 쓸 수 있어야 한다.

기존 compression이나 PEFT 접근은 대개 fixed budget을 전제로 한다. 특정 rank, 특정 pruning ratio, 특정 quantization setting에 맞춰 모델을 만든다. 하지만 실제 deployment에서는 budget이 여러 개다. GPU server, edge device, batch inference, low-latency route가 모두 다르다. FlexRank는 이를 train-once, deploy-everywhere 문제로 바꾼다.

한 줄 요약: FlexRank는 pretrained model의 linear layer를 low-rank factor로 분해하고, 중요도 순서가 있는 nested submodel chain을 찾아 distillation으로 consolidate해, 하나의 weight set에서 여러 budget의 model을 선택적으로 사용할 수 있게 만드는 elastic deployment 방법이다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같다.

  • 모델 압축을 single target compression이 아니라 budget-adaptive deployment 문제로 재정의한다.
  • Low-rank decomposition, dynamic programming rank search, nested submodel training, knowledge distillation을 하나의 pipeline으로 연결한다.
  • 여러 submodel을 따로 학습하지 않고 shared weights 안에서 Pareto frontier에 가까운 submodel chain을 만든다.
  • LLM과 ViT 모두에서 cost-performance trade-off를 다룬다.
  • 실제 inference에서 low-rank factorization을 deployable하게 만들기 위해 GAR reparameterization을 포함한다.

FlexRank는 “작은 모델을 만들자”가 아니라 “하나의 모델이 여러 크기로 접히고 펴질 수 있게 만들자”에 가깝다.

1. Problem Setting

1-1. Problem definition

Large model deployment에는 두 가지 rigidity가 있다.

첫째, model cost가 고정되어 있다. 하나의 checkpoint는 항상 같은 parameter count, 같은 compute path, 같은 latency를 요구한다. 실제 workload가 다양해도 model은 하나의 operating point에서만 동작한다.

둘째, budget별 모델을 따로 만들기 어렵다. 0.4x, 0.6x, 0.8x, 1.0x model을 각각 training하거나 compression하면 비용이 크고, maintenance도 복잡해진다. 각각의 model이 서로 독립적이면 weight sharing이나 cache reuse도 어렵다.

FlexRank의 목표는 다음과 같다.

하나의 pretrained model에서 nested submodels를 추출하고, deployment budget에 따라 필요한 rank components만 활성화할 수 있게 만들자.

여기서 nested는 중요하다. 작은 submodel이 큰 submodel의 prefix가 되어야 한다. 그래야 하나의 factorized weight set 안에서 budget을 늘릴 때 components를 추가하는 방식으로 동작할 수 있다.

1-2. Why previous approaches are insufficient

기존 방법의 한계는 네 가지다.

첫째, fixed-rank low-rank approximation은 한 budget만 겨냥한다. 특정 layer rank를 정하면 그 configuration에서는 좋을 수 있지만, 다른 budget에서는 다시 rank selection과 recovery가 필요하다.

둘째, pruning이나 layer skipping은 architecture heuristic에 의존한다. 어떤 layer를 줄일지, 어떤 depth를 건너뛸지에 대한 선택이 model과 task마다 달라질 수 있다.

셋째, 모든 submodel을 동시에 잘 만들려고 하면 objective conflict가 생긴다. 논문은 all-subspaces learning이 submodel 간 interference를 만들 수 있다고 설명한다. 각 submodel이 다른 rank approximation을 원하면서 shared parameters를 서로 방해한다.

넷째, low-rank factorization은 이론상 parameter를 줄여도 실제 inference speedup이 바로 나오지 않는다. Naive factorization은 kernel efficiency가 낮을 수 있고, dense kernel 대비 overhead가 생긴다. 그래서 deployable reparameterization이 필요하다.

2. Core Idea

2-1. Main contribution

FlexRank의 핵심은 low-rank components를 importance-ordered nested chain으로 만든다는 점이다. 전체 pipeline은 세 단계다.

  1. Layer Decomposition
    • 각 linear layer를 독립적으로 DataSVD 방식으로 factorize한다.
    • Activation data를 사용해 layer별 중요한 low-rank direction을 찾는다.
  2. Nested Submodel Search
    • 각 layer에서 rank를 줄였을 때의 cost saving과 reconstruction error를 probe한다.
    • Dynamic programming으로 budget별 rank configuration을 찾는다.
    • Configuration들은 nested chain을 이루도록 정리된다.
  3. Knowledge Consolidation
    • 단순 decomposition만으로는 layer 간 dependency를 반영하기 어렵다.
    • Base model을 teacher로 두고 distillation을 수행해 shared elastic weights를 refine한다.

이 과정을 통해 하나의 weight set 안에 여러 budget의 submodel이 들어간다.

2-2. Design intuition

FlexRank의 핵심 직관은 다음과 같다.

좋은 elastic model은 모든 possible submodel을 잘 만드는 것이 아니라, Pareto frontier를 따라 이어진 nested submodel chain을 잘 만들어야 한다.

모든 subspace 조합을 잘 만들려고 하면 submodel 간 objective가 충돌한다. 반면 nested chain은 작은 model에서 큰 model로 갈 때 components를 하나씩 추가하는 구조다. 추가되는 component는 이전 submodel이 못 담은 residual knowledge를 담당하면 된다.

이 구조를 수식으로 단순화하면 다음과 같이 볼 수 있다.

\[M_r = U_{1:r} V_{1:r}^{T}\]

rank $r+1$ submodel은 rank $r$ submodel을 포함한다.

\[M_r \subset M_{r+1}\]

이 prefix 구조가 있어야 budget을 늘릴 때 model이 자연스럽게 확장된다. FlexRank는 이 nested property를 model compression의 중심에 둔다.

3. Architecture / Method

3-1. Overview

Item Description
Goal 하나의 pretrained model에서 budget별 nested submodel 생성
Core operation low-rank decomposition of linear layers
Search dynamic programming rank selection
Training knowledge distillation from base model
Key property nested submodel chain
Deployment budget에 따라 rank components 선택
Extra module GAR for deployable inference efficiency
Evaluated models GPT-2, Llama family, DINOv3 ViT

3-2. Module breakdown

1) DataSVD layer decomposition

FlexRank는 각 layer weight $W_l$를 low-rank factors로 분해한다. 일반 SVD와 달리 activation statistics를 반영하는 DataSVD를 사용한다. 이는 weight reconstruction error만 보는 것이 아니라, 실제 input activation이 통과했을 때 중요한 direction을 더 잘 보려는 의도다.

각 layer에서 rank를 낮추면 parameter cost는 줄지만 error가 생긴다. FlexRank는 이 layer-level sensitivity를 모아 global budget allocation 문제로 바꾼다.

Layer가 많고 rank level도 여러 개면 가능한 configuration 수가 폭발한다. Brute force로 모든 layer-rank 조합을 평가하기 어렵다. FlexRank는 각 layer의 candidate와 error estimate를 이용해 dynamic programming으로 Pareto candidate를 찾는다.

이 문제는 multi-choice knapsack처럼 볼 수 있다. 각 layer에서 하나의 rank option을 고르고, 전체 saving budget 안에서 error를 최소화한다. 논문은 DP가 additive error assumption을 사용한다고 설명한다. Deep network에서는 완전한 assumption은 아니지만, ranking consistency를 얻기 위한 tractable probe로 사용한다.

3) Nested submodel learning

논문에서 가장 중요한 이론적 포인트는 Nested Subspace Learning이다. 모든 submodel 조합을 동시에 학습하면 conflict가 생긴다. FlexRank는 각 rank에 대해 하나의 prefix submodel만 최적화하는 구조를 만든다.

이 구조에서는 rank가 커질수록 이전 rank가 이미 배운 부분을 유지하고, 새 component는 residual을 배우면 된다. 그래서 Pareto frontier에 가까운 elastic chain을 만들 수 있다.

4) Knowledge consolidation

Decomposition과 DP search만으로는 충분하지 않다. 각 layer를 독립적으로 factorize하면 cross-layer dependency가 깨질 수 있다. FlexRank는 base model output을 teacher로 두고, submodel을 stochastic하게 sample하면서 distillation을 수행한다.

즉 여러 budget의 submodel이 같은 factorized parameter를 공유하면서 base model behavior를 따라가도록 consolidate한다.

5) Gauge-Aligned Reparametrization

Low-rank factor를 그대로 inference에 쓰면 dense kernel 대비 효율이 낮을 수 있다. FlexRank는 GAR를 통해 selected rank configuration을 실제 inference saving으로 더 잘 연결하려 한다. 이 부분은 compression 논문에서 자주 놓치는 지점이다. Parameter count가 줄어도 kernel이 느리면 deployment 이득이 줄어든다.

4. Training / Data / Recipe

4-1. Data

논문은 NLP와 CV를 모두 다룬다.

  • NLP: GPT-2, Llama-3.2-1B, Llama-3.2-3B, Llama-3.1-8B
  • Calibration: FineWebEdu-10BT
  • CV: DINOv3 ViT models
  • CV downstream: ImageNet1K validation

이 선택은 FlexRank가 특정 architecture trick이 아니라, linear layer를 가진 다양한 Transformer 계열에 적용될 수 있음을 보여주려는 의도다.

4-2. Training strategy

Training recipe는 다음 흐름이다.

  1. Calibration data에서 layer activation을 수집한다.
  2. 각 layer에 DataSVD를 적용해 ordered components를 만든다.
  3. 각 layer-rank candidate의 cost saving과 error를 평가한다.
  4. DP로 nested Pareto configurations를 찾는다.
  5. 여러 budget profile을 sample하며 base model distillation을 수행한다.
  6. Deployment 시 target budget에 맞는 profile을 선택한다.

4-3. Engineering notes

실무적으로 주의할 점은 다음과 같다.

  1. Elasticity는 model compression과 routing 사이에 있다.
    • Budget-conditioned inference나 input-adaptive routing과 결합할 수 있다.
  2. Calibration data quality가 중요하다.
    • DataSVD와 rank sensitivity는 activation distribution에 의존한다.
  3. Distillation budget이 Pareto frontier를 좌우한다.
    • Decomposition만으로는 좋은 deployable submodel이 나오기 어렵다.
  4. Inference kernel efficiency를 확인해야 한다.
    • Low-rank parameter count가 줄어도 actual latency가 같이 줄지 따로 측정해야 한다.
  5. Submodel별 downstream LoRA fine-tuning도 가능하다.
    • 논문은 math, code domain에서 FlexRank submodel에 LoRA를 붙이는 실험을 제시한다.

5. Evaluation

5-1. Main results

FlexRank는 GPT-2 loss, Llama family의 commonsense zero-shot accuracy, DINOv3 ViT의 ImageNet1K validation 등을 통해 평가된다. 핵심 주장은 FlexRank가 기존 rank-based compression보다 더 smooth한 cost-performance trade-off를 만든다는 것이다.

논문은 Llama-3.2-1B와 Llama-3.2-3B에서 math와 code domain LoRA fine-tuning을 보여준다. Relative size가 1.0x, 0.8x, 0.6x, 0.4x로 줄어들수록 성능은 내려가지만, constrained budget에서도 meaningful performance를 유지한다는 점을 강조한다.

또한 LLM-Pruner, LayerSkip 같은 다른 family와도 비교해, rank-based 방법 너머에서도 경쟁력 있는 trade-off를 보인다고 주장한다.

5-2. What really matters in the experiments

FlexRank의 실험에서 중요한 것은 최고 성능 하나가 아니다. 다음이 핵심이다.

  1. Degradation curve가 smooth한가
    • Budget을 줄였을 때 성능이 급격히 무너지는 것이 아니라 점진적으로 낮아져야 한다.
  2. Nested chain이 independent submodel보다 나은가
    • Budget별 model을 따로 학습하면 cost가 커진다.
    • FlexRank는 shared weight 안에서 여러 budget을 지원한다.
  3. DP allocation이 layer importance를 반영하는가
    • 모든 layer를 균등하게 줄이는 것이 아니라, 중요한 module은 rank를 오래 유지해야 한다.
  4. Low-rank saving이 실제 inference saving으로 이어지는가
    • GAR와 latency 측정이 중요한 이유다.

FlexRank의 핵심 가치는 compression ratio 숫자보다 운영 model에 있다. 여러 budget의 model을 따로 관리하지 않고, 하나의 elastic weight set으로 latency tier를 만들 수 있다는 점이 훨씬 실용적이다.

6. Limitations

  1. Adaptive routing은 직접 평가하지 않는다.
    • FlexRank는 budget-conditioned inference를 가능하게 하지만, 어떤 input에 어떤 budget을 쓸지 결정하는 router는 후속 과제로 남긴다.
  2. Training set과 consolidation budget에 민감하다.
    • 논문도 더 긴 training, 더 다양한 data, nested submodel에 특화된 optimization이 Pareto frontier를 개선할 수 있다고 본다.
  3. Training overhead가 있다.
    • Factorized model은 training 중 memory overhead가 생기고, high-rank setting에서는 dense forward보다 느릴 수 있다.
  4. DP의 additive error assumption은 근사다.
    • Deep network의 layer dependency는 비선형적이므로, DP는 exact optimality보다 tractable ranking probe로 이해해야 한다.
  5. Inference speedup은 hardware와 kernel에 의존한다.
    • Theoretical parameter reduction이 항상 wall-clock latency reduction으로 이어지지는 않는다.

7. My Take

7-1. Why this matters for my work

FlexRank는 실제 deployment 관점에서 매우 유용한 질문을 던진다. 모델을 하나의 고정된 artifact로 보지 않고, budget에 따라 여러 operating point를 갖는 elastic artifact로 본다. 서비스에서는 이 관점이 중요하다.

예를 들어 같은 QA system에서도 쉬운 query는 작은 submodel로 처리하고, 어려운 query는 큰 submodel로 보낼 수 있다. Traffic spike에는 더 작은 budget을 쓰고, offline batch에는 큰 budget을 쓸 수 있다. FlexRank는 이런 routing strategy의 model-side 기반을 만든다.

7-2. Reuse potential

재사용 가능성이 있는 부분은 다음과 같다.

  • LLM inference tiering
  • Edge and server shared model deployment
  • Budget-aware model serving
  • Speculative decoding draft model family 구성
  • Input-adaptive compute routing
  • Multi-tenant inference cost control

실무에서는 FlexRank 자체보다도 “nested Pareto chain” 개념을 먼저 가져갈 수 있다. 압축 모델을 만들 때 0.5x 하나만 만들지 말고, 여러 budget을 공유 weight 안에서 함께 관리하는 방향이다.

7-3. Follow-up papers

  • LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
  • LLM-Pruner: On the Structural Pruning of Large Language Models
  • SVD-LLM and SVD-LLM V2
  • ACIP low-rank elastic compression
  • Once-for-All and slimmable network 계열 논문

8. Summary

  • FlexRank는 pretrained model을 importance-ordered low-rank components로 분해한다.
  • DP rank search로 budget별 nested Pareto configuration을 찾는다.
  • Knowledge distillation으로 여러 submodel을 shared weight 안에서 consolidate한다.
  • 핵심은 fixed compression이 아니라 train-once, deploy-everywhere elastic deployment다.
  • 실제 적용에서는 routing policy, calibration data, kernel efficiency를 함께 봐야 한다.

댓글남기기

참고