[논문리뷰]GRACE: Generative Representation Learning via Contrastive Policy Optimization(arXiv, 2025)

Jiashuo Sun, Shixuan Liu, Zhaochen Su, Xianrui Zhong, Pengcheng Jiang, Bowen Jin, Peiran Li, Weijia Shi, and Jiawei Han. 2025. GRACE: Generative Representation Learning via Contrastive Policy Optimization.

Problem Statement

이 논문은 “대규모 언어모델(LLM)을 고품질 임베딩 모델로 학습하되, 생성·추론 능력을 보존하면서 임베딩 품질을 체계적으로 향상시키는 방법”을 제안하는 연구이며, 구체적 태스크는 대규모 문장 임베딩 벤치마크(MTEB) 전 영역에서의 텍스트 임베딩 학습 및 검색/유사도 판단 성능 향상이다. 핵심 아이디어는 기존의 대조학습(contrastive learning) 목표를 “최소화할 손실”이 아니라 생성 정책을 유도하는 보상(reward) 으로 재해석하여 RL 기반 정책최적화로 학습하는 것, 즉 GRACE 프레임워크이다.

[LLM을 ‘블랙박스 인코더’로 쓰는 한계] 기존 대조학습 기반 LLM 임베딩은 LLM을 정적 임베딩 생성 함수로 다루며, 모델의 생성·추론 능력을 버리고 고정 임베딩만 산출하게 만든다는 문제가 있다.

[토큰 비특이적(agnostic) 대조손실의 부작용] 단순 InfoNCE류 손실은 정책(생성)과 분리된 채 임베딩을 직접 밀어붙여 표현 붕괴/스케일 드리프트를 야기하고, LLM의 일반화 능력을 훼손할 수 있다.

Methodology

0. Preliminaries

Preliminaries 1. Policy Gradient Optimization

정책 \(\pi_{\theta}\)가 프롬프트 \(x\)에 대해 시퀀스 \(y\)를 생성하고, 기대 보상(expected reward)을 직접 최대화하는 방식으로 최적화를 진행하는 방식이다. 시퀀스 확률은 토큰 분해로 정의된다.

$$\pi_{\theta} (y \mid x) = \prod_{t=1}^n \pi_{\theta}(y_t \mid x, y_{<t})$$

• \(x\): 입력 프롬프트
• \(y = (y_1, \cdots, y_n)\): 정책이 생성한 응답 시퀀스
• \(\pi_{\theta}(y_t \mid x, y_{<t})\): 시점 \(t\)에서의 조건부 토큰 생성 확률
• \(\theta\): 정책(LLM)의 파라미터

Policy Gradient Optimization에서의 목적함수는 기대보상이다.

$$J(\pi_{\theta}) = \mathbb{E}_{y\sim\pi_{\theta}(\cdot \mid x)} [r(x, y)]$$

• \(r(x, y)\): LLM 응답 \(y\)의 품질을 평가하는 스칼라 reward
• \(\mathbb{E}_{y \sim \pi_{\theta}}\): 정책이 유도하는 분포에서의 기댓값

Policy gradient theorem에 의해, 기댓값의 그래디언트는 다음과 같다.

$$\nabla_{\theta}J(\pi_{\theta}) = \mathbb{E}_{y \sim \pi_{\theta}(\cdot \mid x)}=[r(x,y)\nabla_{\theta}\log\pi_{\theta}(y \mid x)]$$

참고로, 위의 식처럼 목적 함수의 그래디언트에는 바이어스가 없는 unbiased 상태이다. Policy gradient 추정의 분산을 줄이기 위해서 바이어스 \(b(x)\)를 빼서 어드밴티지를 형성하면 다음과 같이 수식이 변한다.

$$\nabla_{\theta}J(\pi_{\theta}) = \mathbb{E}_{y \sim \pi_{\theta}(\cdot \mid x)}=[(r(x,y) - b(x))\nabla_{\theta}\log\pi_{\theta}(y \mid x)]$$

• \(b(x)\): 샘플링된 행동과 독립적인 기준선으로, 분산을 줄위기 위해 사용

Preliminaries 2. GRPO

Group Relative Policy Optimization (GRPO)는 입력 \(x\)당 이전 정책 \(\theta_{\text{old}}\)에서 \(G\)개의 응답을 샘플링하고, 그룹의 평균 대비 정규화된 어드밴티지로 정책을 업데이트하는 방법이다. 먼저 그룹 표준화 어드백티지를 정의한다.

$$\hat{A}_i = \frac{r(x, y_i) - \frac{1}{G}\sum_{j=1}^G r(x, y_i)}{\text{std}(\{r(x, y_j)\}_{j=1}^G)}$$

• \(\{ y_1, \cdots, y_G\} \sim \pi_{\theta_{\text{old}}}(\cdot \mid x)\): 이전 정책에서 샘플링된 그룹 응답
• \(r(x, y_i)\): \(i\)번째 응답의 보상
• \(\text{std}(\cdot)\): 그룹 보상의 표준편차
• \(\hat A_{i, t} = \hat A_i\): 토큰 레벨로 동일한 어드밴티지를 공유

이 때, 중요도 비율(importance ratio)은 토큰 레벨로 정의된다.

$$r_{i, t}(\theta) = \frac{\pi_{\theta}(y_{i, t} \mid x, y_{i, <t})}{\pi_{\theta_{\text{old}}}(y_{i, t} \mid x, y_{i, <t})}$$

중요도 비율의 분자는 현재 정책, 분모는 이전 정책의 조건부 확률이다. KL divergence 항을 제외한 GRPO의 목적함수는 토큰 단위 클리핑을 포함한다.

$$J_{\mathrm{GRPO}}(\theta)= \mathbb{E}\!\left[  \frac{1}{G}\sum_{i=1}^{G}\;  \frac{1}{|y_i|}\sum_{t=1}^{|y_i|}  \min\!\Big( r_{i,t}(\theta)\,\hat{A}_{i,t}\;,\; \operatorname{clip}\!\big(r_{i,t}(\theta),\,1-\epsilon,\,1+\epsilon\big)\,\hat{A}_{i,t}  \Big)\right]$$

• \(\vert y_i \vert\): \(i\)번째 응답의 토큰 길이
• \(\epsilon\): PPO류 클리핑 허용 범위 하이퍼파라미터
• \(\text{min}(\cdot, \cdot)\): 클립 전후의 항 중 작은 값을 선택하여 과도한 업데이트를 방지

요약하면, GPRO는 그룹 표준화 어드밴티지 \(\hat A_i\)를 사용하여 응답 간 상대적 성능 차이를 안정적으로 반영하고, 중요도 비율 클리핑으로 신뢰구간 내에서 정책을 업데이트함으로써 학습을 안정화하는 방법이다.

1. Rationale-Generating Policy

LLM의 파라미터를 정책 \(\pi_{\theta}\)으로 하고, \(\pi_{\theta}\)가 입력 \(x\)에 명시적 추론 흔적(rationale) \(r\)을 생성하도록 하여, 이후의 유사도 판단이 근어에 사용되도록 하는 것이다. 표현 지시문(representation instruction)을 부착하는 프롬프트 함수를 \(P(\cdot)\)로 두면, 정책(=LLM)은 다음과 같이 근거를 생성한다.

$$r \sim \pi_{\theta} \big( \cdot \mid P(x) \big), \quad x \in \{ q, d^+, d^- \}$$

명시적 추론 흔적은 핵심 의미 특징, 개념, 관계를 표면화하여 이후 표현 추출과 보상 계산의 투명성을 높이는 핵심적인 역할을 수행한다.

2. From Rationale to Representation

생성된 명시적 추론 흔적, 근거 \(r\)을 입력 \(x\)와 결합하여 LLM에 입력시켜 인코딩하고, Masked mean pooling으로 최종 표현 \(h\)를 얻는다.

$$E = \pi_{\theta} \big( P(x) \oplus r \big)$$ $$h = \frac{1}{\vert M \vert} \displaystyle\sum_{t \in M} E_t, \quad M = \{ t: L_{\text{sys}} < t \leq L, \text{mask}_t = 1 \}, \quad h \in \mathbb{R}^d$$

• \(E\): 마지막 레이어의 히든 스테이트 행렬
• \(L\): 시퀀스 길이
• \(d\): 차원수
• \(M\): 시스템/지시문 토큰을 제외하고 유효 토큰만 남긴 마스크된 인덱스 집합
• \(L_{\text{sys}}\): 시스템 프롬프트에 해당하는 앞부분 길이
• \(h\): Masked mean pooling으로 얻은 최종 임베딩 벡터

근거 \(r\)에 정박된 표현 추출은 의미 정보에 집중하면서 지시문/형식적 토큰의 잡음을 배제하여, 더 신뢰도 높은 유사도 판단을 가능하게 함.

3. Constrastive Rewards as Policy Guidance

본 프레임워크는 대조 신호를 손실이 아닌 정책 보상으로 재해석하여, 정책경사 기반 최적화를 수행하도록 설계되어 있음이다. 설명의 편의상 GRPO를 채택하지만, Policy Gradient 알고리즘에 일반적으로 적용 가능하도록 구성되어 있다.

3.1. Rollout Strategy

탐색 다양성과 계산 효율을 절충하기 위해 비대칭 롤아웃을 사용한다. 각 학습 인스턴스 \((q_i, d^+_i, d^-_i)\)에 대해 텍스트 종류에 따라 다음과 같이 생성한다.

$$y_{q_i} \sim \pi_{\theta}\!\big(\,\cdot \mid \mathcal{P}(q_i)\big),\qquad y_{d_i^-} \sim \pi_{\theta}\!\big(\,\cdot \mid \mathcal{P}(d_i^-)\big),\qquad \{\,y_{d_i^+}^{(k)}\,\}_{k=1}^{K} \sim \pi_{\theta}\!\big(\,\cdot \mid \mathcal{P}(d_i^+)\big)$$

• \(y_{q_i}, y_{d_i^-}\): 쿼리와 음성 문서를 1회 샘플하여 기준 표현을 형성함
• \(\big\{ y_{d_i^+}^{(k)} \big\}_{k=1}^K\): 양성 문서는 \(K\)회 확률적 샘플로 다양한 해석 관점을 탐색함
• \(K\): 양성 문서에 대한 롤아웃 개수(탐색 강도)

이로써 동일 문서의 다중 근거 간 일관성/다양성을 모두 고려한 보상 계산이 가능해지고, GRPO 내 어드밴티지 추정의 기반이 된다.

3.2. Reward Design

대조학습 목표를 정책 지침으로 변환한 합성 보상을 설계한다. 쿼리 \(q_i\)의 음성 집합을 \(D_i^- = \{ d_{i, m}^- \}_{m=1}^{M_i}\)라 하면, 양성 문서의 \(k\)번째 롤아웃에 대해 다음의 보상을 계산한다.

[Contrastive Reward] 쿼리-양성 문서간의 정렬을 높이고 음성 문서와의 상관을 낮춤

$$\mathcal{R}^{(i,k)}_{\mathrm{CL}} = \operatorname{sim}\!\big(h_{q_i},\, h^{(k)}_{d^{+}_i}\big) - \sum_{m=1}^{M_i} \operatorname{sim}\!\big(h_{q_i},\, h_{d^{-}_{i,m}}\big)$$

• \(\text{sim}(\cdot, \cdot)\): 유사도
• \(h_{q_i}, h_{d_i^+}^{(k)}, h_{d_{i, m}^-}\): §2의 마스크드 풀링으로 얻은 임베딩
• \(M_i\): 쿼리 \(q_i\)에 대응하는 음성 문서 수

[Consistency Reward] 동일 양성의 다중 롤아웃 간 표현 정합성 강화

$$\mathcal{R}^{(i,k)}_{\mathrm{consist}} = \frac{1}{K-1}\sum_{j\ne k}^{K} \operatorname{sim}\!\big(h^{(k)}_{d^{+}_i},\, h^{(j)}_{d^{+}_i}\big)$$

• \(K\): 양성 문서에 대한 롤아웃 개수(탐색 강도)

[Hard-negative Reward] 배치 내 다른 예시의 가장 헷갈리는 양성 문서를 억제(= distractor)

$$\mathcal{R}^{(i)}_{\mathrm{hard}} = -\,\frac{1}{B-1}\sum_{\substack{j=1\\ j\ne i}}^{B} \max_{1\le l\le K}\; \text{sim}\!\big(h_{q_i},\, h^{(l)}_{d^{+}_j}\big)$$

• \(B\): 배치 크기

in-batch에서 다른 샘플의 양성 문서에 대한 \(K\)개의 롤아웃 중 가장 유사도가 큰 것을 선택하고, 그 값은 패널티로 주는 것이다.

최종적으로 이 Reward들을 가중합하여 최종 보상을 정의한다.

$$\mathcal{R}^{(i,k)}_{\mathrm{total}} = \mathcal{R}^{(i,k)}_{\mathrm{CL}} + \lambda_{1}\,\mathcal{R}^{(i,k)}_{\mathrm{consist}} + \lambda_{2}\,\mathcal{R}^{(i)}_{\mathrm{hard}}, \quad \quad \widehat{R}^{(i,k)}_{\mathrm{total}} = \frac{R^{(i,k)}_{\mathrm{total}}}{\tau}$$

4. Policy Optimization Objective

GRPO 절차에 따라, 동일 쿼리 \(i\)의 양성 롤아웃 \(K\)개에 대한 그룹 기준선(baseline)을 사용해 어드밴티지를 산출하고, 이 어드밴티지로 정책 로그확률을 가중하여 목표함수를 최적화하는 과정을 정의한다.

GRPO 프레임워크를 따르되, 그룹 기준선 대비 어드밴티지를 계산하되 표준편차 정규화는 제거한다.

$$A^{(i,k)}= \mathcal{R}^{(i,k)}_{\mathrm{final}} - \frac{1}{K}\sum_{l=1}^{K} \mathcal{R}^{(i,l)}_{\mathrm{final}}$$

• \(\mathcal{R}_{\text{final}}^{(i, k)}\): 실제 어드밴티지 산출에 사용하는 최종 보상
• \(A^{(i, k)}\): 그룹 평균 보상을 기준선으로 한 어드밴티지

최종적으로 정책은 advantage-weighed likelihood를 최적화하는 방향으로 진행된다.

$$L_{\mathrm{total}}= -\,\mathbb{E}_{(q,d^{+},d^{-})\sim\mathcal{D}}\Bigg[  \sum_{i=1}^{B}\sum_{k=1}^{K}  A^{(i,k)}\,  \log \pi_{\theta}\!\big(    y^{(k)}_{d^{+}_{i}} \,\big|\, \mathcal{P}(d^{+}_{i})  \big)\Bigg]$$

• \(\mathcal{D}\): 학습 데이터

5. Unsupervised Learning Extension

명시적 쿼리–문서 페어가 없는 비지도 환경에서도, 같은 텍스트의 서로 다른 해석(롤아웃) 간 자기 정합(self-alignment) 과 배치 내 하드 네거티브를 활용해 동일한 정책 최적화 틀을 적용한다.

비지도 설정에서, 텍스트 배치 \(B = \{ x_i \}_{i=1}^B\)에 대해 \(x_i\)는 앵커 해석과 \(K\)개의 롤아웃을 생성한다.

$$y^{\text{anchor}}_{x_i} \sim \pi_{\theta}\!\big(\,\cdot \mid \mathcal{P}(x_i)\big), \qquad\{\,y^{(k)}_{x_i}\,\}_{k=1}^{K} \sim \pi_{\theta}\!\big(\,\cdot \mid \mathcal{P}(x_i)\big)$$

• \(y_{x_i}^{\text{anchor}}\): 입력과 같은 텍스트의 앵커 해석(SimCSE의 same-sentence positive에 상응)

여기서 앵커는 같은 입력 텍스트에 대해 정책이 한 번 생성한 기준용 근거 시퀀스를 의미한다. 이 앵커를 기준으로, 동일한 입력 텍스트에서 추가로 샘플링한 \(K\)개의 롤아웃과의 자기 정합을 보상으로 선택한다는 것이다. 요지는 “같은 텍스트에서 나온 표현들끼리는 가깝게” 만듦으로써 정답 레이블 없이도 안정적인 표현 학습을 유도하는 것이다. 앵커 표현 \(y_{x_i}^{\text{anchor}}\)과 롤아웃 표현 \(h_{x_i}^{(h)}\)간의 자기 정합 보상은 다음과 같다.

$$R^{(i,k)}_{\text{self}}= \operatorname{sim}\!\big( \,h^{\text{anchor}}_{x_i},\, h^{(k)}_{x_i}\,\big)$$

그 외 동일 텍스트 내 일관성 보상과 배치 내 하드 네거티브 보상은 §3.2의 정의와 동일하며(식 생략), 가중 결합은 지도 설정과 동일한 형태로 구성된다.

Experiments

Main Results

GRACE는 4개 디코더 백본 전반에서 MTEB 평균을 +11.52% 끌어올렸고, 특히 Retrieval과 Pair Classification에서 큰 폭의 개선을 보였으며, Classification/Clustering/STS/Summarization도 고르게 향상되었다. 명시적 근거+대조 보상이 결합된 정책 최적화는, 단순 대조 미세튜닝과 달리 생성 능력을 보존하면서 임베딩 품질을 광범위하게 향상시키는 것으로 나타났다.

Ablation Study

\(\lambda_1\)은 Consistency 가중치이고, \(\lambda_2\)은 하드 네거티브 가중치이다. 둘 다 0일 때는 구조화된 보상을 제거했을때로(consistency, hard) 성능이 안좋게 나온 것을 볼 수 있다. 하드 네거티브 판별력이 전체 성능을 좌우하는 핵심 드라이버이며, 일관성 보상은 보조적으로 양성 롤아웃 간 표현 정합성을 높여 안정성을 높이는 것을 알 수 있다.

Alternative RL Algorithm

동일 프레임워크를 여러 RL 알고리즘(예: ReMax, REINFORCE++, DAPO, GRPO 등)에 적용해도 일관된 성능 개선을 확인. 즉, GRACE의 개선은 특정 옵티마이저에 종속되지 않는다.

Efficiency

Figure 4는 임베딩 접근들의 품질–지연 절충을 비교한 결과로, 생성형 파이프라인(Base·GRACE)은 지연이 주로 생성 단계 $\text{T_gen}$에 의해 결정되며 디코딩 길이를 256에서 512로 늘리면 성능 gain은 체감하지만 지연(latency)은 크게 늘어나므로, 논문에서는 최종적으로 256을 선택하였다. 같은 디코딩 길이에서 GRACE는 base대비 동일 지연에서 더 높은 정확도를 달성하였다.

Rerpresentation 방법에 따른 효과

Figure 6는 표현 추출 방식에 따라 임베딩 품질이 크게 달라지며, Mean Pooling이 일관적으로 우수하다. 네 가지 GRACE 변형(GRACE-1.5B, GRACE-3B(Qwen), GRACE-3B(LLaMA), GRACE-4B)에 대해 감독 학습(a)과 비지도 학습(b) 모두에서 Last Layer(LL)와 Penultimate Layer(PL, L-1번째 레이어)의 Mean Pooling이 EOS 토큰 사용이나 Max Pooling보다 항상 높은 점수를 보인다.

LL과 PL의 성능 차이는 미미해 두 층 모두 안정적이며, 이는 정보가 특정 토큰(EOS)에 집중되기보다 여러 토큰에 분산되어 있어 평균 풀링이 문장 의미를 더 견고하게 집약한다고 해석한다.

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