LLM-JEPA 논문 리뷰 - 토큰 생성과 표현 예측을 함께 학습하기

LLM-JEPA는 Transformer를 새로 설계하지 않는다. 기존 Next Token Prediction(NTP)에, 같은 의미를 가진 두 표현의 hidden representation을 맞추는 보조 loss를 더한다.

자연어 요구사항과 정규식, 질문과 정답처럼 의미가 대응하는 pair가 있을 때 효과가 컸다. 반면 이 pair가 없는 일반 말뭉치에서는 별도의 view를 만들어야 하고, 학습 계산량도 늘어난다. 범용 대체재로 부르기에는 아직 실험 범위가 좁다.

두 학습 신호를 함께 쓴다

논문은 같은 내용을 서로 다른 형식으로 표현한 두 view를 Text와 Code라고 부른다.

과제	Text	Code
NL→Regex	정규식 요구사항	정답 정규식
Text→SQL	DB 스키마와 자연어 질의	정답 SQL
GSM8K	수학 문제	풀이와 정답
NQ-Open	질문	짧은 정답

Text와 Code는 개념적인 이름이다. 실제 구현에서는 모델별 chat template를 사용한다.

전체 loss는 token 생성 loss와 표현 정렬 loss의 합이다.

\[\mathcal L_{\mathrm{LLM-JEPA}} =\gamma\mathcal L_{\mathrm{NTP}} +\lambda\mathcal L_{\mathrm{JEPA}}\] \[\mathcal L_{\mathrm{JEPA}} =\mathbb E\left[1-\cos(z_{\mathrm{pred}},z_{\mathrm{target}})\right]\]

• $\mathcal L_{\mathrm{NTP}}$: 정답 output token에 대한 cross-entropy
• $z_{\mathrm{pred}}$: Text 뒤에 predictor token을 붙여 얻은 마지막 hidden state
• $z_{\mathrm{target}}$: Code 전체를 읽은 뒤의 마지막 hidden state
• $\lambda$: JEPA loss의 비중

두 vector의 방향이 같아질수록 JEPA loss는 0에 가까워진다.

NTP:  Text ───────────────→ Code token 생성

JEPA: Text + [PRED] × k ─→ z_pred
      Code 단독 입력 ────→ z_target
                           └─ cosine distance

[PRED]는 생성할 정답이 아니다

[PRED]는 Text 표현을 더 처리하기 위한 학습용 special token이다. 그 자리에 정답 token을 넣거나 [PRED] 자체를 출력하도록 학습하지 않는다.

\[z_{\mathrm{pred}} =\operatorname{LastHidden}(\operatorname{LLM}(\mathrm{Text}+[\mathrm{PRED}]^k))\]

k=0이면 Text의 마지막 hidden state를 그대로 쓴다. 논문은 $k=0$부터 $4$까지 탐색했지만 모든 모델과 과제에 통하는 최적값은 찾지 못했다.

I-JEPA와의 차이도 있다. 별도의 target encoder나 EMA teacher를 두지 않고, 학습 중인 LLM 하나를 양쪽 encoder와 predictor로 사용한다. Target 쪽 gradient도 끊지 않는다.

Block attention mask가 필요한 이유

Text와 Code를 한 sequence로 그대로 이어 붙이면 뒤쪽 Code 표현이 앞쪽 Text를 볼 수 있다. 그러면 독립적인 target representation이 아니다.

LLM-JEPA는 두 block을 한 context에 넣되 서로 attention하지 못하도록 가린다.

Text block → 자기 block의 이전 token만 참조
Code block → 자기 block의 이전 token만 참조
Text ↔ Code → 참조 금지

최신 공개 구현은 stream을 batch 축으로 묶어 모델 호출 횟수를 줄인다. 그래도 처리해야 할 token은 남으므로 같은 step의 계산량은 표준 NTP의 약 두 배다.

어디에서 효과가 컸나

주요 결과는 5개 seed 평균이다.

설정	NTP	LLM-JEPA	차이
Llama-3.2-1B, NL-RX-SYNTH	57.29	71.46	+14.17%p
Gemma2-2B, NL-RX-SYNTH	33.65	43.12	+9.47%p
Llama-3.2-1B, GSM8K	32.36	36.36	+4.00%p
Llama-3.2-1B, Spider	47.52	50.55	+3.03%p
Llama-3.2-1B, NQ-Open	20.12	21.59	+1.47%p
Llama-3.2-1B, HellaSwag	69.40	70.51	+1.11%p

정규식처럼 두 view의 대응이 분명하고 baseline이 약한 과제에서는 차이가 컸다. QA와 reasoning에서는 대체로 작았다. JEPA loss만 남기고 NTP 비중을 0으로 만들면 모델이 빈 문자열만 생성해 정확도가 0%가 됐다. 표현 정렬이 token 생성을 대신할 수 없다는 뜻이다.

일반 말뭉치에는 pair가 필요하다

NTP는 임의의 token sequence만 있으면 학습할 수 있다. LLM-JEPA는 같은 의미를 가진 두 view가 필요하다.

논문은 같은 의미의 문장 다섯 개를 묶은 paraphrase dataset에서 문장 $i$로 문장 $i+1$의 표현을 예측했다. 이후 감성 분류에 일반 NTP fine-tuning을 적용했을 때도 작은 향상이 남았다. 다만 일반 corpus 전체에 적용하려면 의미를 보존하는 paraphrase를 대량으로 만들거나 원래부터 pair가 있는 데이터를 확보해야 한다.

JEPA loss를 batch의 일부에서만 계산하는 dropout도 제안한다.

\[\text{relative compute}=2-\alpha\]

$\alpha=0.75$라면 batch의 25%에서만 JEPA 경로를 계산해 평균 비용이 약 $1.25\times$가 된다. Task마다 $k$와 $\lambda$를 찾는 탐색 비용은 이 계산에 포함되지 않는다.

표현이 달라졌다는 것과 생성이 좋아졌다는 것은 다르다

학습 뒤 Enc(Text) - Enc(Code)의 singular value가 작아지고, 두 view 사이의 구조가 더 선형적으로 정렬되는 현상이 관찰됐다. JEPA loss가 hidden geometry를 바꿨다는 근거다.

하지만 geometry 변화가 정확도 향상의 원인이라는 증명은 아니다. 후속 연구는 비슷한 표현 제약이 hidden space를 바꿔도 decoded exact-match를 안정적으로 높이지는 못할 수 있다고 보고했다. 원 논문의 모든 설정을 직접 반박한 재현 연구는 아니지만, 표현 지표만 보고 생성 품질까지 일반화하면 안 된다는 경고로 읽을 수 있다.

LLM-JEPA가 잘 맞는 곳은 의미가 대응하는 pair가 이미 있고, 데이터가 적으며, 추가 학습 비용을 감당할 수 있는 구조화 과제다. 추론 비용은 늘지 않지만 학습 데이터와 hyperparameter 조건은 NTP보다 까다롭다.

참고 자료

• LLM-JEPA 논문
• LLM-JEPA 코드
• Representation Without Reward