[논문리뷰]UniICL: An Efficient Unified Framework Unifying Compression, Selection, and Generation (ACL, 2025)

Jun Gao, Qi Lv, Zili Wang, Tianxiang Wu, Ziqiang Cao, and Wenjie Li. 2025. UniICL:An Efficient ICL Framework Unifying Compression, Selection, and Generation. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (Eds.). Association for Computational Linguistics, Vienna, Austria, 500–510. https://doi.org/10.18653/v1/2025.acl-long.24

Problem Statement

[Context Length 폭증] 여러 demonstration을 단순히 연결(concatenation)하면 LLM의 입력 시퀀스 길이가 급격히 증가하여, 메모리 사용량이 크게 늘어나고 하드웨어 부담이 커진다.

[Indiscrimitnate Compression의 한계] 기존 Prompt Compression 기법들은 demonstration과 query를 구분하지 않고 모두 압축하기 때문에, 생성 과정에서 중요한 의미 정보를 잃어버리며 잘못된 응답을 유도할 수 있다. 또한 압축기의 훈련 비용이 크고, 추론 속도가 느려지거나 원래 모델의 입력 한계를 여전히 넘지 못하는 문제가 있다.

[Shallow-relevant Demonstration Selection의 한계] 기존 Retrieval-based ICL Example Selection (RICES) 기법은 단순한 lexical similarity 기반으로 demonstration을 선택하여, 실제 task에 도움이 되지 않는 얕은 관련성(shallow relevance)을 가진 예시들을 포함시키게 된다. 이로 인해 LLM이 오답을 출력하거나 성능이 저하된다.

(a) Prompt compression methods

• 기존 방법들은 demonstrations와 queries를 구분하지 않고 모두 virtual tokens로 압축한다. 이 경우 generator는 compressor가 만든 virtual tokens만을 기반으로 응답을 생성하게 되고, 그 결과 문법적으로 잘못된 답변을 내놓을 수 있다.

(b) Retrieval-based demonstration selection methods

• 기존 retrieval 방식은 off-the-shelf tools(e.g., 사전 학습된 Open source LLM)를 사용해 query와 얕은 수준에서 비슷한 demonstrations를 선택한다. 하지만 이러한 lexical demonstrations는 실제 label과 상반될 수 있어, LLM이 잘못된 답변을 생성하게 된다.

(c) UniICL

• UniICL은 demonstrations만을 discriminately compress하고, 그 compression 결과를 기반으로 demonstration selection을 수행한다. 이를 통해 shallow-relevant 예시 선택 문제를 완화하고, 효율적인 ICL generation을 가능하게 한다.

Methodology

Overview

UniICL은 i) demonstration compression, ii) demonstration selection, 그리고 iii) response generation을 통합한 Unified In-context learning(ICL) 프레임워크이다.

Frozen LLM을 기반으로 compressor와 generator를 공유하며, adapter와 학습 가능한 임베딩인 Memory Slot [M]을 도입한다. Demonstration은 Memory Slot을 통해 독립적으로 Memory Tokens로 압축되고, query와 demonstration 간 similarity를 바탕으로 relevant examples가 선택된다. 이후 선택된 Memory Tokens와 query가 generator로 입력되어 최종 응답이 생성된다. 또한, Demonstration Bank(DB)를 통해 이미 압축된 demonstration의 Memory Tokens를 캐시하여, 동일한 demonstration의 반복 압축을 피하고 효율성을 크게 향상시킨다.

Demonstration Compression

• 입력: Demonstration \(D_i\)
• 출력: Memory Tokens \(C_i = (c^i_1, ..., c^i_k)\)

i. 각 demonstration 뒤에 \(k\)개의 Memory Slots \(M = k \times [M]\)을 부착한다.

$$ H^i = f_{\theta}(D_i^{L_i \times d} \oplus M^{k \times d}), \quad H^i \in \mathbb{R}^{k \times d} $$

ii. Frozen Vicuna를 통해 forward propagation을 수행하고, Memory Slot 위치에서 hidden states \(H^i = (h^i_1, h^i_2, \ldots, h^i_k)\)를 얻는다.

$$ c_j^i = W_p^{d \times d} \cdot h_j^i $$

iii. Linear adapter \(W_p \in \mathbb{R}^{d \times d}\)를 적용하여 Memory Tokens \(C_i = (c^i_1, \ldots, c^i_k)\)를 얻는다.

Demonstration Compression 단계에서는 각 demonstration \(D_i\)를 독립적으로 Memory Tokens \(C_i\)로 변환하는 과정이다. 첫 단계에서는 demonstration 뒤에 \(k\)개의 learnable Memory Slots \(M\)을 부착한다. 이렇게 하면 실제 텍스트 토큰과 Memory Slot이 하나의 입력 시퀀스로 결합된다. 두 번째 단계에서는 Frozen Vicuna 모델을 통해 forward propagation을 수행하여, Memory Slot 위치의 hidden states \(H^i = (h^i_1, ..., h^i_k)\)를 얻는다. 이 hidden states는 demonstration의 의미적 정보를 요약해 담는다. 세 번째 단계에서는 Linear adapter \(W_p\)를 적용하여 hidden state를 Memory Tokens \(C_i\)로 변환한다. Adapter는 compressor와 generator의 표현 공간을 정렬해주며, 생성된 Memory Tokens는 이후 selection과 generation 단계에서 재사용된다.

Demonstration Selection

• 입력: Query \(Q\), Candidate Demonstrations \({D_1, \ldots, D_n}\)
• 출력: 선택된 Demonstrations \({D_{\text{sel}}}\)

$$ \bar C^i_{\{Q, D\}} = \frac{1}{k}\displaystyle\sum_{j=1}^k c_j $$

i. Query와 Demonstration 각각의 Memory Tokens를 average pooling하여 representation \(\bar{C}\)를 얻는다.

$$ S_i = \cos(\bar C_Q, \bar C_D^i) $$

ii. Query와 Demonstration representation 간 cosine similarity를 계산하여 saliency score \(S_i\)를 얻는다.

Demonstration Selection 단계에서는 query와 각 demonstration의 Memory Tokens를 average pooling하여 compact representation을 얻는다. 이 representation은 \(\bar{C}\)로 표기된다. 그런 다음 Query의 representation \(\bar{C}_Q\)와 각 demonstration의 representation \(\bar{C}^i_D\) 사이의 코사인 유사도를 계산하여 saliency score \(S_i\)를 산출한다. 이 score는 query와 demonstration의 의미적 유사도를 정량화한다. 마지막으로 score가 높은 demonstration을 선택해 최종 ICL 입력에 사용한다. 이 과정을 통해 shallow lexical similarity에 의존하는 기존 retrieval과 달리, semantic representation 기반의 selection이 가능해진다.

Response Generation

• 입력: 선택된 Demonstration Memory Tokens \({C_1, \ldots, C_m}\), Query \(Q\)
• 출력: 응답 \(y = (y_1, y_2, \ldots)\)

i. 선택된 demonstration Memory Tokens \({C_1, \ldots, C_m}\)를 concatenate한다.
ii. Query와 함께 generator \(g_\theta\)에 입력한다.

$$ y_i= g_{\theta}(C^1, \cdots, c; Q; y_{<i}) $$

iii. Auto-regressive generation으로 응답을 생성한다.

Response Generation 단계에서는 선택된 demonstration Memory Tokens를 가로로 concatenate한 후, 이를 실제 query와 함께 generator에 입력하여 응답을 생성한다. 이때 Memory Tokens는 virtual embeddings의 형태로 입력되며, query는 자연어 토큰 그대로 사용된다. Generator는 Memory Tokens와 query를 동시에 활용하여 auto-regressive 방식으로 응답 \(y = (y_1, y_2, ...)\)를 생성한다. 이 방식은 demonstration에서 압축된 의미와 query의 실제 정보를 모두 반영할 수 있다는 장점이 있다.

Training

$$\mathcal{L} = \mathcal{L}_{lm} + \mathcal{L}_{ctr}$$

i. Training objective는 Language Modeling Loss와 Contrastive Loss의 joint loss로 정의한다.

$$ \mathcal{L}_{lm} = -\frac{1}{|y|} \sum_{t=0}^{|y|} \log P(y_t \mid x_{ui}, C; y_{<t}), $$

ii. Language Modeling Loss (Next-token prediction) \(L_{lm}\)은 “메모리 토큰 \(C\)”을 조건으로 실제 쿼리 토큰과 과거 출력 \(y_{<t}\)가 주어졌을 때 정답 토큰 \(y_t\)의 log-likelihood를 최대화하도록 학습한다. iii. PPL (Perplexity) 기반 contrastive examples mining으로 \(D^+\)와 \(D^-\)를 자동 선정한다.

• 이 때, Demonstration 없이 정답 시퀀스 \(y\)에 대한 perplexity를 계산한 기준값을 \(ppl^{Q}\)로 정의한다.
• 각 Demonstration \(D_i\)를 제공한 상태로 같은 정답에 대한 perplexity를 \(ppl_i^D\)로 정의한다.
• 후보의 상대 PPL 이득을 다음과 같이 정의한다.

$$ \tilde{ppl}_i^D = ppl^Q - ppl^D_i, \; i \in [1, n] $$

• 후보의 상대 PPL 이득이 클 경우 해당 demonstration은 \(D^+\), 작으면 \(D^-\)로 선택한다.

$$ \mathcal{L}_{ctr} = \frac{\exp(\cos(C_Q, C_D^+))}{\exp(\cos(C_Q, C_D^+)) + \exp(\cos(C_Q, C_D^-))}. $$

iv. 선택된 \(D^+ , D^-\)에 대해 메모리 토큰에 대한 mean pooling으로 representation을 얻으면 이를 InfoNCE를 활용하여 대조 학습 (contrastive learning)에 이용한다.

Training 과정은 Adapter \(W_p\)와 Memory Slot \([M]\)을 학습하는 데 초점을 맞추며, 이는 전체 모델 대비 약 0.24%에 해당하는 17M 파라미터이다. Loss 함수는 두 가지 항으로 구성된다. Language Modeling Loss \(L_{lm}\)는 query와 demonstration Memory Tokens를 포함한 입력이 주어졌을 때 정답 시퀀스 \(y\)를 생성하도록 모델을 유도한다. 이 때, concatenation compression을 이용하여 긴 입력 \(x_c\)가 윈도 한계를 초과할 때 \(x_c\)를 분할하여 각각을 압축한 뒤, 생성 단계에서 사용할 하나의 글로벌 가상 토큰 집합으로 다시 결합한다. 이렇게 얻은 가상 토큰은 \(L_{lm}\)에 그대로 투입되어 정답 시퀀스 복원 학습을 진행하므로, 입력 길이 제약을 우회하면서 압축–생성 정합을 유지한다. 전 과정에서 학습되는 파라미터는 Adapter \(W_p\)와 Memory Slot \([M]\)뿐이어서 경량이며, 이로써 inference 품질과 selection 정합을 동시에 확보한다.”

대조 학습(Contrastive learning)을 하기 위해서는 positive와 negative의 정의가 필요하다. 이를 위해 논문에서는 perplexity를 사용하여 demonstration의 실질적 유용성에 따라 positive와 negative demonstration을 정하게 된다. 먼저 Demonstration 없이 정답 시퀀스 \(y\)에 대한 perplexity를 계산한 기준값을 \(ppl^{Q}\)로 정의하고, 각 Demonstration \(D_i\)를 제공한 상태로 같은 정답에 대한 perplexity를 \(ppl_i^D\)로 정의한다. 최종적으로 기준 perplexity대비, 각 demonstration을 제공했을 때, perplexity가 얼마나 감소하는지를 계산하여 가장 크게 낮추는 demonstration을 \(D^+\), 오히려 높이는 demonstration를 \(D^-\)로 자동 발굴한다. 이 demonstration들은 메모리 토큰과 마찬가지로 mean pooling으로 vector representation을 구하고 코사인 유사도 기반 InfoNCE를 통해 대조 학습한다.

Demonstration Bank (DB)

i. 각 demonstration \(D_i\)를 한 번만 압축하여 메모리 토큰 \(C_i\)를 얻는다.
ii. 얻은 \(C_i\)를 DB에 캐싱하고, 관련 점수(예: \(\tilde{ppl}\) 기반 지표 또는 과거 saliency)와 함께 인덱싱한다.
iii. 동일한 demonstration가 재사용될 때는 재압축 없이 DB에서 \(C_i\)를 즉시 불러와 selection 및 generation에 사용한다.

DB는 demonstration 압축이 쿼리와 독립적이라는 성질을 활용해 반복 압축을 제거하기위한 모듈이다. 1단계에서 각 demonstration \(D_i\)를 한 번만 압축해 \(C_i\)를 얻고, 2단계에서 이를 캐시해두면, 3단계처럼 후속 쿼리에서 동일한 demonstration가 필요할 때 즉시 조회하여 latency를 줄이고 throughput을 높일 수 있습니다. 이때 캐시 엔트리는 demonstration 식별자에 매핑된 \(C_i\)와, 선택 과정에서 활용 가능한 부가 점수(예: \(\tilde{ppl}\) 기반 랭킹 또는 과거 saliency)로 구성되어, selection 재계산을 가볍게 하고 재현성 있는 재사용을 가능하게 합니다. 결과적으로 DB는 대규모 demonstration 풀에서도 메모리·시간 비용을 선형적으로 억제하면서, UniICL의 압축–선택–생성 파이프라인을 안정적으로 확장시키는 역할을 수행합니다.

Experiments

Main Results

Table 3은 CoLA-dev, SST-2-dev, IMDb 데이터셋에서 고자원(high-resource) 및 저자원(low-resource) ICL 설정에서의 성능을 비교한 결과이다.

• 괄호 안 수치는 저자원 ICL 결과를 의미한다.
• UniICL은 단순히 S-BERT로 사전 선별된 예제를 사용하는 경우보다, 자체 선택(♠)을 적용했을 때 성능이 더 좋아졌다.
• 또한, 대조 학습 기반 선택 강화(+Lctr)를 적용하면 성능이 추가로 향상되었으며, 특히 시연 수(shots)가 늘어날수록 선택 강화의 이점이 커졌다.
• 전반적으로 UniICL은 AutoCompressor, LLMLingua, ICAE 등 기존 압축 기반 방법들보다 높은 정확도(Acc.)와 ROUGE 점수(R-1, R-2, R-L)를 달성하였다.

Efficiency Comparison

Figure 8은 CoLA 데이터셋에서 shots 수를 0→64까지 늘려가며 메모리 사용량(GB)과 처리량(Throughput, iter/s)을 비교한 그림이다.

• Baseline (원본 Vicuna)는 8-shot에서 이미 메모리 폭발(memory explosion)이 발생했다.
• AutoCompressor와 ICAE는 최대 32-shot까지만 처리 가능했지만, UniICL은 64-shot까지 안정적으로 확장할 수 있었다.
• Demonstration Bank(DB)를 활용한 UniICL+Caching은 이미 압축된 예제를 재사용하여 추가 지연(latency)을 제거했고, 처리량도 다른 방법들과 비교해 크게 저하되지 않았다.

결과적으로 Figure 8은 UniICL이 메모리 효율성과 확장성 측면에서 가장 우수하며, DB를 활용하면 지연까지 최소화할 수 있음을 시각적으로 보여준다

Computational Latency Analysis

Table 7은 Vicuna, Vicuna-1k, UniICL의 계산 효율성을 GPUHours, TFLOPs, TMACs 지표로 비교한 것이다.

• Vicuna: 1.5 GPUHours, 8,620 TFLOPs, 4,309 TMACs.
• Vicuna-1k: 긴 입력 윈도우(1k)를 지원하지만, 그 결과 1.9 GPUHours, 31,664 TFLOPs, 15,832 TMACs로 연산량과 비용이 크게 증가하였다.
• UniICL: 1.6 GPUHours, 22,437 TFLOPs, 11,218 TMACs로, 기본 Vicuna 대비 다소 증가했지만, Vicuna-1k 대비 크게 효율적이었다.

즉, UniICL은 입력 압축을 통해 윈도우 한계를 극복하면서도 Vicuna-1k보다 훨씬 효율적인 계산 성능을 보였다.

NR 카테고리 내 다른 글 보러가기