Matryoshka Representation Learning

Embedding의 dimension을 줄이면 저장과 dot-product 비용은 줄지만, 아무 vector나 앞부분만 남겨도 의미가 유지되는 것은 아니다. MRL은 학습할 때부터 작은 prefix만으로도 task를 풀도록 요구한다. 한 모델의 vector 안에 여러 크기의 유용한 표현을 중첩해 두고, 배포 환경에 맞는 정확도–비용 operating point를 선택하는 방법이다.

1. Full vector encode
2. trained prefix slice
3. re-normalize
4. small-prefix search
5. larger-prefix rerank

Nested Prefixes

MRL은 하나의 full2048d forward에서 8/16/32/…/1024/2048 같은 nested prefixes에 loss를 준다. 256d만 남기면 raw storage/dot-product coordinates는 2048d의 $\frac18$이다. 재인코딩 없이 slice하지만 prefix를 re-normalize하고 dimension별 ANN index와 task quality를 검증해야 한다. 학습하지 않은 cut point의 품질은 보장하지 않는다.

Cross-Model Comparisons

OpenAI2024 보고의 3-large256d가 ada0021536d보다 MTEB 평균이 높다는 것은 cross-model 비교다. 같은 모델을 줄여 성능이 높아졌다는 뜻이 아니다.

Float32 1M vectors의 coordinates만 1536d≈6.1GB,256d≈1.0GB이며 ANN graph/metadata/replication은 별도다. 강의 당시 API option 설명은 현재 제품 안내와 구분한다.

Adaptive Retrieval

Adaptive retrieval 예시: full2048d는 보관하고 16d prefix ANN으로 1Mdocs→K200(0.02%)→2048d prefix로 rerank한다. Dimension 증가이지 float32→float64의 precision 변화가 아니다. API에서 256d만 받으면 나머지 coordinates는 되살릴 수 없다. 1단계 recall이 cascade 전체 상한이며 2단계도 vector similarity라 Cross-Encoder와 다르다.

Why the Prefix Becomes Useful

하나의 2048d vector $z$를 얻은 뒤 256d, 512d, 1024d, 2048d prefix에 각각 loss를 준다고 하자. Full vector만으로 답을 맞추면 짧은 prefix의 loss를 줄일 수 없으므로, 앞부분만으로도 유용한 정보를 담아야 한다. 여러 크기의 목표가 겹쳐서 학습되는 것이 일반적인 사후 truncation과 다른 점이다. 슬라이드에서는 더 촘촘한 8/16/32/… cut points도 예로 든다.

Slice는 두 번째 encoder forward를 필요로 하지 않지만, 처음의 full-vector encoding 비용까지 없어지는 것은 아니다. 절약하는 것은 주로 보관하는 coordinates와 이후 similarity 계산이다. Full vector가 unit norm이어도 일부 좌표만 남긴 prefix의 norm은 일반적으로 1이 아니므로, 비교할 prefix를 별도로 Normalize해야 한다. 여기서의 정규화는 coordinate의 단순 합을 1로 만드는 과정과 혼동하지 않는다.

Two Different Cost Decisions

256d만 보관하면 2048d 대비 coordinates는 1/8이지만, adaptive retrieval에서 나중에 2048d로 rerank하려면 full vector를 계속 보관해야 한다. 작은 prefix를 index로 쓰는 것이 전체 저장도 1/8로 만든다는 뜻은 아니다. API에서 256d만 받은 뒤 버린 coordinates를 복원할 수도 없다.

1M documents에서 16d로 200개를 먼저 찾으면 후보는 전체의 0.02%다. 이 200개에 2048d similarity를 적용하면 모든 문서를 큰 vector로 비교하는 비용을 피할 수 있다. 그러나 작은 prefix가 놓친 정답은 큰 prefix로 다시 볼 기회가 없으므로, $D_s$, $K$, $D_r$를 corpus에 맞추어 조정해야 한다. 두 단계 모두 vector similarity이며, query와 document tokens를 함께 읽는 Cross-Encoder reranking과는 다른 방식이다.

$$\text{Raw vector bytes}=N\times d\times\mathrm{bytes\_per\_value}$$ $$2048d\longrightarrow256d:\quad\frac18\text{ coordinates}$$ $$\operatorname{Normalize}(z_{1:D_s}),\quad\operatorname{Normalize}(z_{1:D_r})$$ $$D_r>D_s\quad\text{(각 prefix를 별도로 Normalize)}$$

Dimension reduction은 dtype precision 변경이 아니다. Full vector가 없어지면 더 큰 prefix를 나중에 복원할 수 없다.

Trained prefix를 slice한 뒤 다시 normalize하고 dimension별 index를 쓴다.

Check

MRL vector를 잘라 바로 dot product하면 안 되는 이유와 추가 확인 항목은?

Answer

Slice 후 norm이 달라져 re-normalize해야 한다. dimension별 ANN index를 만들고 해당 prefix의 downstream quality를 검증한다.

Adaptive retrieval의 큰 prefix rerank가 Cross-Encoder와 같지 않은 이유는?

Answer

같은 embedding의 더 많은 coordinates로 similarity를 다시 계산할 뿐 $q$,$d$ token-level joint interaction은 없다. 첫 후보에 없는 정답은 복구하지 못한다.

MRL vector를 짧게 slice하면 Encoder forward 비용도 같은 비율로 줄어드는가?

Answer

그렇게 단정할 수 없다. Full vector를 encode한 뒤 prefix를 쓰는 설계에서는 encoding 비용이 남고, 저장·검색의 dimension 비용이 줄어든다. 학습된 prefix를 선택하고 다시 normalize하며 해당 dimension의 index를 사용한다.

하나의 full vector에 nested prefix별 loss를 주고 배포 때 잘라 재사용하는 학습 방식은?

Answer

Matryoshka Representation Learning (MRL). Slice 후 re-normalize와 dimension별 검증이 필요하다.

MRL의 작은 prefix 검색 후 더 큰 prefix rerank 과정을 말하라.

Answer

Full vector 보관 → small prefix를 normalize해 ANN → top-$K$ → larger prefix도 따로 normalize해 similarity rerank. 놓친 후보는 되살리지 못한다.

Related

• MRL Evaluation
• Text Embedding
• Contrastive Training