Behrouz, A., Li, Z., Kacham, P., Daliri, M., Deng, Y., Zhong, P., Razaviyayn, M., & Mirrokni, V. (2025). Atlas: Learning to optimally memorize the context at test time. arXiv. https://arxiv.org/abs/2505.23735

이전편을 읽고오자

4. DeepTransformers: Transformers with Deep Memory

변형한 메모리 용량과 Omega rule이 실제 Transformer에 어떤 영향을 미치는지를 설명

4.1 Online and Local Context Optimization of Memory

Sliding Window Attention과 연결

Softmax block은 Nadaraya-Waston estimator을 이용해서 non-parametric solution으로 re-formulated 될 수 있다.

• $M^* = arg min_M \displaystyle \sum^L_{i=1}s(k_i, q) ||v_i -M||^2_2 = {{\sum^L_{i=1} s(k_i, q)v_i}\over{\sum^L_{j=1}s(k_j,q)}}$

이를 다시 과거 c개의 토큰으로 제한하면

• $M^* = arg min_M \displaystyle \sum^t_{i=t-c+1}s(k_i, q) ||v_i -M||^2_2 = {{\sum^L_{i=t-c+1} s(k_i, q)v_i}\over{\sum^L_{j=t-c+1}s(k_j,q)}}$
• 이는 Sliding Window Attention과 동일함
• Attention은 non-parametric한 방법으로 global한 최적화를 수행
  • 최신 recurrent model은 online-learning 방식
    • non-parametric: 정보를 압축하지 않고 데이터를 그대로 활용
    • online learning: 새로운 토큰에 대해 직전 요약된 메모리만 참조
  • Sliding window RNN과 Omega rule은 과거 토큰의 일정한 context window에 대해 메모리를 최적화해서 context 자체를 기억하게 만듦

Deep Linear Attention

linear attention의 행렬 메모리를 MLP로 대체

• deep memory 모듈로 메모리를 graident descent로 최적화
• $M_t = \alpha_t M_{t-1} - \eta_t ∇l(M_{t-1}; \Phi(k_t),v_t)$
  • $l(M_{t-1}; \Phi(k_t), v_t) = <M_{t-1}(\Phi(k_t)),v_t>$
• chunk 단위 병렬화로 효율적으로 사용할 수 있음

Sliding Window Linear Attention

• Associative memory에서 linear attention에 따라 dot product similarity를 attnetional bias로 사용
• $l(M_t;k_i,v_i) = <M_t(k_i),v_i>$
• gradient descent로 최적화
  • $M_t = \alpha_t M_{t-1} - \eta_t ∇ \displaystyle \sum^t_{i=t-c+1} l(M_{t-1}; \Phi(k_i), v_i)$
  • $M_t = M_{t-1} + \displaystyle \sum^t_{i=t-c+1}\gamma^{(t)}_i v_i \Phi(k_i)^\top$
  • c = 1, $\Phi(x) = x$일때 기존 linear attention과 동일

4.2 Memory Capacity and Exponential Kernels

기존 Transformer의 softmax attention

• $y_i = {{1}\over{\sum^i_{l=1}exp(q_i^\top k_l / \sqrt{d_{in}}}}\displaystyle \sum^i_{j=1}exp(q^\top_i k_j / \sqrt{d_{in}})v_j$
  • 이때 exp는 분리되지 않기에 recurrent한 표현이 어려움
• exp는 $\Phi_p$와 비교했을 때 입력을 무한차원 공간으로 매핑하는 feature map으로 볼 수 있음
  • $\Phi^*(x) = \begin{pmatrix} 1 \ {{x}\over{\sqrt 1}} \ {{x^{\otimes2}}\over{\sqrt{2!}}} \ . \ .\end{pmatrix}$
  • $\Phi_p(x) = x^{\otimes p}$
    • $x^{\otimes p}$는 Kronecker 곱을 이용한 self-tensoring 연산
    • $exp(q_t^\top k_t) = \Phi^* (q_t)^\top \Phi^* (k_t)$
• 최종적으로 $y_i = {{1}\over{\sum^i_{l=1} exp(q^\top_ik_l)}} \displaystyle \sum^i_{j=1} v_j \Phi^(k_j)^\top \Phi^(q_i)$ $={{1}\over{\sum^i_{l=1} exp(q^\top_i k_l)}}(\displaystyle \sum^i_{j=1} \Phi^(v_jk_j)^{\top})\Phi^(q_i) = M_i \Phi^* (q_i)$

이를 활용해 Deep Linear Attention에서 $\Phi^*$를 활용해 DeepTransformers 정의

• $M_t = M_{t-1} - ∇<M_{t-1}(\Phi^*(k_t)),v_t>$
• 선형인 경우: $M_t = M_{t-1} - ∇<M_{t-1}\Phi^(k_t), v_t> = M_{t-1} + v_t\Phi^(k_t)^\top$ $=\displaystyle \sum^t_{i=1}v_i\Phi^*(k_i)^\top$
• Output: $y_t = M_t\Phi^*(q_t) = \displaystyle\sum^t_{i=1} v_i exp(q_i^\top k_i)$

4.3 Deep Omega Tranformer(DoT)

위의 DeepTransformer에서 Hebbian Rule대신 Omega learning Rule 사용

• $M_t = M_{t-1} - ∇ \displaystyle \sum^t_{i = t-c+1} \gamma^{(t)}_i ||M(\Phi^*(k_i)) - v_i||^2_2$

Linear memory: unnormalized된 형태를 얻음

• $M_t = (I- \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}i \Phi^(k_i) \Phi^(k_i)^\top)M{t-1} - \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i v_i \Phi^*(k_i)^\top$

• 최종 출력: $y_t = M_t\Phi^*(q_t)$

$= (I-\displaystyle\sum^t_{i=t-c+1} \gamma^{(t)}i \Phi^(k_i) \Phi^(k_i)^\top) M{t-1}\Phi^*(q_t) - \displaystyle\sum^t_{i=t-c+1} \gamma^{(t)}_i v_i exp(q_t^\top k_i)$

Online Case: c = 1

• $M_t = (I-\eta_t \Phi^(k_t)\Phi^(k_t)^\top)M_{t-1} - \eta_t v_t \Phi^*(k_t)^\top$
• 출력: $y_t = M_t\Phi^(q_t)$ $= (I-\eta_t \Phi^(k_t)exp(q_t^\top k_t)) M_{t-1} - \eta_t v_t exp(q_t^\top k_t)$
• feature map을 사용하기에 이전 메모리 상태를 이용해 value를 replace하는 기능도 가능함

5. ATLAS: A Locally Optima lMemory with High Capacity

Omega Rule

• 모델이 context를 기억하도록 만듦 + feature mapping의 메모리 용량 증가
• gradient descent에 의존 $\rightarrow$ local optimum에 머무르게 함

Muon optimizer

• 2차 정보를 근사 + 행렬 곱 연산으로 구성되어 병렬화 가능
• $M_t = \alpha_t M_{t-1} - \eta_t NewtonShulz-k(S_t)$
• $S_t = \theta_t S_{t-1} + ∇\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i ||M(\Phi^*(k_i)) - v_i||^2_2$
  • c: local cantext 길이
  • k: Newton-Schulz 반복 횟수
• 성질: $k \rightarrow \infty$일때 $NewtonShulz-k(S_t)$는 $S_t$에 가장 가까운 semi-orthogonal 행렬로 수렴
  • 2차 정보를 더 정확하게 반영
  • 즉 k가 크면 더 정확한 메모리 업데이트, 작은 k는 더 빠른 계산

5.1 Parallel Training

Atlas의 병렬화 방법

• c = 1로 가정
• 시퀀스를 여러 개의 청크로 나눔
• 청크의 마지막 상태를 기준으로 gradient를 계산

Muon 제외, momentm만 사용

• $M_t = \alpha_t M_{t-1} + S_t$
• $S_t = \theta_t S_{t-1} - \eta_t ∇l(M_{t'}, k_t, v_t)$

gradient 정의 변경

• $u_t = ∇l(M_{t'}, k_t, v_t)$
• $M_t = \alpha_t M_{t-1} + S_t$
• $S_t = \theta_t S_{t-1} - \eta_t u_t$
• 전개: $S_t = \theta_t S_{t-1} - \eta_t u_t$

장점

• 모든 graident $u_t$를 병렬로 계산 가능

6. Experiment

Task

• Language Modeling: 일반적인 언어 모델링에서 컨텍스트가 길어져도 기존 모델보다 우수
• Common-Sense Reasoing: Transformer 수준 or 그 이상
• Recall-Intensive Tasks: 긴 문맥에서 정보를 찾아내는 과제: Atlas가 가장 강력
• Long Context Understanding: 매우 긴 시퀀스(수백만 토큰)에서 안정적
  • 특히 BABILong 벤치마크에서 10M context에 대해 +80% 정확도 향상
• Needle-in-a-Haystack: Atlas와 OmegaNet이 가장 잘찾음

Ablation Study

• Memory Architecture: Deeop memory 성능이 크게 향상
• Feature Mapping: 고차원 feature일수록 성능 차이가 큼
• Omega Rule: online, sliding window, global과 비교할 때 context 기반 학습이 가장 좋음
• Optimizer: Muon이 가장 안정적

Behrouz, A., Li, Z., Kacham, P., Daliri, M., Deng, Y., Zhong, P., Razaviyayn, M., & Mirrokni, V. (2025). Atlas: Learning to optimally memorize the context at test time. arXiv. https://arxiv.org/abs/2505.23735

Titans 논문을 읽고 오는 것이 좋다.

Introduction(기존 모델의 한계)

Transformer(Attention module)

• associative memory처럼 작동
• 토큰 간 쌍별 의존성 -> K-V mapping, QK 유사도로 검색
  • 시간 복잡도가 O($N^2$)으로 긴 문맥에서 제한됨

Modern Recurrent Model

• attentional bias를 최적화

  • 고정 크기 메모리를 사용해서 KV cache를 관리

• 더 정교한 메모리 관리 필요

  • 학습 규칙: DeltaNet $\rightarrow$ Delta rule
    • 단순 가산: $M_t = M_{t-1} + v_t k_t^\top$
      • 새로운 정보를 더함. 메모리가 누적되고 오래된 정보 제거 불가
    • DeltaNet: $M_t = M_{t-1} + \beta_t(v_t-M_{t-1}k_t)k^\top_t$
      • 예측과 실제 값 차이를 줄이도록 업데이트
  • forget gate: RetNet의 비의존적 gating $\rightarrow$ 적응적 gating(Titans)
    • $M_t = \alpha_tM_{t-1} + update$
    • RetNet: $\alpha$가 입력과 무관하게 동일한 방식으로 감소
    • Titans: $\alpha_t$가 입력에 따라 변화
  • 메모리 구조: 벡터 기반 메모리 $\rightarrow$ deep-neural memory
    • 벡터 행렬 기반: $M_t \in \mathbb R^d$ or $M_t \in \mathbb R^{d \times d}$
      • 표현력이 부족하고 저장가능한 패턴 수 제한
    • Deep Memory Module: $M(k) = MLP(k)$
      • 추상화된 정보를 저장할 수 있지만 문맥 단위 학습 부족

• 한계

  • online update: 메모리가 현재 토큰에 대해서만 최적화
  • 제한된 메모리 용량: KV 매핑에 의해 저장 가능한 정보량 제한
  • 메모리 관리 표현력 부족: 1차 정보에 기반한 graident descent로 비효율적 mapping

Memory Perspective

Associative Memory: 서로 다른 개체, 사건을 연결하는 능력

• 메모리는 입력에 의해 갱신되는 신경적 상태
• surprise가 클수록 더 큰 영향을 미침
• 개별 토큰 중심의 surprise라는 한계점이 존재함
  • 토큰이 아닌 문맥의 surprise를 측정하는 방식

Test time memorization

• 모델의 학습된 파라미터를 업데이트 하지 않음
• 주어진 입력 문맥 내에서만 정보를 저장, 검색
• 즉, 메모리가 초기화되면 새로운 문맥에는 이전 학습이 유지되지 않음

Contribution

• 메모리 용량에 대한 이해, 개선
  • 입력에 대해 higher-order mapping을 사용해 메모리 용량 증가
• 새로운 학습 규칙: Omega rule
  • 문맥 전체를 고려해 메모리 업데이트
• Transforemr의 일반화
  • DeepTransformers를 제시해 기존 Transformer을 포함하는 더 일반적인 모델군
• 향상된 메모리 관리 구조
  • OmegaNet과 Atlas
• 다양한 작업에서 성능 향상

2. Preliminaries

Notation

• input: $x \in \mathbb R^{N \times d_{in}}$
• 시간 t에서의 메모리 상태: $M_t$
• key, value, query: $K, V, Q$, 각 벡터 $k_t, v_t, q_t$
• attention bias: $l(M_t;k_t,v_t)$: 메모리가 무엇을 잘 기억해야 하는가
• Memory module: $L_M \geq 1$인 MLP + residual connection
• Memory module parameter: $\theta_M := {W_1, ..., W_{L_M}, ...}$

2.1 Background

Attention

입력 x에 대해

• $Q = xW_Q$, $K=xW_K$, $V=xW_V$
• causal attention: $y_i = {{\sum^i_{j=1}exp({{q^\top_ik_j}\over{\sqrt{d_{in}}}})v_j}\over{\sum^i_{l=1} exp({{q^\top_ik_l}\over{\sqrt{d_{in}}}})}}$
  • $W_Q$, $W_K$, $W_V \in \mathbb R^{d_{in} \times d_{in}}$은 학습 가능한 파라미터
  • 분모는 normalization term
• 각 토큰마다 최소 $N \times d$만큼의 연산이 필요하기에 긴 문맥에서 확장성 제한

Recurrent Models

• 장점
  • 선형 시간 학습
  • 병렬화 가능
  • Transformer와 유사한 성능
• $M_t = A_t * M_{t-1} + v_t k_t^\top$
  • $M_t \in \mathbb R^{d \times n}$: 메모리
  • $k_t,v_t \in \mathbb R^d$: 입력
  • $A_t$: 감쇠, 게이팅
  • $*$: 임의의 연산자
• 단점
  • 메모리 계속 누적

Deep Memory Module

• 메모리는 단순 저장이 아닌 학습되는 함수
• $M^* = argmin_M L(M(K); V)$
• 반복적인 알고리즘으로 수행: 메모리 업데이트 규칙
• 최적화 구조
  • inner Loop: 메모리 파라미터 $\theta_M$ 최적화
  • Outer Loop: 나머지 파라미터 학습

3. Learning to Memorize the Context at Test Time

3.1 Associataive Memory with Super Linear Capacity

모델은 서로 상관이 없는 KV 쌍을 최대 몇 개 까지 저장할 수 있는가

• matrix memory + l2 loss를 사용
  • $l(M_t; k_t, v_t) = ||M_t(k_t) - v_t||^2_2$
  • graident descent로 최적화

    Capacity of $l_2$

    Memory M: $d_v \times d_k$크기 행렬
  • linear independent key를 가지는 (k, v)쌍을 $O(d_k)$만큼 저장
  • 메모리 파라미터 수에 비해 저장 가능한 정보량이 sub-linear
    • 메모리 크기가 M이면 저장 가능한 패턴 수는 $c\times M$보다 작음

      Theorem 1: Effect of Deep Memory

      $L_M \geq 2$층의 MLP형태 메모리 M에 대해
  • input dimension: $d_k$, hidden dimension: $d_h$
  • 메모리 저장 범위: $O(d_k, d_v)에서 O(d_kd_v \displaystyle \sum^{L_M}{i=1} min{d^{(j)}_h}{j \geq i}d^{(j+1)}_h)$
    • MLP 형태 메모리가 몇 개의 서로 다른 key를 linearly independent하게 value로 매핑할 수 있는가
    • MLP는 선형이 아닌 여러 층을 지나며 비선형적으로 변형됨
    • 하한: 입력 차원 $\times$ 출력 차원
    • 상한
      • bottleneck: 가장 작은 차원의 layer가 전체 capacity를 제한함($\displaystyle min_{j \geq i}d^{(j)}_h$
      • 각 레이어는 변환 능력을 추가함: $d^{(j)}_h \times d^{(j+1)}_h$
  • 표현력과 용량이 증가하지만, 여전히 super-linear하지 않음

단순히 $d_k$의 차원을 늘리면

• 파라미터, 메모리 사용량, 계산 비용이 증가함
• Kernel Attention Perspective를 이용

Polynomial Feature Mapping 입력 x에 대해: $\Phi_p(x) = [x^{\beta}]_{|\beta| \leq p}$

• 최대 차수가 p까지의 polynomial feature로 확장 메모리 학습: $L(M(\Phi(K));V)$

  Memory Capacity with Polynomial Mapping

  다항 매핑 $\Phi_p$를 사용하면 행렬 메모리는 최대 $O(d^p_k)$개의 (k, v)쌍 저장
• 차수 p에 따라 용량이 급격히 증가

attentional bias를 gradient descent로 최적화할 때 메모리 업데이트

• 내적 기반: $l^{(1)}(M_t; k_t, v_t) = <M_tk_t, v_t>$
• $l2$오차 기반: $l^{(2)}(M_t; k_t, v_t) = ||M_t \Phi (k_t) - v_t||^2_2$ 이를 gradient descent로 최적화

Hebbian Rule: $M_t = M_{t-1} + \eta_tv_t\Phi(k_t)^\top$

• 새로운 정보 $(k_t,v_t)$를 그대로 메모리에 추가
• $v_t\Phi(k_t)^T$: key-value관계를 outer product로 저장
• 같이 나타난 것들은 연결해서 기억

Delta Rule: $M_t = (I-\eta_t\Phi(k_t)\Phi(k_t)^\top)M_{t-1} + \eta_t v_t \Phi(k_t)^\top$

• 1항: 기존 메모리 수정(망각도 포함)
• 2항: 새로운 정보 추가

Hebbian Rule의 Kernel Attention Perspective Attention의 exponential kernel: $exp(q^\top_ik_j)$

• 테일러 전개: $exp(q_i^\top k_j) \approx 1+q_i^\top k_j + {{(q^T_ik_j)^2}\over{2!}} + {{(q^\top_i k_j)^3}\over{3!}}+...$
• 일반화: $exp(q^\top_ik_j) \approx \Phi_p(q_i)^\top\Phi_p(k_j)$
• 다항 매핑은 softmax attention의 근사

해석 기존 attention의 exponential은 비선형이고 분리가 불가능함

• $exp(q^\top_ik_j) \neq f(q_i) \cdot g(k_j)$
• 이로 인해 attention의 time complexity가 $O(N^2)$이 됨
• 이를 $\Phi$로 이루어진 두 식으로 변형하면서 고차원 공간에서의 단순 내적으로 변형시킴
• 즉 attention은 고차원 feature 공간에서의 내적임

Input gating $\Phi_p(x) = a_0 + a_1x + a_2x^2 +... + a_px^p$

• $a_i$는 특정 차수의 특징을 켜거나 끔

3.2 Long-term Memory with Context Memorization

기존의 Recurrent Model은 현재 입력에 대해서만 attentional bias를 최적화하고 이전 메모리 상태는 유지함

• $min_M l(M;k_t,v_t)+Ret_t(M,M_{t-1})$
  • 이때 Ret은 보존 게이트
  • 계산이 간단하고 빠르지만 문맥 전체를 제대로 기억하지 못함
• 매 시점마다 전체 입력 시퀀스에 대해 최적화
  • $min_M \displaystyle\sum^t_{i=1}l(M;k_i,v_i)$
    • 모든 과거 토큰을 보기에 비용 증가

Sliding window recurrent model $min_M \displaystyle \sum^t_{i=t-c+1}\gamma ^{(t)}_i l(M;k_i, v_i)$

• c: 문맥 길이
• $\gamma^{(t)}_i$: 감쇠 게수

OmegaNet

Omega Rule

• $min_M \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i ||M(k_i) -v_i||^2_2$
  • 문맥 전체 기준 업데이트 c = 1일때 기존 Delta Rule

OmegaNet

• Update: $M_t = \alpha_tM_{t-1} - ∇\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i ||M(\Phi(k_i)) - v_i||^2_2$
• 선형$M_t = (diag(a_t) - \displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}i \Phi (k_i) \Phi(k_i)^\top )M{t-1} -\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i v_i \Phi(k_i)^\top$
  • $diag(a_t)$: 기존 메모리 유지 정도
  • $\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i \Phi (k_i) \Phi(k_i)^\top$: 특정 방향의 정보를 제거. 즉 현재 문맥과 충돌하는 기억을 지움
  • $-\displaystyle \sum^t_{i=t-c+1} \gamma^{(t)}_i v_i \Phi(k_i)^\top$: 새로운 key-value관계를 slide로 추가

Gradient Descent외의 예시

최적화 알고리즘

• c = 1, momentum: $S_t = \theta_t S_{t-1} - \eta_t ∇l(M_{t-1};k_t,v_t)$: Titans
• c = context length: $M_t=arg min_M \displaystyle \sum^t_{i=1} ||MK_i - v_i||^2_2$
• 이런 케이스들은 병렬화가 어렵고 계산 비용이 큼

3.3 Parallelizing Omega Rule

원래 Omega Rule

• c개의 gradient $∇l \in \mathbb R^{d_{in} \times d_{in}}$을 모두 계산해야 함
  • 메모리 사용량과 I/O 비용 모두 증가

청크 분할

• 입력 시퀀스 길이 L을 $b \ge 1$인 청크로 나눔
  • $S_i = {x_{(i-1)b+1}, ..., x_{ib}}$
    • 이전 청크의 마지막 메모리 상태를 기준으로 gradient 계산
• $\gamma^{(t)}_i = \eta_t$라고 가정
  • b = 1일때(청크 x): $M_t = \alpha_tM_{t-1} - \eta_t \displaystyle \sum^t_{i=t-c+1} ∇l(M_{t-1}; k_i,v_i)$
  • $b \ge 1$
    • $t' = t - mod(t,b)$, t'은 현재 청크 시작 지점
    • $M_t = \alpha_t ... \alpha_{t'}M_{t'} - \displaystyle \sum^t_{n=t'}(\alpha_t ... \alpha_{n+1})\eta_n \displaystyle^n_{i=n-c+1} ∇l(M_{t'}; k_i, v_i)$
  • graident계산시 sliding window mask인 $M_s$ 사용
  • 효과
    • 메모리 사용량 감소
    • 병렬화 가능
    • 계산 효율 유지

길어져서 렉걸리는관계로 다음 게시글로

사전 세팅

arag/scripts/build_index.py

function: semantic search용 embedding index를 미리 만드는 스크립트

• chunks.json을 입력받음
• load_chunks: 파일을 입력받아서 dictionary형태로 반환
  • {"id": ..., "text": ...}형태로 전부 반환
• split_sentences: chunk를 문장 단위(.!?\n)로 split
  • 이때 빈 문자열과 길이가 10자 이하인 짧은 조각은 제거함

    build_index

• load_chunk로 청크단위로 만듦
• split_sentences를 이용해 청크를 문장 단위로 나누고 나눠진 청크에 맞춰 chunk id를 sentence_to_chunk에 순서대로 매핑
• embedding model: 기본값은 all-MiniLM-L6-v2
• output: pkl파일로 sentences, embeddings, chunk id, chunk, model name을 반환
  • pkl: pickle 라이브러리의 기능으로 데이터를 binary 형태로 변환

arag/src/arag/tools/registry.py

• ToolRegistry: 도구를 등록해두는 역할
  • LLM이 tool(keyword, semantic, read_chunk)를 문자열로 선택하게 만듦
  • get_all_schemas: 어떤 tool들을 사용할 수 있는지를 알리는 역할

arag/scripts/batch_runner.py

_init_shared_tools

• KeywordSearchTool, ReadChunkTool을 ToolRegistry에 저장
• build_index에서 만든 pkl 파일을 이용해 embedding 모델 로드 후 SemanticSearchTool도 저장

_load_question: limit개 만큼의 query만 처리(테스트용) _load_completed_qids: 이미 처리된 질문은 스킵

• _append_prediction: 이미 처리된 질문 리스트 제작

_create_agent: 위의 tool바탕 LLM 호출 _process_one: 한 질문에 대해서 처리 run: 이미 처리된 질문들 제외, pending questions에 대해서 처리

main tools

keyword_search

tiktoken

• LLM 사용을 위해 텍스트를 토큰화하고 토큰 수를 계산하거나 토큰을 다시 텍스트로 변환해주는 도구

실제 작동

• _load_chunks와 _split_sentences는 위와 동일
• name, get_schema: LLM에게 keyword_search에 대해서 전달

execute

• top k: 기본값 5, 최대 20
• chunk의 text와 입력받은 keyword를 모두 lower로 변환하고 count를 이용해서 chunk별로 해당 keyword가 얼마나 들어있는지를 확인
  • 최종 점수는 키워드 개수 $\times$ 키워드의 길이
• 만약 chunk 전체에 키워드가 한 개라도 포함되어 있다면(최종 점수가 0보다 크다면) 해당 키워드가 포함된 sentence를 찾음
  • 이후 scored_chunks에 chunk id와 점수, match된 sentence를 저장하고 sort를 이용해 top_chunk들을 찾음
  • 이때 토큰 관리를 위해 match된 sentence는 5개만 저장
• top_chunks에 대해 snippet을 생성
  • 정확히 그 단어가 들어간 문장만 따로 저장
  • 이때 앞뒤로 ...을 붙여서 snippet임을 알림
• 이후에 matched sentence를 리스트형태로 저장하고, 이를 sentences_text라는 하나의 텍스트로 결합한 뒤 token의 개수를 구해 저장합니다.
• return은 토큰의 개수와 top_chunk를 전달합니다.

semantic_search

앞쪽은 keyword_search와 거의 동일

• _load_index를 통해 임베딩을 가져옴

execute

• top k: 기본값 5, 최대 20

• query를 임베딩 + normalization

  • 이때 multi-thread에서도 안전하게 실행되도록 lock을 걺

• similarities: chunk별 임베딩값과 query 임베딩값을 dot product하여 cosine similarites 계산

• top_indices

  • argsort: similarities 바탕 오름차순 전환 $\rightarrow$ [::-1]: 내림차순으로 뒤집음 $\rightarrow$ top_k의 3배를 가져옴(노이즈가 존재하기에)

• 다시 청크화

  • top_indices들에 대해 sentences, 유사도, top_indices에서의 position을 각각 저장
  • 이후 동일한 청크는 같은 키값의 chunk_sentences에 저장

• chunk_sentences에서 청크별로 max similarity를 구함

  • 이후 chunk_scores에 chunk id, max similarity, sent(sentences, 각각 유사도, position)를 tuple형태로 저장
  • chunk_scores는 청크별로 max similarity를 기준으로 다시 sort
  • 이후 top_k를 뽑아냄

• k개에 대해서

  • 청크 내에서 유사도 기준이 아닌 원래 문장 순서를 기준으로 재정렬
  • 이후 snippet 생성
  • result_parts에 chunk ID, similarity, matched text를 저장

• 이후 최종 top_k개의 청크와 토큰 개수, 찾은 청크 수 반환

read_chunk

• keyword와 동일하게 청크를 dictionary형태로 변환

  execute

• arguments
  • context: 현재 에이전트의 상태
  • chunk_ids: 여러 청크를 한 번에 읽기 위한 리스트
  • chunk_id: 하나의 청크만 읽을 때 사용하는 단일 ID -> backward compatibility
• chunk_ids $\rightarrow$ str형태로 변환(키값으로 쓰기 위함)
• 이미 읽은 청크인지 확인(is_chunk_read)
  • 출력은 하되 내용은 다시 보여주지 않음.
  • 이미 읽음 표시만 남김
  • continue로 이후 로직을 건너뜀
• 새로운 chunk인 경우
  • 실제 chunk 내용을 가져옴
  • 새로 읽은 chunk는 토큰 계산
  • 이후 읽음 상태로 기록(mark_chunk_as_read)
• 이후 too_result로 합치고 log(토큰 개수, 새로 읽은 청크 개수, 이미 읽은 청크 개수)와 함께 반환

1. Introduction

기존의 RAG

• 알고리즘을 설계하여 여러 문서를 한 번에 검색하고 이를 모델 입력에 결합하는 방식
• 사전에 정의된 workflow를 바탕으로 모델이 여러 단계에 걸쳐서 수행하도록 유도하는 방식

A-RAG

• corpus 내의 정보가 본질적으로 서로 다른 수준의 단위로 구성되어 있음
  • 키워드 수준 - 문장/청크 수준
• 에어전트가 다양한 수준의 정보에 접근할 수 있도록 검색 도구 집합 설계

2. Related Work

Basic RAG

• retrieval이 모델이 외부 지식을 활용해 질문에 더 정확하게 답하도록 도와줌
• query rewriting, adaptive routing strategy, retrieval quality evaluation, reranking

  Graph RAG

• Corpus로부터 entity-relation graph를 구성
  • 대규모 지식 기반을 전체적으로 이해할 수 있도록 도움
• 실제 예시
  • RAPTOR: 재귀적 요약 + 계층적 트리 구조
  • LightRAG: 지식 그래프 + vector retrieval
  • HippoRAG: 해마의 기억 인덱싱 방식 모방
• 한계
  • 사전에 정의된 검색 알고리즘에 의존
  • 초기 검색 결과가 충분하지 않은 경우 추론 바탕의 추가 정보 수집이 불가능함

Workflow RAG

• RAG + Agent: 사전에 정의된 agent workflow에 의존함
• Non-learning
  • FLARE: 생성 신뢰도가 낮아질 떄 retrieval을 수행
  • IRCoT: CoT와 retrieval을 교차 수행
  • RA-ISF: self-feedback을 통해 복잡한 query를 분해
  • MA-RAG: 전문화된 에이전트들이 협력적으로 CoT 수행
  • RAGentA: 하이브리드 검색 + 인용 추적
• Learning
  • 작은 모델도 효과적인 검색 전략 학습 가능
  • 설계 시점의 워크플로우가 고정 -> 작업의 특성에 따라 전략을 바꾸는것이 불가능함

3. Methodology

3.1 Hierarchical Index Construction

• corpus 정보를 서로 다른 추상화 수준에서 조직

  chunking

• LinearRAG의 설정에 따라 corpus를 100 token 크기의 청크로 분할 (영단어 기준 700~800단어)
• semantic 일관성: 청크의 경계와 문장 경계가 일치하도록 설정
• 청크별로 독립적인 의미 단위로 작동

  Emedding

• 각 청크 $c_i$를 문장 단위로 분해: {$s_{i,1}, s_{i,2},\cdots,s_{i,n_i}$}
  • 규칙 기반 문장 분할로 수행됨
• 사전 학습된 sentence encoder $f_{emb}$로 dense vector 계산
  • $v_{i,j} = f_{emb}(s_{i,j})$
  • semantic 기반 매칭을 가능하게 함 + 문장이 속한 원래 청크와의 연결 유지

    Keyword-Level

• 사전 인덱스 없음
  • 오프라인: inverted index, knowledge graph 구축
  • 질의 시점: 텍스트 매칭을 수행

3.2 Hierarchical Retrieval Interfaces

keyword Search

• 어휘 매칭을 통해 특정 용어를 포함하는 청크를 찾음

  • agent는 키워드 집합 $K = {k_1, k_2, \cdots, k_m}$과 반환할 결과 수 k를 입력으로 제공
  • 관련성 점수: $Score_{kw}(c_i, K) = \displaystyle \sum_{k \in K} count(k, T_i) \cdot |k|$
    • $count(k, T_i)$: k가 청크 텍스트 $T_i$에서 등장한 횟수
    • $|k|$: 키워드의 문자 길이(긴 키워드 == 구체적)
  • $Snippet(c_i, K) = {s \in Sent(c_i) | \exists k \in K, k \subseteq s}$
    • 키워드를 포함하는 문장만 추출하여 축약
    • $Sent(c_i)$: 청크 $c_i$에 포함된 문장들의 집합
  • 상위 k개의 청크 ID와 snippet을 반환함

• 추가적인 부분

  • Score상 청크의 문장이 무조건 길수록 점수를 가져가기 유리한 구조임 -> snippet을 통해 완충한 것으로 파악됨 + 후보를 찾는 수준으로만 작동

Semantic Search

• dense search를 통해 의미적으로 유사한 문장을 찾아냄
  • query q에 대해 임베딩: $v_q = f_{emb}(q)$
  • 모든 문장 임베딩과 코사인 유사도 계산: $Score_{emb}(S_{i,j}, q) = {{v^T_{i,j}v_q}\over{||v_{i,j}||||v_q||}}$
  • 상위 점수를 가진 문장들을 선택하고 이를 해당 문장이 속한 청크 단위로 묶음
  • 상위 k개의 청크 ID와 청크에서 매칭된 문장들을 snippet으로 반환
  • 관련성 점수는 가장 높은 점수를 가진 문장에 의해 결정

Chunk Read

• keyword search와 semantic search에서 반환된 snippet기반, 에이전트가 어떤 청크를 전체적으로 읽어야 하는지를 판단하도록 함
• 에이전트: 선택한 청크의 전체 내용 확인, 필요에 따라 인접한 청크를 함께 읽음
• 에이전트가 자신의 판단에 따라 다른 수준에서 코퍼스 정보에 접근할 수 있도록 함

3.3 Agent Loop

Agent Loop

• ReAct와 유사한 프레임워크: 추론과 도구 호출을 번갈아 수행
  • 반복마다 하나의 도구를 선택하여 호출 $\rightarrow$ 결과를 관찰한 뒤 다음 행동 결정
  • 기타 복잡한 설계를 제외: 서로 다른 인터페이스 구성이 에이전트의 행동에 미치는 영향 관찰
  • 최대 반복 횟수: 지금까지 수집한 정보를 바탕으로 최종 답변을 생성하도록 유도

Context Tracker

• 이미 읽은 청크를 기록: 중복된 정보 검색 + 불필요한 토큰 사용 방지
  • $C_{read} = {c_{i1},c_{i2},\cdots,c_{ik}}$
  • 이미 읽은 청크를 다시 읽으려고 할 경우 읽은 적 있다는 알림 메시지 반환

4. Experiments

• 평가 지표
  • LLM-Acc: 의미적으로 정답인가(LLM 기반 평가)
  • Contain-Acc: 정답이 텍스트에 포함되는가
• A-RAG
  • Naive: embedding 하나만 적용
  • Full: hierarchical tools
• 결과
  • 대부분 벤치마크에서 최고 성능
  • GPT-5-mini에서 더 크게 향상

    Ablation Study

    

5. Analysis and Discussion

Test-Time Scaling

• step 수가 증가하면 성능이 일관되게 증가함

• reasoning effor(한번에 사용하는 추론 연산량)가 증가하면 성능이 일관되게 증가함

• 강한 모델일수록 성능이 증가함

  Context Efficient

  

• 비슷하거나 더 적은 토큰으로 더 높은 성능 달성

  Failure Mode Analysis

  

• 기존 RAG: 문서 자체를 못찾음

• A-RAG: reasoning 문제나 entity confusion

6. Conclusion

• 향후 연구가 복잡한 검색 알고리즘 설계보다는 모델이 활용하기 쉬운 인터페이스 설계에 더 집중해야 함을 시사함
• LM과 외부 지식 간 새로운 상호작용 패러다임 탐구가 중요함

설명

PySceneDetect

• 연속된 영상 파일에서 장면 전환 지점을 자동으로 탐지하는 파이썬 라이브러리
• Cut이나 fade in/out, 디졸브 등 전환을 색상 변화나 밝기 변화를 기준으로 감지함
  • 감지된 장면을 기준으로 Video splitting도 가능

Content Detect

• 인접한 두 프레임 간 content의 차이가 threshold를 넘길 경우 장면 전환으로 판단함
• 다른 detector들에 비해서 빠르고 직관적인 편
• 카메라가 빠르게 움직이는 등의 전체 프레임이 계속 크게 흔들리는 경우에는 오탐이 늘어날 수 있음
• threshold, min_scene_len(최소 프레임) 등을 조절할 수 있음
  • 기본값 threshold == 27, min_scene_len == 15

코드

_mean_pixel_distance

def _mean_pixel_distance(left: numpy.ndarray, right: numpy.ndarray) -> float:
    """Return the mean average distance in pixel values between `left` and `right`.
    Both `left and `right` should be 2 dimensional 8-bit images of the same shape.
    """
    assert len(left.shape) == 2 and len(right.shape) == 2
    assert left.shape == right.shape
    num_pixels: float = float(left.shape[0] * left.shape[1])
    return numpy.sum(numpy.abs(left.astype(numpy.int32) - right.astype(numpy.int32))) / num_pixels

두 이미지 간(두 프레임 간) 픽셀 값 차이를 평균을 반환하는 함수

• assert (조건): 뒤의 조건이 맞지 않을 경우 에러를 반환
  • 두 이미지 모두 2차원인지 확인(흑백)
  • 두 이미지의 크기가 같은지 확인
• num_pixels: 전체 픽셀 수 (height * width)
• return
  • astype: 이미지를 int32로 반환(음수값 대비)
  • 이후 차를 절댓값(abs)한 값들을 합산 후 나누기

_estimated_kernel_size

def _estimated_kernel_size(frame_width: int, frame_height: int) -> int:
    """Estimate kernel size based on video resolution."""
    # TODO: This equation is based on manual estimation from a few videos.
    # Create a more comprehensive test suite to optimize against.
    size: int = 4 + round(math.sqrt(frame_width * frame_height) / 192)
    if size % 2 == 0:
        size += 1
    return size

filter(kernel)의 크기를 해상도에 맞기 설정

• size: 전체 픽셀 수(width * height) / scale 조정 값(hyperparameter)
  • 해당 값을 sqrt하여 대략적인 길이를 구함
  • +4는 최소값 보정
• kernel은 중심이 필요하기에 홀수값으로 맞춰줌

class ContentDetector

class Components

class Components(ty.NamedTuple):
        """Components that make up a frame's score, and their default values."""

        delta_hue: float = 1.0
        """Difference between pixel hue values of adjacent frames."""
        delta_sat: float = 1.0
        """Difference between pixel saturation values of adjacent frames."""
        delta_lum: float = 1.0
        """Difference between pixel luma (brightness) values of adjacent frames."""
        delta_edges: float = 0.0
        """Difference between calculated edges of adjacent frames.

ContentDetector의 경우에는 BGR값을 HSV로 변환해서 작동함

• Hue: 색의 종류(color wheel의 각도)
• Saturation: 채도(색이 얼마나 선명한가)
• Luma: 밝기(흑백 기준과 유사)
• 각각 delta를 붙여서 두 프레임의 픽셀 간 차이를 나타냄
• 만약 Luma-only를 킬 경우 delta_lum만 1.0, 나머지는 0.0으로 변환

_calculate_frame_score

def _calculate_frame_score(self, timecode: FrameTimecode, frame_img: numpy.ndarray) -> float:
        """Calculate score representing relative amount of motion in `frame_img` compared to
        the last time the function was called (returns 0.0 on the first call)."""
        # TODO: Add option to enable motion estimation before calculating score components.
        # TODO: Investigate methods of performing cheaper alternatives, e.g. shifting or resizing
        # the frame to simulate camera movement, using optical flow, etc...

        # Convert image into HSV colorspace.
        hue, sat, lum = cv2.split(cv2.cvtColor(frame_img, cv2.COLOR_BGR2HSV))

        # Performance: Only calculate edges if we have to.
        calculate_edges: bool = (self._weights.delta_edges > 0.0) or self.stats_manager is not None
        edges = self._detect_edges(lum) if calculate_edges else None

        if self._last_frame is None:
            # Need another frame to compare with for score calculation.
            self._last_frame = ContentDetector._FrameData(hue, sat, lum, edges)
            return 0.0

        score_components = ContentDetector.Components(
            delta_hue=_mean_pixel_distance(hue, self._last_frame.hue),
            delta_sat=_mean_pixel_distance(sat, self._last_frame.sat),
            delta_lum=_mean_pixel_distance(lum, self._last_frame.lum),
            delta_edges=(
                0.0 if edges is None else _mean_pixel_distance(edges, self._last_frame.edges)
            ),
        )

        frame_score: float = sum(
            component * weight
            for (component, weight) in zip(score_components, self._weights, strict=True)
        ) / sum(abs(weight) for weight in self._weights)

        # Record components and frame score if needed for analysis.
        if self.stats_manager is not None:
            metrics = {self.FRAME_SCORE_KEY: frame_score}
            metrics.update(score_components._asdict())
            self.stats_manager.set_metrics(timecode, metrics)

        # Store all data required to calculate the next frame's score.
        self._last_frame = ContentDetector._FrameData(hue, sat, lum, edges)
        return frame_score

이전 프레임과 비교해서 얼마나 움직임이 있었는지를 계산하는 함수

• cv2.cvtColor(frame_img, cv2.COLOR_BGR2HSV): BGR값으로 표현된 이미지 벡터를 HSV형태로 변환
  • cv2.split: 리스트 형태로 hue, sat, lum을 분리함
• calculate_edges: edge계산이 필요한지 확인
  • 만약 필요하다면 _detect_edges로 넘어감
• 만약 첫 프레임이라면 0.0을 반환
• score_components: _mean_pixel_distance를 활용해서 hue, sat, lum, edges의 변화량 계산
• frame_score: 각 component(hue, sat, lum, edges)에 weight를 곱한 뒤(zip 활용), 합을 나눠서 평균
  • 어떤 변화가 더 중요한지를 조절함
• 필요할 경우 stats_manager로 통계 기록을 남김

process_frame

def process_frame(
        self, timecode: FrameTimecode, frame_img: numpy.ndarray
    ) -> ty.List[FrameTimecode]:
        """Process the next frame. `frame_num` is assumed to be sequential.

        Args:
            frame_num (int): Frame number of frame that is being passed. Can start from any value
                but must remain sequential.
            frame_img (numpy.ndarray or None): Video frame corresponding to `frame_img`.

        Returns:
           ty.List[int]: List of frames where scene cuts have been detected. There may be 0
            or more frames in the list, and not necessarily the same as frame_num.
        """
        self._frame_score = self._calculate_frame_score(timecode, frame_img)
        if self._frame_score is None:
            return []

        above_threshold: bool = self._frame_score >= self._threshold
        return self._flash_filter.filter(timecode=timecode, above_threshold=above_threshold)

계산한 score가 threshold를 넘기는지 확인 + filter 처리

• _flash_filter: 노이즈나 순간적인 반짝임들을 제거
  • 연속된 변화만을 남김

_detect_edges

def _detect_edges(self, lum: numpy.ndarray) -> numpy.ndarray:
        """Detect edges using the luma channel of a frame.

        Arguments:
            lum: 2D 8-bit image representing the luma channel of a frame.

        Returns:
            2D 8-bit image of the same size as the input, where pixels with values of 255
            represent edges, and all other pixels are 0.
        """
        # Initialize kernel.
        if self._kernel is None:
            kernel_size = _estimated_kernel_size(lum.shape[1], lum.shape[0])
            self._kernel = numpy.ones((kernel_size, kernel_size), numpy.uint8)

        # Estimate levels for thresholding.
        # TODO: Add config file entries for sigma, aperture/kernel size, etc.
        sigma: float = 1.0 / 3.0
        median = numpy.median(lum)
        low = int(max(0, (1.0 - sigma) * median))
        high = int(min(255, (1.0 + sigma) * median))

        # Calculate edges using Canny algorithm, and reduce noise by dilating the edges.
        # This increases edge overlap leading to improved robustness against noise and slow
        # camera movement. Note that very large kernel sizes can negatively affect accuracy.
        edges = cv2.Canny(lum, low, high)
        return cv2.dilate(edges, self._kernel)

edge를 별도로 구분하는 이유

• HSV만으로 판단하기에는 애매한 경우가 많음
  • 밝기만 많이 변하는 경우(조명 유무)
  • 색만 변하는 경우(색 필터 등)
• edge
  • 형태를 남겨서 조명 등에 영향을 덜받게함

코드 설명

• _estimated_kernel_size로 kernel 생성
  • 이때 lum만 사용(edge 검출은 밝기 변화가 제일 명확함)
• sigma: canny edge 검출에 사용될 구간 값 결정
  • median을 기준으로 +sigma, -sigma까지 검출
• cv2.Canny: canny edge detect algorithm 활용
• cv2.dilate: edge를 더 두껍게 함
  • edge 비교가 더 안정적이도록 만듦(노이즈와 느린 움직임에 더 강해짐)

결론

Detection에 있어서 HSV > RGB

• RGB의 경우에는 밝기와 색이 섞인 형태
  • 어느쪽의 변화인지 구분하기가 어려움

색 정보만이 아닌 edge 검출을 별도로 이용

• HSV정보만을 사용하면 밝기나 필터 등 변수가 많음
  • edge를 별도로 검출하여 형태를 바탕으로도 구분해냄

1. Introduction

Goal: Untrimmed video에서 행동 instance의 시작점과 끝점을 찾고 동시에 행동의 class를 예측하는 것

기존의 방식

• 행동 instance의 후보를 추출
• 각 후보를 개별적으로 처리
• 한계
  • 행동이 복잡한 움직임을 포함함
  • 동일 클래스에도 변동성이 큼
  • 조명, 시점, 복잡한 배경으로 인한 어려움

Diffusion

• 노이즈가 섞인 불확실한 이미지에서 노이즈를 제거하는 과정을 통해 고품질 이미지 생성
• 이를 이용해 행동 탐지 이미지를 생성하는 문제로 치환
  • 행동 클래스 예측을 $N \times C$(N은 프레임 수, C는 class 수)
  • 시작점과 끝점 예측을 $N \times 2$

ADI-Diff

• Diffusion을 통해 Action class, 시작점, 끝점에 해당하는 AD 이미지 생성
• forward process: 정답 AD 이미지에 점진적으로 노이즈 추가
  • Discrete Action-Detection Diffusion Process: 이산 확률 분포를 생성
• reverse process: 노이즈가 포함된 입력에서 노이즈를 제거하고 고품질 AD 이미지 생성
• Row-Column Transformer: 행 방향은 시간적 순서성, 열 방향은 클래스간 관계

2. Related Work

Anchor

• Anchor-based

  • 비디오 전반에 걸쳐 미리 정의된 길이를 가지는 다수의 anchor 생성
  • multi-tower network, temporal feature pyramid network

• Anchor-free

  • 비디오 프레임에 대해 actionness scroe / 행동 경계 신뢰도 점수를 예측

시공간 정보 인코딩

• 기존: RNN, Graph, Transformer
• proposal-free

3. ADI-Diff Framework

Formulation of AD images

Goal: AD의 세 가지 출력을 세 개의 이미지로 재구성할 수 있음을 관찰

Action-class AD image $x^a$

• $x^a = N\times C$, [0, 1]범위 N = 비디오 프레임 수, C = 행동 클래스 개수
• $x^a$의 n번째 행, c번째 열의 값: n번째 프레임에서 행동 c가 발생할 확률

Starting/Ending point AD image $x^s, x^e$

• $x^s \in [0,1]^{N \times 2}$, $x^e \in [0,1]^{N \times 2}$
• 각 행은 하나의 프레임, 행마다 2개의 픽셀
  • 첫 번째 값: 해동 프레임에서 시작점/끝점이 존재할 확률
  • 두 번째 값: 존재하지 않을 확률

결론적으로 $x^a, x^s, x^e$ 모두 grayscale의 이미지로 image diffusion으로 생성될 수 있음

Discrete Action-Detection Diffusion Process

표준 이미지 기반 모델

• 완전히 노이즈가 섞이고 불확실한 이미지에서 시작
• 여러 단계를 거쳐 점진적으로 노이즈와 불확실성 제거
• Forward process: Gaussian noise 점진적으로 추가
• Reverse process: Gaussian noise를 제거하는 과정 학습

AD생성

• Gaussian이 아닌 이산 확률 분포를 나타냄
• 정답인 이상적인 이산 확률 분포와 완전히 불확실한 이산 확률 분포를 연결
  • Discrete Action-Detection Diffusion Process
  • 가장 불확실한 상태: Uniform distribution(Property 1)
• 모델 학습을 위한 추가적인 성질
  • Property 2: 여러 단계의 Forward Process를 건너뛸 수 있는 $q(z_t|z_0)$가 존재해야 함
  • Property 3: Forward Process posterior인 $q(z_{t-1}|z_t, z_0)$를 계산할 수 있어야 함

Forward Process

• 정답 AD 이미지를 초기값으로 설정
• 각 행을 점진적으로 Diffusion하여 T단계에 걸쳐 Uniform distribution으로 수렴

정의

• $z_t$: 확산 과정의 t번째 단계에서의 이상 확률 분포 (길이 C인 벡터)
• $z_0$: 초기 상태, 정답 클래스 위치에 1, 나머지는 0인 one-hot vector
• 총 T단계 구성, 점진적으로 노이즈 추가, $z_T$가 균등분포에 가까움
• 각 단계 t에서 노이즈 $v_t$ 추가
  • $v_t$는 균등 확률을 가지는 Multinomial distribution에서 샘플링

실제 과정(단일 단계) $z_t = (1-\beta t)z{t-1} + \beta _t v_t$

• 각 단계 t에서 작은 확률로 임의의 클래스를 동일한 확률로 선택할 수 있도록 불확실성 증가
• $\beta_t$: 해당 단계에서의 랜덤성 증가 정도(hyperparameter)
• $v_t$는 길이 C인 벡터, 모든 원소가 0이상, 합이 1
  • $v_t$는 $MN_K(K, {{1}\over{C}1})$를 샘플링
  • 샘플링 결과를 K로 나눠서 합이 1이 되도록 함

즉, 단일 단계: $q(z_t|z_{t-1}) = MN_K({{z_t -(z-\beta_t)z_{t-1}}\over{\beta_t}})$

t단계 전체를 고려했을 때 $z_t = \bar \alpha_t z_0 + (\displaystyle \Pi_{r=2}^t \alpha_{\tau})\beta_1 v_1 + ... + \beta_t v_t$

• $\alpha_t = 1-\beta_t$
• $\hat \alpha_t = \Pi^t_{r=1} \alpha_\tau$

따라서 $q(z_t|z_0) = MN_{B_tK}({{z_t-\hat \alpha_t z_0}\over{1- \hat \alpha_t}})$

• 중간 단계 분포를 직접 생성할 수 있기에 효율적인 학습이 가능함

이때 $\beta_t$가 충분히 크다면 $\hat \alpha_t \rightarrow 0$이 되어서 최종 분포 $z_T$는 균등 분포로 수렴 $E[z_T] = {{1}\over{C}}1$

property 3. forward process posterior

• Markov chain에 따라서 $q(z_{t-1}|z_t, z_0) = {{1}\over{\sigma_t}}\cdot$(두 Multinomial 분포의 곱)
• 이때 $\sigma_t$는 정규화 상수. 실제에선 hyperparameter

Reverse Process

• Forward Process를 통해 중간 분포 ${z_1, ..., z_T}$를 생성할 수 있음
• 중간 분포를 활용해 reverse process 학습
• T 단계에 걸쳐서 AD image 출력

정보 추출

• Diffusion 과정을 보조
• 학습된 feature extractor 이용해서 비디오의 snippets에서 feature 추출
• $f_{ST} \in \mathbb R^{N \times C_{ST}}$
  • N은 프레임 수, $C_{ST}$는 채널 수
  • $f_{ST}$는 condition으로 사용

과정

• 초기 input: 완전히 불확실한 분포 $Z_T$
• Reverse Process를 수행하며 최종 예측 $\hat z_0$를 예측
• 단계별 $\hat z_{t-1} = d_\phi (\hat z_t, f_{ST}, f_t)$, $t \in {1,...,T}$
  • $f_t$는 t번째 diffusion 단계를 나타내는 임베딩, 사인 함수를 이용해 생성

Multi-row Processing

• 동일한 Diffusion 과정은 전체 AD image의 모든 행에 동시에 적용이 가능함
• 순방향 과정에서 모든 행에 동시에 노이즈를 추가 {$x^1_a, ..., x^T_a$}를 생성
• 역방향 과정에서 각 단계에서 $x^t_a$로부터 $x^{t-1}_a$를 생성하도록 학습

Row-Column Transformer Architecture

• Row(frame): temporal한 관계
• Column(class): 서로 다른 클래스 간 관계
• 기존의 Diffusion network: 2D 공간의 local negihbor을 처리하는 방식

input

Diffusion network d의 t번째 단계

• input image: $x_a \in \mathbb R^{N \times C}$
• 시공간 특징: $f_{ST} \in \mathbb R^{N \times C_{ST}}$
• Diffusion 단계 임베딩: $f_t \in \mathbb R^{N \times 1}$
  • $f_t$는 사인 함수 기반 임베딩
• concat: $x \in \mathbb R^{N \times (C + C_{ST} + 1)}$

Row-Column Block

Column

• 클래스 간 관계는 장거리 관계를 가질 수 있기에 Multi-Head Self-Attention(MHSA) 사용
  • 입력 x의 각 열을 하나의 토큰으로 간주
  • 총 $C+C_{ST} + 1$개의 토큰 생성, 각 토큰의 길이는 N
• 중간 출력 $u_{col} \in \mathbb R^{N \times (C+C_{ST}+1)}$을 얻고 두 개의 MLP Layer 통과

Row

• Temporal convolution(TC) + MHSA
• $1 \times 3$ Temporal Convolution을 적용, local 관계를 반영한 $u_row$ 생성
• 각 행을 하나의 토큰으로 간주하여 MHSA 수행

Combined Image Processing

• 동일한 방법으로 $x_s$와 $x_e$에도 적용
• 각각 따로 생성 or 세 이미지를 하나로 결합하여 동시에 처리도 가능
  • $x^{combined}_t \in \mathbb R^{N \times (C+4)}$
  • 효율성 증가
  • 분류와 위치 추정 사이의 정보를 공유하여 성능 향상

Inference and Training Pipeline

Inference Pipeline

모델이 출력한 세 개의 AD 이미지 $\hat x^0_a, \hat x^0_s, \hat x^0_e$를 이용해 최종 행동 탐지 결과 생성

• 시작점/끝점 후보 찾기
  • $\hat x^0_s, \hat x^0_e$의 첫 번째 열 확인
  • 값이 임계값 이상인 프레임들을 후보로 선택
• grouping
  • 연결된 프레임들을 하나의 그룹으로 묶고
  • 평균 위치를 실제 시작점 / 끝점으로 결정
• candidate 생성
  • 시작점에 대해 끝점과 짝을 지음
  • 시작점 ~ 끝점 구간을 하나의 행동 후보로 생성
• 행동 클래스 결정
  • 해당 구간에 해당하는 $\hat x^0_a$의 행들을 평균
  • 가장 높은 값을 가진 클래스 선택
• Soft-NMS
  • 겹치는 후보들을 Soft-NMS로 제거
  • 최종 행동 탐지 결과 생성

    Training Pipeline

• 비디오 특징 추출
  • 사전 학습된 모델을 사용해 $f_{ST}$ 추출
  • 학습 중에는 이 특징 추출기 고정
• 모델 초기화
  • Diffusion model d는 랜덤 초기화
• Forward Process
  • 정답 이미지 $x_0$에 노이즈 추가
  • 중간 단계 분포 생성: 학습용 정답 역할
• Reverse Process
  • 확산 모델이 노이즈를 제거하여 중간 단계 및 최종 AD 이미지 에측
• Loss 계산
  • 각 단계에서 모델의 예측 AD 이미지와 Forward Process에서 생성된 정답 이미지 간 MSE Loss 계산하여 학습

4. Experiment

Ablation Study

Discrete Diffusion

• 기존 diffusion보다 성능이 크게 향상됨

Row-Column Transformer

• 일반 구조 대비 성능 상승

AD image 구성

• 시작/끝 이미지가 없다면 성능이 떨어짐

stitching

• 속도와 성능 모두 향상 병렬 처리 + 정보 공유의 효과

결론적으로 속도는 기존 최신 모델과 유사하면서 성능은 더 높음

1. Introduction

Robotic Manipulation Data

• 관절 매핑 기반 teleoperation tools: 사람이 컨트롤러나 장치를 조종 후 사람의 관절 움직임을 로봇 관절에 직접 매핑
• Exoskeleton: 사람 몸에 기계 장치를 착용
• VR device: 가상 환경에서 로봇을 원격 조작
• 로봇 장비와 숙련된 조작자가 필요하다는 단점이 존재함

인간의 영상을 통한 Manipulation Data

• Hand-Object Interaction: 인간의 시선은 다음 행동의 의도를 드러냄
  • 인간의 시선을 이용해 손과 물체의 상호작용을 예측하거나 생성
  • 인간의 조작 의도를 장기적으로 예측할 수 있음
• 학습 데이터의 양을 크게 늘릴 수 있음 + 작업과 환경의 다양성 증가

EgoVLA

• 인간의 action space와 로봇의 action space간 차이는 크지 않으며 기하학적 변환을 통해 근사할 수 있음
  • 시각 관측 프레임, 언어 지시, 현재 손 자세를 입력으로 주면 몇 단계 미래의 인간 동작을 예측함
  • 손목 위치: Inverse Kinematics
  • 인간 손 관절: retargeting
• 일부 예측 오류가 존재함
  • 이는 소량의 teleoperation data로 fine-tuning하여 보정 가능

2. Related Work

Dexterous Manipulation

초기 연구

• Control-Based Method
  • 학습이 아닌 물리 법칙과 수학적 모델을 이용해 로봇 제어
    • Inverse Kinematics: 목표 위치를 주면 로봇 관절 각도 생성
    • Grasp Planning: 물체를 잡기 위한 손가락 위치와 힘 계산
  • 정밀한 조작은 가능하지만, 일반화 능력이 제한적
• Learning-Driven Approaches
  • Pose Vector: 로봇 손의 자세를 직접 벡터 형태로 생성
  • Intermediate Represnetations: 중간 단계의 표현만 먼저 예측
  • Contact Maps: 손과 물체 사이의 접촉 위치를 먼저 예측

Egocentric Human Video

• 특정 작업에 대한 정책 학습
  • 일반적 작업이 아닌 작업이 한정되어있다는 한계

VLA(Vision-Langauge-Action)

VLMs를 대규모 로봇 데이터로 fine-tuning

• perception과 action 통합
• 데이터 요구량이 매우 큼
  • 대규모 teleoperation data
  • 스크립트 기반 실행 데이터

    Egocentric Vision

    기존의 한계
• 데이터 규모와 다양성의 부족
• 복잡한 데이터 셋: 현재 로봇이 수행하기 어려운 활동들
• 단순한 데이터 셋: hand pose와 같은 정밀한 annotation이 부족

  Learning from In-the-Wild Video

  자연 환경의 영상
• Affordance, 상호작용 단서 추출
• 인간 영상을 이용해 표현을 pre-train하고 transfer이 가능
• 대부분 비지도 학습에만 초점을 둠

3. Learning Manipulation Skills from Ego-Centric Human Videos

Datasets

Language model과 VLMs에서 알 수 있듯, 데이터셋 구조가 모델의 성능에 큰 기여를 한다.

데이터셋 조건: 조작 행동을 포함하는 egocentric video + pose annotations

• Egocentric RGB observations
• wrist pose
• hand pose
• camera pose

실제 데이터셋

• HOI4D: 한 손 조작 작업
  • pick-and-place
  • re-orientation
  • articulated object interaction: 문, 서랍 등을 회전하거나 미는 동작
• HOI3D: 33개의 rigid objects와 상호작용
• HoloAssist: 양손 상호작용
  • 배터리 교체, 가구 조립, 기계 설정 등
  • 노이즈는 많지만 양손 상호작용을 풍부하게 포함
  • 다른 데이터셋에 비해 데이터가 많기에(166시간) 10% 균등 샘플링
• TACO: tool-action-object

Data Processing: Egocentric video는 카메라가 움직이기에 학습이 어려움

• World frame에서 미래 손목 위치를 계산하고, 이를 현재 카메라 기준으로 변환하여 사용

EgoVLA Model

base

• backbone: NVILA-2B
  • Vision-Language 이해 능력이 높아서 intention inference에 좋음
  • 비교적 작은 모델 크기로 fine-tuning에 좋음
• 입력
  • 현재 및 과거의 egocentric 시각 관측: head camera시점 6 frame
    • Visual Observation: 0.2초 간격 sampling(총 1초의 과거)
  • Language Instruction: High-level이 아닌 즉각적인 행동 지시
  • Action Query Tokens: 행동을 예측하도록 요청하는 special tokens
  • Human Proprioception: 손목의 이동, 회전, 손 자세
  • 해당 입력들을 backbon에서 인코딩하여 action head로 처리

Action Representation

정보

• 손목 자세
  • 3D translation
  • rot6D rotation
• hand joint angles
  • MANO hand model: 상위 15개 PCA 성분

학습(regression)

• 미래 손목 자세, 손 관절 파라미터
• $L = \lambda_{wrist_trans}L_{wrist_trans}+\lambda_{wrist_rot}L_{wrist_rot}+\lambda_{joint}L_{joint}$

  Action Head & Action Query Tokens

  구성
• 6개의 encoder layer
• hidden size 1536

input

• 인간/로봇의 proprioception state
• action query token의 latent embedding

output

• 미래 행동 시퀀스
• $A_t = [a_t, a_{t+1}, ..., a_{t+H}]$

Prediction Range

• 1초 미래(30 future steps at 30 Hz)

Action Query Token

• vocabulary의 마지막 30개의 word ID

Training Details

1. 인간 데이터 pre-training

• Ego-Centric Human Manipulation Datasets, 20epochs

2. 로봇 데이터 post-training

• 115 epochs, 100 epoch 이후 learning rate 감소

Transferring EgoVLA to Humanoid Robot

사람과 휴머노이드 모두 두 팔과 손을 이용한다는 점은 동일하지만, 카메라의 위치, 손의 형태, 시각적 외형의 차이로 인해 바로 적용은 힘듦

• 소량의 로봇 데이터를 이용해 fine-tuning
• 이땐 Unified Action Space를 활용

Retargeting robot data to human representation

로봇의 action space를 인간의 action representation으로 변경하여 정렬

• end-effector pose: 3D 변환으로 로봇과 인간의 좌표계 정렬
• hand configuration: MANO hand model parameter
  • $\displaystyle min_{\Theta}L(\Theta) = {{1}\over{5}}\displaystyle \sum^5_{i=1} SmoothL1(J_{pred}(\Theta)i, j{obs, i})$
  • $\Theta \in \mathbb R^{15}$: MANO 손 파라미터
  • $J_{pred}(\Theta)$ MANO forward kinematics로 게산된 손가락 끝
  • $J_{obs} \in \mathbb R^{5 \times 3}$ 로봇 손가락 끝 실제 위치

    Mapping human hand to robot hand

    Inference: EgoVLA가 예측한 손목, 손 자세를 robot actuerator instructure
• 예측한 손목 자세 -> 3D 변환으로 End-effector로 변환 -> Inverse Kinematics로 로봇 팔 관절 계산

손 동작 변환

• 예측된 MANO 파라미터로 3D 손 keypoint 계산
• MLP 모델이 손 keypoint로 로봇 손 관절 명령 예측
  • MLP는 로봇 시연 데이터로 학습
• 원래 시연 데이터로도 정상적으로 수행되며, 작은 오차가 큰 영향을 주진 않음

  4. Ego Humanoid Manipulation Benchmark

  필요성

  실제 로봇 실험의 문제
• 비용이 매우 비쌈
• 실험 시간이 오래 걸림
• 실험 환경을 동일하게 재현하기 어려움
• 안전 문제

이에 시뮬레이션 기반 평가를 많이 사용함 해당 벤치마크는 sim2real transfer보단 통제되고 재현 가능한 테스트 환경으로 시뮬레이션 활용

Short Horizon Tasks(단일 동작 중심)

• Push-Box, Filp-Mug, Pour-Balls, Close-Drawer, Open-Drawer, Open-Laptop, Stack-Can

Long Horizon Tasks(여러 단계 행동 결합)

• Sort-Cans, Insert-Cans, Unload-Cans, Insert-And-Unload-Cans, Stack-Can-Into-Drawer

Observation & Action Space

벤치마크 관측 정보

• robot joint position
• end-effector poses
• contact force
• egocentric RGB-D visual input

EgoVLA

• Egocentric 시각 정보
• End-effector pose
• hand joint actuation
• task descriptions

로봇 제어

• 팔: end-effector control
• 손: PD joint control
  • 6개의 active, 6개의 mimic joints -> 12개의 DoFs

    PD control: 목표 관절 위치에 도달하도록 오차를 이용해 힘 계산 $\tau = K_p(q_{target} - q) + K_d(\dot q_{target}- \dot q)$ $q$는 위치, $\dot q$는 관절 속도

최종적으로 36차원 action space

• arm inverse kinematics
• direct hand actuation
• Control frequency = 30Hz
• 각 작업마다 매 단계 성공 여부, sub-task 완료 여부 확인

Diverse Visual Background

5개의 방, 5개의 테이블 텍스쳐 = 25개의 시각적 환경

• Generalization 성능 평가

Demonstrations

Imitation Learning을 지원하기 위해 전문가 데이터 제공

• teleoperation system으로 수집
• Room 1, 2, 3, Table 1
• 총 100개의 성공 시연

5. Experiments

Human Manipulation Modeling

인간 손 움직임을 얼마나 잘 모델링하는가

• 인간 손목과 미래 위치 예측 오차는 약 8cm
• 언어 지시의 변경에 따라 경로도 정상적으로 변경됨
• 행동 뒤에 있는 semantic intent도 정상적으로 학습함

Evaluation

평가 지표

• Success Rate(SR): 작업을 완전히 성공한 비율
• Progress Rate(PSR): 전체 단계 중 얼마나 많은 단계까지 수행했는지

Baseline

• EgoVLA: 인간 영상 사전 학습 + 로봇 데이터 fine-tuning
• EgoVLA-NoPretrain: 로봇 데이터만 학습
• ACT: 각 작업마다 따로 학습되는 전문 모델

  Evaluation Setup

• Seen environment: 학습에서 본 환경
• Unseen environment: 완전히 새로운 환경(일반화 평가)

Result

1. 인간 데이터 pretraining이 성능을 향상

• short-horizon, long-horizon 모두에서 EgoVLA > EgoVLA-NoPretrain
• 특히 Stack-Cans, Sort-Cans, Insert-And-Unload-Cans, Flip-Mug에서 큰 차이
  • 인간 데이터는 일반적인 조작 기술 학습에 큰 도움
• Long-Horizon Task에서 더 큰 차이
• 새로운 환경에서의 일반화에 큰 영향

2. Specialist < generalist

• EgoVLA의 경우 low-level manipulation skill을 공유할 수 있음

3. 여전히 로봇 데이터는 필요함

• 인간 영상만으로는 zero-shot 성공률이 0%
• human과 robot간 차이가 존재함

6. Conclusion & Limitation

EgoVLA

• 대규모 egocentric human manipulation dataset으로 VLM을 pretraining한 후 소량의 로봇 데이터로 fine-tuning
• 이때 차이를 극복하기 위해 unfied action space 사용

Limitation

• hand and wrist pose annotation이 필수적임
• unified action space를 사용했음에도 추가적인 로봇 데이터로 fine-tuning 해야함

1. Introduction

인간의 지식은 로봇에 물리적인 지능을 부여하는 토대

• VLM을 통해 인간이 생성한 텍스트-이미지로 robot policies 초기화
• 로봇 원격 조작을 통해 인간이 생성한 행동 모방
• 누군가 과제를 수행하는 영상을 시청함으로써 인간의 경험으로부터 직접 학습

최신 LM: 데이터 소스를 활용하는 능력은 모델 규모와 밀접하게 관련됨

• 명시적인 alignment(매핑과정)없이 인간 비디오 데이터로부터 기술을 학습하는 능력은 규모와 관련이 있는가
• 이를 증명하기 위해 co-training 활용

  co-training: 인간 비디오 데이터 또한 로봇 데이터와 동일한 하나의 embodiment로 취급. 이때 인간 데이터도 로봇 데이터와 동일한 목적함수로 학습

1. 3D 손 추적 정보로 low level end effector trajectories를 예측
2. dense language annotation을 활용해서 high level의 하위 과제를 예측
3. 인간 데이터와 로봇 데이터를 혼합 및 co-finetuning
4. 인간 데이터에만 존재하는 설정에서 평가

결과: 로봇 데이터의 다양성이 확장될수록 Pre-training된 VLA는 인간 비디오를 더 효과적으로 활용할 수 있게 됨 사전학습의 다양성이 증가할수록 인간 데이터와 로봇 데이터 사이의 latent representation이 alignment됨. 즉, 데이터 범위가 확보되면 시각적, 운동학적 영역 차이가 크더라도 모델이 embodimnet에 구애받지 않는 표현 생성

2. Related Works

Learning from Humans

초기: 비디오 데이터를 활용해 Vision Encoder 학습

• 시각 표현은 학습할 수 있지만 행동 예측을 직접적으로 향상시키진 못함

Developed proxy

• Keypoint Tracking: 사람의 손, 물체의 keypoints를 추적
• Latent Actions: 명시적 행동 값 없이 latent space에 압축해서 표현
• Reward Modeling: goal state를 학습(success signal)
• Affordance Prediction: 물체가 무엇을 할 수 있는지를 예측

Explicit Alignment

• Overlayed robots: 인간이 수행하는 장면 위에 로봇의 외형/관절 모델을 겹쳐서 표시
• AR/VR: 인간이 AR/VR장비를 착용하고 작업하면서 시스템이 인간의 동작을 직접 측정/유도

한계점

• 규모가 작아지면 불안정함
• Alignment에 의존함

  Heterogeneous Vision-Language-Action Models

  로봇의 원격 조작 데이터 + 웹 규모의 Vision-Language 데이터, 언어 주석
• Vision-Language Backbone: 인간이 생성한 이미지 + 텍스트에서 폭넓은 의미 이해
• 대규모 원격 조작 데이터에 대한 행동 모방 학습

Cross-Embodiment Training

• 하나의 정책을 통합된 아키텍처와 행동 표현으로 여러 로봇 Embodiement를 제어하도록 학습.
• 공유된 관측 및 행동 공간 외 특별한 alignment없이 embodiement 간 Transfer이 가능함

  Scaling Alternative Data Collection Strategies

  로봇 원격조작 대용
• 손으로 조작하는 휴대용 하드웨어 사용
• UMI: 평행 집게. 자체 움직임 추적

  3. Preliminaries

  VLA 모델을 사용하여 범용 정책 학습
• Vision-Language 모델의 아키텍처와 pretraining weight를 계승
• 연속적인 로봇 제어 출력을 생성하도록 학습
• Demonstration dataset: $D = (o_t, l_t, a_{t:t+H})$에 대해 behavior cloning
• 관측, 언어 명령을 미래 행동 궤적으로 매핑: $\pi_\theta(a_{t:t+H} | o_t, l_t)$ 행동
• 이산적 행동 토큰: next-token prediction. FAST 행동 토큰 예측
• 연속적 행동 토큰: flow-matching objective function. 연속 행동을 복원하는 action expert

$\pi0.5$: Objective function(subtask 예측, 객체 탐지, VQA) + co-training하여 일반화 성능이 향상됨

• subtask: 시각 관측과 상위 언어 명령이 주어졌을 때, subtaskd 문자열을 예측 $p(l_t^{subtask} | a_t, l_t)$
• 이후 모델에 다시 입력되어 행동 생성 조건으로 사용 $\pi_\theta(a_{t:t+H} | a_t, l_t)$

  4. Fine-tuning with Emergent Human-Robot Alignment

  최종 목적

• Explicit alignment없이 활용
• 인간 설계 휴리스틱보다 대형 모델이 다양한 데이터 소스로부터 관련 정보를 흡수할 수 있는 능력에 의존

  Human Data Collection Pipeline

  

Data Collection Device

• 고해상도 head-worn camera
• 2개의 손목 장착 카메라: end-effector와 조작 대상 간의 상호작용을 자세히 관찰할 수 있음

Data Collection Protocol

• 인간 데이터를 에피소드 형태로 수집
• 인간과 로봇 사이의 차이를 시각적, 운동학적 차이로만 제한할 수 있음
• 수집자들은 각 과제에 대해 반복적인 시연
• 카메라 시야 안에서 유지되도록 요청

Data processing & Annotation

• 시점 t에서 head-worn camera의 6D 움직임 $e_t \in \mathbb R^6$(세계 좌표계 기준)

  6D: Translation 3개(x, y, z) + Rotation 3개(Roll, Pitch, Yaw)

• 두 손 카메라 17개 3D Keypoints $h_t^e \in \mathbb R^{3 \times 17}$ (head-worn camera 기준: 카메리를 원점으로 하는 좌표계)
• 팔의 동작을 설명하는 텍스트 기반 subtask 주석

Action Space

• 인간-로봇 행동 표현을 유사하게 정렬
  • Joint-position, End-effector pose trajectory 이 중 Joint-position은 정확한 추정이 어렵기에 End-effector based trajectory
• Action Chunk(길이 H): $[a_0, a_1, ..., a_H]$ 이때 각 $a_i$는 현재 상태의 6DoF pose
• 전체 Action Space
  • 왼팔: 6DoF + Gripper(1차원)
  • 오른팔: 6DoF + Gripper
  • 베이스 이동: 2 dimension(x, y 이동)
  • 총합: $a \in \mathbb R^{H \times 16}$

Training Objectives

• low level action prediction
  • Discrete Fast tokens에 대한 next-token 예측
  • Continuous action에 대한 flow matching loss
  • $\pi_\theta(a_t | o_t, l_t^{subtask})$
• high level subtask prediction
  • subtask language tokens에서 next-token을 예측
  • $\pi_\theta(l_t^{subtask} | o_t, l_t)$

Training Mixture

• fine-tuning시에 기존 능력을 유지 + 인간 데이터의 새로운 개념을 도입
• simple recipe 사용
  • human data for generalization task: 50%
  • 가장 유사한 robot task: 50%
• 위의 혼합으로 $\pi_{0.5}$를 fine-tuning하여 egocentric data를 통합한 모델을 $\pi_{0.5} + ego$라고 지칭

Experimental Findings

Benchmark

목적: 로봇 데이터에는 없고 인간 데이터에만 존재하는 개념을 로봇이 수행할 수 있는가

Scene Transfer(Spice, Dresser)

• 로봇 데이터: 특정 몇 개의 집에서만 task 인간 데이터: 완전히 새로운 집에서 동일한 task
• 평가: 새로운 집에서 작업 성공 여부

Object Transfer(Bussing)

• 로봇 데이터: 특정 종류의 물체 정리 인간 데이터: 새로운 물체 유형 포함
• 평가: 올바르게 정리한 물체 개수

Task Generalization(Eggs)

• 로봇 데이터: 달걀을 상자에 넣기만 학습 인간 데이터: 색깔별 달걀 분류라는 새로운 개념 포함
• 평가: 올바르게 분류한 달걀 개수

Result

위의 3가지 벤치마크에 대해서 모두 성능이 향상된 모습을 보임

해석

• 기존의 robot-only($\pi_{0.5})$는 기본적인 조작 능력만 존재하고 새로운 의미 구조를 파악하지 못함
• 인간 데이터와 co-finetuning하면($\pi_{0.5}+ego$)하면 인간 데이터의 개념을 실행 가능
• 인간 데이터는 단순 보조 데이터가 아닌 새로운 능력을 실제 행동으로 전이

주장

충분히 다양한 pretraining이 존재할 때 인간 데이터를 제대로 활용할 수 있다.

실험 설계

• 0%: 기본 VLM

• 25%, 50%, 75%, 100%: 점점 더 다양한 로봇데이터

• 100% + X-Emb: 여러 로봇 Embodiment 포함

  X-Emb(Cross-Embodiment): 서로 다른 몸 구조 간 학습(인간 - 로봇, 서로 다른 로봇)

  결과

• pre-training이 적은 경우(0%, 25%): 인간 데이터를 추가해도 성능 향상 없음

• 중간 수준(50%): 일부 task에서만 효과가 나타남

• 충분히 다양한 경우(75%): 인간 데이터 효과가 급격하게 증가함

• 최대 다양성(100% + X-emb): Transfer 효과가 가장 강함

  작업별

  

• Sort Eggs: 인간 데이터가 있을 때만 급격히 성능 상승. $\rightarrow$ 새로운 의미 구조는 인간 데이터에만 의존함

• Dresser: 50% 이하에선 인간 데이터가 오히려 방해가 됨

해석

t-SNE 분석

• 사전학습이 적으면 인간과 로봇의 표현이 분리됨
• 다양성이 증가하면서 두 표현이 겹치기 시작함

즉, 인간에서 로봇으로 transfer은 단순 데이터 추가 효과가 아닌 규모 의존적인 창발 현상임

인간 데이터의 가치

Upper Bound: Target Robot(ARX)과 비교

• Sort Eggs, Dresser: 인간 데이터 성능 $\approx$ target robot 데이터 성능
• Bussing: 인간 데이터 성능 $<$ target robot 데이터 성능

  다른 로봇 embodiment data(UR5)

  
• Bussing: 인간 to ARX transfer $\approx$ UR5 to ARX
• 인간 데이터는 다른 embodiment data처럼 작동함

  결론

  일반적으론 다른 로봇 embodiment와 유사한 수준의 Transfer을 보이고, 일부 작업에선 Target Robot Data에 근접한 성능을 보여줌

High-Level과 Low-Level에서 Transfer

High-Level(HL)과 Low-Level(LL)

HL: subtask prediction LL: 실제 행동의 trajectory 예측(6DoF 등)

실험

1. Robot-only HL + Robot-only LL
2. Human data HL + Robot-only LL
3. Robot-only HL + Human data LL
4. Human data HL+LL

   Result

• HL만 인간 데이터 사용: LL이 명령을 제대로 따르지 못함(잘못된 위치에 배치)
• LL만 인간 데이터 사용: HL이 잘못된 subtask 예측(엉뚱한 행동 지시)
• 모두 인간 데이터를 사용한 4번에서 가장 높은 성능
• 즉, 인간 데이터 Transfer은 HL과 LL 모두에서 나타난다.

손목 카메라의 효용성

• Dresser, Bussing에 대해서는 성능이 향상되었고, Spice, Eggs에선 큰 차이가 발생하지 않았다.
• 즉, 정밀한 물체 조작이 중요한 작업에서는 도움이 되지만, 비교적 global한 정보가 중요한 경우에는 큰 차이가 발생하지 않는다.

6. Discussion

• 제한된 pretraining에선 인간 데이터의 Transfer이 실패함
• Pretraining이 임계점을 넘어서면 Transfer이 창발적으로 나타남
• 해당 연구는 VLA 모델 확장에 있어서 유망한 방향을 제시함

뭔가 많이 추가됐다. 장기기억 + 출력 실험용으로 들어가있는 파일들도 있기에 필요한 것들만 보자.

• .env: api키가 저장되어있다. 시스템 변수로 저장해두긴했지만, 일단 별도로 저장도 해두었다.

1. config.py

LLM api를 호출하고 기본 세팅을 하는 파일이다.

from google import genai
from google.genai import types
MEM_PATH = "short_term_memory.json"
LONG_PATH = "integrated.json"

TOP_K = 4

STM_MAX = 10
PROMOTE_START = 5
PROMOTE_END = 10
client = genai.Client()
instruction = (
    "You will receive context in this structure:\n"
    "_user query\n<user_input>\n_short_term memory\n<short_term memory>\n"
    "Each section is labeled and appears on its own line. "
    "Use short_term memory for recent turns. Answer the user query based on the provided context."
)

api_model = "gemini-2.5-flash"

setting_config = types.GenerateContentConfig(
    thinking_config=types.ThinkingConfig(thinking_budget=0), # Disables thinking
    system_instruction=instruction,
    temperature=0.9,
    max_output_tokens=200,  
    stop_sequences=["User:", "사용자:", "\n\n"], 
    )

우선 단기기억(MEM_PATH)와 장기기억(LONG_PATH)의 경로, TOP_K를 몇으로 지정할지가 저장되어있다.

STM_MAX와 PROMOTE_START, PROMOTE_END는 이후 설명할 단기기억을 장기기억으로 넘기는 개수 기준이다. 즉, 단기기억이 10개가 쌓일 경우 5번째에서 10번째까지의 기억을 장기기억으로 넘길 것이다.

instructure에선 short_term을 넘겨주는 사실을 명확하게 해두었고, config는 기존과 크게 달라진점은 없다.

2. memory_format.py

def build_contents(turns, user_input=None):
    contents = []
    turns = turns or []

    for t in turns:
        contents.append(f"User: {t['user']}")
        contents.append(f"Assistant: {t['assistant']}")

    if user_input is not None:
        contents.append(f"User: {user_input}")

    return "\n".join(contents)

build_contents를 STM와 LTM 모두에서 사용하기에 별도로 분리해두었다. 추후에 GPT API에게 전송하기 위해서 User과 Assistant를 명시적으로 표시해두었다.

3. short_term.py

import os
import json
from datetime import datetime, timezone


def load_memory(path):
    if not os.path.exists(path):
        return []

    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, list):
            return []
        turn = []
        for t in data:
            if not isinstance(t, dict):
                continue
            if "user" not in t or "assistant" not in t:
                continue
            created_at = t.get("created_at")
            turn.append({
                "id": len(turn),
                "user": str(t["user"]),
                "assistant": str(t["assistant"]),
                "created_at": created_at,
            })
        return turn
    except (json.JSONDecodeError, OSError):
        return []


def save_memory(path, memory):
    tmp_path = path + ".tmp"
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(memory, f, ensure_ascii=False, indent=2)
    os.replace(tmp_path, path)


def new_turn(user_input, assistant_reply):
    return {
        "id": 0,
        "user": user_input,
        "assistant": assistant_reply,
        "created_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }

단기기억을 저장하는 파일이다. 대부분 전에 만들었던 부분과 동일하고, created_at이라는 해당 QA가 만들어진 시간을 저장하는 부분을 추가했다.(추후에 사용할 예정)

4. llm_api.py

from google import genai
from google.genai import types
import memory_format
from config import *

def call_gemini(user_input, memory, long_term_turns=None):
    recent_memory = memory[-TOP_K:]
    sections = []

    user_section = "_user query\n" + user_input

    short_term_contents = memory_format.build_contents(recent_memory)
    short_term_section = "_short_term memory\n" + (short_term_contents if short_term_contents else "(empty)")

    sections.append(user_section)
    sections.append(short_term_section)

    user_content = "\n".join(sections)

    response = client.models.generate_content(
        model=api_model,
        contents=user_content,
        config=setting_config,
    )

    return response.text.strip()

실제 api를 호출하는 구간이다.

5. main.py

from config import MEM_PATH, LONG_PATH, STM_MAX, PROMOTE_START, PROMOTE_END, TOP_K
import short_term
import llm_api
import long_term
import stm_to_ltm

def main():
    memory = short_term.load_memory(MEM_PATH)

    print("대화 시작 (종료는 Ctrl+C) \n")

    try:
        while True:
            user_input = input("Gemini에게 물어보기: ").strip()
            if not user_input:
                continue


            reply = llm_api.call_gemini(user_input, memory)
            print(f"Gemini: {reply}\n")

            new_turn = short_term.new_turn(user_input, reply)
            new_turn["id"] = len(memory)
            memory.append(new_turn)
            short_term.save_memory(MEM_PATH, memory)
            memory = stm_to_ltm.promote_stm_to_ltm(
                MEM_PATH,
                LONG_PATH,
                max_stm=STM_MAX,
                promote_start=PROMOTE_START,
                promote_end=PROMOTE_END,
            )

    except KeyboardInterrupt:
        print("\n종료")

if __name__ == "__main__":
    main()

앞의 부분들을 통합 + 다음에 설명될 장기기억으로 변환을 통합해둔 파트이다.

STM to LTM

이제 단기기억이 일정 이상 쌓이면 장기기억으로 변환해주는 파트를 만들어볼 것이다. 3가지 함수로 구현해보았다.

_next_ltm_index

def _next_ltm_index(items):
    if not items:
        return 0
    last = items[-1]
    item_id = last.get("id", "")
    if isinstance(item_id, int):
        return item_id + 1
    if isinstance(item_id, str):
        if item_id.isdigit():
            return int(item_id) + 1
    return len(items)

STM이 아닌 LTM에서 해당 QA의 index를 구하는 파트이다. 현재 LTM에는 항상 오름차순으로 기억이 추가만 된다고 가정해둔 상태이다.(즉, 기억의 중간 삽입이나 삭제가 존재하지 않는다.) 따라서 가장 마지막 item의 id값 + 1로 현재 기억의 id값을 결정한다.

append_to_long_term

def append_to_long_term(integrated_path, turns):
    """
    turns: list of {id, user, assistant, created_at}
    """
    if not os.path.exists(integrated_path):
        data = {"items": []}
    else:
        try:
            with open(integrated_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            if "items" not in data or not isinstance(data["items"], list):
                data = {"items": []}
        except (json.JSONDecodeError, OSError):
            data = {"items": []}

    next_index = _next_ltm_index(data["items"])
    for t in turns:
        record = {
            "id": next_index,
            "user_query": t["user"],
            "answer_query": t["assistant"],
            "created_at": t.get("created_at"),
        }
        data["items"].append(record)
        next_index += 1

    tmp_path = integrated_path + ".tmp"
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    os.replace(tmp_path, integrated_path)

파일에 ltm을 추가하는 파트이다. integrated_path(LTM 경로)의 파일을 열고, 위의 _next_ltm_index에서 가져온 id와 STM에서 필요한 정보들을 가져와서 그대로 LTM에 넣는다. 데이터 안정성을 위해서 tmp에 저장해두었다가 모두 완료된 이후 원래의 파일에 덮어쓰는 방식을 사용한다.

promote_stm_to_ltm

def promote_stm_to_ltm(stm_path, ltm_path, max_stm=10, promote_start=5):
    turns = short_term.load_memory(stm_path)

    if len(turns) <= max_stm:
        return turns

    excess = len(turns) - max_stm
    promote_count = promote_start + excess + 1
    if promote_count < 0:
        promote_count = 0
    if promote_count > len(turns):
        promote_count = len(turns)

    promote = turns[:promote_count]
    remaining = turns[promote_count:]

    append_to_long_term(ltm_path, promote)

    for i, t in enumerate(remaining):
        t["id"] = i

    short_term.save_memory(stm_path, remaining)
    return remaining

stm쪽을 정리해주는 함수다. stm의 최대 개수인 max_stm(여기선 10)을 초과한 개수를 구하고, 가장 오래된 0번째 기억부터 (promote_start(여기선 5) + 초과된 개수)까지의 기억을 ltm으로 넘긴다. 이후 남은 stm들의 id를 재정리해준다.

Conclusion

테스트를 해보고싶지만 STM에 어떤 데이터를 넣을지를 고민중이다. 테스트한 이후에 추가할 수 있다면 추가하도록 하겠다.

1. Introduction

VLA Pre-training

• VLA(Vision-Language-Action): 카메라나 센서 데이터 등 시각(Vision)정보와 자연어 명령(Language)을 통해 실제 환경에서 특정한 작업(Action)을 수행하도록 설계된 AI모델.

일반적으로 VLA데이터셋은 실험실 환경에서 teleoperation으로 수집된다. 이는 데이터의 가치는 매우 높지만, 수집 비용이 너무 높다는 단점이 존재한다.

웹 데이터의 사용

웹에는 일상적인 행동과 그에 따른 물리적 상호작용을 다루는 많은 실제 인간 영상이 존재한다. 그러나 이런 데이터들은

• 대본이 존재하지 않음
• 분할되어 있지 않음
• 길이와 작업 세분성(작업을 나누는 단위)가 제각각임
• 노이즈가 많고 목적에 관련 없는 행동도 많음
• 언어적 지시나 3D 행동 레이블도 없음

이런 웹 데이터를 VLA 학습 데이터로 사용하기 위해, 사람의 손을 AI 모델의 End-Effector로 간주하고,

• Task alignment: 기존 로봇 데이터의 구성 방식을 따름과 동시에 원자 수준의 단기작업을 분할하고 필터링해야 한다.
• Label alignment: Dense Action Label을 제공하기 위해서 Metric Space 상에서 3D 손 움직임을 복원해야 한다. 실제 웹 데이터는 한 개의 카메라를 보정없이 움직이며 촬영하기에 난이도가 높다. 추가로 행동을 설명하는 Language Instrcution Label도 필요하다.

2. Related Works

VLA model pretraining

• VL은 pretraining module이나 backbone을 통합: 대부분 OXE(Open X-Embodiment) 데이터 셋 활용.
• Action은 proxy(대체 표현)을 활용함: latent action을 비지도 학습하고, 이를 proxy로 활용. 영상의 미래 프레임을 예측 대상으로 활용하기도 함.

Egocentric한 인간 영상에서 3D 손행동 레이블을 사용하여 사전학습을 진행하기도 하지만, 이는 통제된 환경에서 수집된 손-물체 상호작용 영상에만 의존한다.

Dexterous Hand Manipulation

기존에 손가락 로봇 손은 시뮬레이터에서 강화학습을 통해 학습되었다. 하지만 이는 보상 설계가 복잡하고, 실제 환경 적용성이 제한된다. 인간의 원격조작 시연 데이터를 이용한 imitation learning을 이용한 연구들도 존재하지만, 이는 소수의 작업에 특화된 모델이다.

Robot Learning from Human Videos

Egocentric한 인간 영상을 활용해 Vision 및 Language 표현을 학습한다.

• 모션 캡쳐나 웹 영상의 명시적인 인간 행동을 Imitation Learning Framework에서 로봇 정책 학습에 사용
• 명시적 움직임 대신 Affordance(객체가 제공하는 가능한 행동 가능성에 대한 정보), Point trajectories(영상 속 특정 점이 시간에 따라 이동하는 경로), hand-object mask(손과 물체 영역을 픽셀단위로 분할한 정보)를 학습한다.

  Temporal Action Segmention for Videos

  긴 인간 영상에서 행동 구간을 탐지하고 분류. 초기에는 정의된 행동 클래스에 대해서만 초점을 맞췄지만, 현재는 광범위한 행동 이해 능력을 갖춤. 하지만 여전히 정확도 측면에서 어려움이 존재함.

3. Transforming Human Hand Video to VLA Data

VLA Data: 간단하고 단기적인 작업으로 구성됨

• High-level Planner: 원자단위의 행동을 목표를 위해 순서대로 계획하는 시스템. 이를 이용해 VLA Data의 단기 작업이 장기 작업으로 조합됨
• 각 data episode는 언어 지시, 영상 프레임 시퀀스, end-effector의 3D Action chunk를 포함함

3D Motion Labeling

목표: 영상에서 두 손과 카메라의 움직임을 포함한 3D 움직임을 추출

1. Background Optical Flow: 카메라가 정지 상태인지 이동 중인지를 판별하는 알고리즘
2. 카메라 내부 파라미터(intrinsics) 추정

• Moving인 경우 DroidCalib: 움직이는 영상의 내부 파라미터(초점 거리, 렌즈 왜곡, 카메라 중심 위치 등)추정
• Static인 경우 MoGe-2(monocular 영상에서 깊이 및 기하 구조 추정), DeepCalib(정지 카메라 영상의 내부 파라미터 추정)
• 왜곡이 큰 경우 왜곡 보정(Undistortion)

3. 손 재구성 및 카메라 자세 추정

• 손 재구성(HaWoR): 프레임별 카메라 좌표계 3D 손 복원, 손목의 6D 자세와 MANA hand parametric model 기반의 관절 각도를 포함함
• 카메라 자세 추정(MegaSAM): MoGE-2의 깊이 데이터를 바탕으로 깊이 추정 기반 SLAM 사용

4. 3에서 얻은 두 정보를 활용해 word-space의 3D 손 시퀀스를 얻은 후 spline smoothing과 outlier removal을 적용

spline smoothing: 노이즈가 포함된 데이터 점들을 부드러운 곡선으로 보정. interpolation과 달리 모든 점을 정확히 통과하지 않고 전체 추세를 부드럽게 근사

Atomic Action Segmentation

목표: 긴 영상에서 원자 수준 행동 시퀀스를 분리 실제 인간 손 행동의 자연스러운 beats를 이용: 사람은 행동 전환 시 손 속도 변화가 발생하며, 속도의 최소값이 행동의 전환 지점인 경우가 많음.

• world space 손목 속도 최소값을 탐지해서 이를 분할 지점으로 사용
• hand trajectory를 smooth해서 일정한 window 내 local speed minima를 선택

이때 좌우 손은 독립적으로 분할하고, 한 손의 움직임을 분석할 때 다른 손의 움직임은 무시한다. 추가 모델 추론이나 사전 텍스트 라벨 없이도 계산 효율성이 높다. 일부 행동에서는 over-segmentation이 발생할 수 있지만, 이후 언어 라벨링에서 병합된다.

Instruction Labeling

위 과정들을 통해 얻은 영상 분할과 3D 손 행동 시퀀스를 활용해 시각화 자료를 생성하고, GPT-4.1을 활용해 행동 캡션을 생성

• 손 움직임 분석
• 의미 있는 조작 여부 판단: 의미 없는 행동은 N/A로 라벨링
• 자연어 명령형 설명 생성

이때 고정 길이 분할보단 원자 수준의 영상 클립이 캡션 정확도가 더 높다.

Hand V-L-A Dataset Construction

Ego4D, Epic-Kitchen, EgoExo4D, Something-Something-V2 등 Egocentric한 인간 영상 데이터셋을 처리하여 대규모 인간 손 V-L-A Dataset을 구축 기존의 라벨링은 모델의 요구 조건과 일치하지 않기에 사용하지 않는다.

4. Dexterous Hand VLA Model

VLA model $\pi$, 현재 시간 관측 $o_t$, 로봇의 proprioceptive state $s_t$, language instruction $l$, 미래의 end-effector 행동 시퀀스 $a$ $\pi : (l,o_t,s_t) \rightarrow (a_t, a_{t+1}, ..., a_{t+N})$

4.1 VLA 모델 설계

Model Architecture

VLM: PaliGemma-2(SigLIP encoder + Gemma-2 Language model)

• FoV정보를 추가 토큰으로 입력(이미지 종횡비, 카메라 내부 파라미터 해석)
• 학습 가능한 cognition token을 VLM입력에 추가: VLM의 출력이 매우 많고, Action expert에게 요약을 전달해주기 힘들기에 특정 장면을 한 벡터로 요약하는 토큰을 만듦
• cognition token을 통해 생성된 하나의 특징 벡터 $f_c$를 Action expert condition으로 사용: Diffusion기반 행동 생성에서 denosing 과정 전체가 해당 지시에 맞는 행동으로 유도되도록 만듦

Action Expert: Diffusion Transformer(DiT)

• 입력: $f_c$, 손의 상태 $s_t$, noise가 추가된 action segment($a^i_t, a^i_{t+1},...,a^i_{t+N}$)
• $s_t$: 현재 영상에서 카메라 좌표계 손목 위치 및 회전, 손 관절 각도
• 행동의 유효성 여부를 나타내는 action mask도 입력
• $f_c$는 AdaLN을 통해 주입되어 조건 정보 강화
• 학습: $\mathcal L_{MSE}=E_{\epsilon \sim N(0,1),i}||\hat\epsilon_i - \epsilon||^2$, 이때 $\hat\epsilon$은 예측 노이즈, $\epsilon$은 실제 노이즈

Action expert, VLM, cognition token은 end-to-end로 학습되고, vision encoder은 frozen

Hand Action Space

시점 $t$에서 행동 $a_t$, $a_t = [\Delta t_l,\Delta r_l, \theta ^h_l, \Delta t_r, \Delta r_r, \theta^h_r] \in \mathbb R^{102}$ 이때

• $\Delta t \in \mathbb R^3$: 연속 프레임 간 손목 위치 변화
• $\Delta r \in \mathbb R^3$: 손목 회전 변화(Euler angle)
• $\theta_h \in \mathbb R^{15 \times 3}$: MANO 기반 15개 관절 각도
• l, r은 왼손, 오른손

  Unified Single and Dual-Hand Action Prediction

  VLM은 항상 Left hand: <왼손 행동> Right hand: <오른손 행동> 과 같은 형식으로 언어 지시를 입력받음 만약 해당 손 행동이 없다면 None

Action expert는 항상 양손 행동 노이즈를 입력받음 한 손 레이블만 존재할 경우 action mask를 action dimension에 맞춰 추가입력받는다. mask가 0이면 해당 행동은 0으로 설정 후 손실 계산에서 제외한다.

Causal Action Denoising

실제 인간의 손은 빠르게 움직이면서 pre-training data action clip은 약 1초 수준(30 frame)이기에 예측 길이를 16 frame정도로 잡게 되면 padding이 필요하다. 이때 단순히 0으로 패딩하면 부자연스럽기에 action denoising에 causal attention을 적용하고, 각 행동 토큰이 이전 행동만 참조하도록 제한한다.

4.2 Pretraining with Human Hand VLA Data

목표: 이전에 구축한 데이터셋을 이용하여 Pre-training 수행 일반화 능력 향상을 위해 trajectory-aware augmentation 적용

trajectory-aware augmentation: 일반적인 data augmenation을 적용하면 좌표와 이미지 간 불일치가 발생하기에, 이미지 변환과 동시에

• 손의 3D위치 변환
• 카메라 파라미터 재계산
• 행동 벡터 변환
  을 동시에 수행한다.

구체적으로

• random cropping
• 시야각, 종횡비, 중심 변화 perspective 변환
• 카메라 중심 유지
• trajectory crop 영역 내부 유지
• 이미지 좌우 반전 및 대응 행동, 언어 수정
• 색상 정보가 없다면 color jitter 적용

4.3 Fine-tuning for Robotic Dexterous Manipulation

사전학습 후 로봇 데이터로 fine-tuning을 진행하여 실제 적용시킨다.

• 로봇 end-effector 6D 자세를 활용해서 $\Delta t$와 $\Delta r$을 계산

• 로봇 손의 관절을 사람 손 관절 topology 기반($\theta_h$)에 대응하여 매핑

• 대응이 안된다면 mask 0으로 패딩

  5. Experiment

  Pretraining Data Analysis

  다양한 실제 생활 활동을 데이터로 이용할 때

• Visual Diversity: 실험실에 한정된 데이터들에 비해서 실제 환경 데이터가 더 높은 장면 다양성을 보이고 일반화에 더 유리한 모습을 보임 (h: 최소 h번 이상 등장한 단어가 h개 이상 존재할 때 최대 h값 / i100: 100번 이상 등장한 서로 다른 단어 개수)

• Instruction Diversity: GPT-4.1을 활용해 지시문에서 명사, 동사, 형용사 분포를 확인했을 때 언어 표현 범위가 더 넓은 모습이 관찰됨

  Human Hand Action Prediction

  Pretraining된 VLA 모델이 보지 못한 환경에서 사람 손 행동을 제대로 예측하는가

• 기존 모덷를보다 더 정확한 행동 예측이 가능했음

• 데이터 다양성이 높을수록 성능이 증가

• Explicit 3D 행동 예측 방식이 latent action보다 더 안정적임

  Real-World Robot Dexterous Manipulation

  실제 로봇에 적용시킨 결과

• unseen setting에서 기존의 모델들에 비해 우수한 일반화가 이루어짐

• 인간 손 VLA데이터가 다양할수록 실제 로봇 성능의 향상으로 이어짐

• 인간 손 행동 예측 정확도와 로봇 조작 성공률 간 명확한 양의 상관관계가 보임

  6. Discussion / 7. Conclusion

  정해진 실험실 환경에서의 데이터 뿐 아니라 Egocentric한 인간 활동 영상을 활용하여 VLA model을 pretraining 할 수 있는 방법에 대한 설명이다. 실제로 이런 방식으로 pretraining된 모델들은 unseen환경에서 강력한 zero-shot 성능을 보임.

1. Introduction(Background)

Text-embedding: 임의 길이의 텍스트를 저차원 벡터로 표현. lexical mismatch를 극복함과 동시에 효율적인 검색과 매칭을 가능하도록 한다.

BERT(Encoder-only)나 GPT(Decoder-only)도 전이(Transfer) 가능한 표현을 생성할 수는 있지만, 효율성 및 범용성에서 대조학습(Contrasitve Learning)을 바탕으로 하는 GTR이나 Sentence-T5가 더 이상적이다.그러나 이런 방식은 품질이 낮은 경우가 많기에 BM25 기준 모델 성능을 넘지 못한다.

2. Related Work

초기 연구들

Latent Semantic Indexing(LSA) + Latent Dirichlet Allocation(LDA)

• LSA: 단어-문서 co-occurrence matrix를 분해하여(SVD) 문서 임베딩을 생성

  co-occurrence matrix: 문서 집합에서 각 단어가 각 문서에 얼마나 등장하는지를 수치로 정리한 행렬

• LDA: 확률 그래픽 모델을 이용해 주제 분포 학습

Sentence Embedding

짧은 텍스트에 초점을 맞춰 언어모델을 fine-tuning하여 연속적인 임베딩 출력

SimCLR: contrastive loss가 기존의 분류 기반 손실보다 더 효과적 LaBSE, LASER, CLIP: 병렬 문장, 이미지-텍스트 쌍을 활용해 다국어 및 멀티모달 환경으로 확장 ICT: 문단 내 임의 문장을 pseudo-query로 선택하고 나머지를 positive sample로 간주 Contriever: Data Augment를 활용한 Random cropping이 ICT보다 효과적 OpenAI Text Embedding: 인접 텍스트를 Positive sample로 간주, 모델 규모 확장 Oguz: 도메인 일치 사전학습으로 특정 도메인의 성능 개선 SPAR: BM25를 Teacher model로 dense retriever 학습. 그러나 이런 방식으로 얻은 합성 데이터는 품질이 낮은 경우가 많고, 라벨도 따로 fine-tuning 해줘야 한다.

3. CCPairs

웹 기반 대규모 고품질 텍스트쌍 데이터셋

semi-structured data sources

query와 passage로 구성된 (q, p) 텍스트 쌍.

• Reddit의 게시글-댓글 쌍
• Stackexchange의 질문-답변 쌍
• en Wikipedia의 (엔티티 이름+제목)-본문 쌍
• 과학 논문의 제목-초록/인용 관계 쌍
• Common Crawl 웹페이지 및 뉴스 데이터의 제목-본문 쌍

대략 13억 개 텍스트 쌍

Consistency-based filter

위에서 얻은 노이즈가 포함된 텍스트 쌍으로 모델을 학습한 뒤, 학습된 해당 모델을 사용하여 각 텍스트 쌍을 100만 개 랜덤 문서 풀과 비교해서 ranking 작업을 한다. 이때 텍스트 쌍이 top-k(k=2) 안에 포함될 때만 유지하는 방식으로 모델 예측과 학습 라벨이 일관성일 보이는 데이터만 유지한다.

이는 신경망의 memorization behavior에서 나온 방식

memorization behavior: noise label이 섞인 학습에서 초반에는 상대적으로 맞는 예제부터 먼저 잘 맞추는 경향이 존재하지만, 모델 용량이 충분해지면 noise까지 overfitting되어 잘못된 라벨만 존재하는 현상

4. Method

Contrastive Pre-training with Unlabeled Data

Contrastive pre-training: 관련 있는 데이터 쌍을 관련 없는 데이터 쌍과 분리

식을 먼저 살펴보면, 텍스트 쌍 집합: ${(q_i, p_i)}^n_{i=1}$ i번째 예제에 대한 Negative Passage 집합: ${p_{ij}^-}^m_{j=1}$ InfoNCE contrastive loss: min$L_{cont} = -{{1}\over{n}}\displaystyle\sum_i \log{{e^{s_\theta(q_i, p_i)}}\over{e^{s_\theta(q_i, p_i)}+\sum_j e^{s_\theta (q_i, p_{i,j}^-)}}}$ $s_\theta (p, q) = \cos(E_q, E_p) / \tau$

. $s_\theta(q,p)$부터 살펴보면, 파라미터 $\theta$로 표현되는 query q와 passage p 사이의 점수 함수로 $E_q$와 $E_p$는 사전 학습된 Transforemr encoder + Average Pooling을 적용한 고정 크기의 텍스트 임베딩이다. 여기에 temperature hyperparameter($\tau$, 확률 분포의 날카로움을 조절하는 계수, 해당 논문에선 0.01)와 코사인 유사도를 적용하여 계산한다.

이때 인코더는 동일하되 query와 passage 접두사를 추가한다(query와 passage가 명확하지 않은 경우는 랜덤). 이는 대칭성을 깨기 위한 방법으로 이런 비대칭 설계가 paraphrase가 존재하는 검색에서 중요하다.

Negative Sample의 경우에는 in-batch negative 방식을 사용한다.

in-batch negative: 같은 batch 내에서 다른 텍스트쌍의 passage를 전부 negative sample로 사용하는 방법

Fine-tuning with Labeled Data

Labeled Data를 추가하여 임베딩 품질을 향상시킨다. 해당 논문에서는

• MS-MARCO, NQ(Natural Question): Retrieval 전이. mined hard negative와 CE teacher model로부터 knowledge distillation을 사용
• NLI(Natural Language Inference): STS 및 linear probing 전이. contradiction을 hard negative로 활용.

Loss function: hard label에 대한 contrastive loss인 $L_{cont}$와 KL divergence(teacher모델의 soft label 증류)의 결합 min $D_{KL}(p_{ce}, p_{stu}) + \alpha L_{cont}$

• $p_{ce}$: cross-encoder teacher 모델 확률
• $p_{stu}$: 학생 모델의 확률
• $\alpha$: 두 loss간 균형을 맞추는 hyperparameter

  Text embedding 작업에 적용

• Zero-shot Retrieval: 대상 corpus의 문서 임베딩을 사전에 계산 후 query에 대해 query 임베딩만 계산 + 상위 k개 반환
• Few-shot Text classification: 고정된 임베딩에서 소량의 라벨 데이터를 이용해 선형 분류기 학습. 작업마다 분류 헤드 파라미터만 학습 + 저장
• Zero-shot Text Classification: 입력, 라벨 텍스트를 사람이 작성한 프롬프트 템플릿으로 문장 형태로 변환. 이후 임베딩 공간에서 가장 가까운 라벨을 예측 결과로 선택
• Semantic Texture Similarity: 두 텍스트 임베딩 간 코싸인 유사도를 이용해서 의미적 유사도 측정. 보통 rank correlation coeffcient로 평가
• Text Clustering: 표준 cluster알고리즘 그대로 적용

5. Experiment

결과

• 비지도학습: 라벨 없이도 BM25 기준 모델 능가. CCpairs가 성능 향상에 크게 기여함
• 지도 fine-tuning: 대부분 데이터셋에서 성능이 추가 향상됨. 기존의 임베딩 모델보다 높은 평균.
• MTEB benchmark: 파라미터가 훨씬 많은 모델들과 경쟁력이 유지됨

추가 분석 결과

• 배치가 클수록 negative sample이 증가하며 성능이 개선됨
• MS-MARCO + NQ는 검색 작업, NLI는 의미 유사도, 분류 작업에 유리함. 둘을 결합시킬 때 가장 균형 잡힌 성능을 보임
• 데이터 필터링: 노이즈 제거 시 성능 향상, 학습 효율 개선
• Negative Sampling: in-batch negative가 가장 안정적
• BM25, dense retrieval: dense모델이 발전해도 BM25의 장점이 여전히 존재

6. Conclusion

기본적으로 Sentence-T5를 따라가되,

• CCPairs라는 대규모 weak supervision 데이터 + consistency filtering + 소량의 supervised fine-tuning
• prefix trick(query, passage)로 비대칭을 주어 성능 향상

즉, Dense-only Embedding으로 범용성에 초점이 맞춰진 모델이다. 이전의 M3-Embedding의 경우에는 Sparse + Dense + Multi-Vector로 대규모 RAG에 더 적합하고, 대화형의 경우에는 lexical overlap 정보가 적고 token interaction benefit도 적기에 E5 + 낮은 수준의 Sparse면 된다고 판단했다.

1. Introduction

Sentence Embedding: 문장의 의미를 압축된 형태로 표현하여 분류, QA, Semantic retrieval 등 다양한 언어 처리 작업에 사용 가능

최근 연구들에선 모델의 파라미터를 늘리고, pre-trained된 모델을 사용하는 것이 성능 향상에 효과적인 접근법임을 보여준다.

해당 논문에선 T5에서 Sentence Embedding을 추출하는 방법에 대해서 탐구한다. 기본적으로 T5는 span corruption 방식을 사용하기에 파라미터의 확장이 가능하지만, retrieval과 clustering에서 효율적인 적용이 어렵다.

span corruption: 문장에서 연속된 구간을 통째로 masking

• cross-attention inference cost: Query 수 N과 Candidate 수 M에 대해, 검색 후보마다 decoder을 추론하기에 $O(N \times M)$이란 계산량
• 의미적 유사성 기반 임베딩은 학습하지 않음
• encoder-decoder 구조의 한계: 원래 임베딩용이 아닌 문장 이해(encoding) 후 생성(decoding)이 목적이기에 의미를 반영하기 위해선 encoder 결과만이 아닌 decoder interaction도 포함되어야 함

해당 논문에서는 3가지 방식을 시도한다. a) Pretrained T5 b) encoder의 첫번째 토큰 표현을 임베딩으로 사용 c) encoder의 모든 토큰 표현의 평균을 임베딩으로 사용 d) decoder의 첫번째 토큰 표현을 임베딩으로 사용

2. Sentence T5

Model Architecture

위 방식의 b를 Encoder-only first(ST5-Enc fisrt), c를 Encoder-only mean(ST5-Enc mean), d를 Encoder-Decoder first(ST5-EncDec first)라고 명명한다.

ST5-Enc first와 ST5-Enc mean의 경우에는 BERT와 같은 encoder-only pre-trained model의 풀링과 유사하지만, CLS 토큰이 존재하지 않는다.

ST5-EncDec first에선 Encoder-Decoder 구조에서 decoder가 첫 토큰을 생성할 때 전체 입력 문장의 의미를 반영하고 있다고 가정하기에 첫 토큰을 임베딩으로 활용한다.

Encoder의 학습에는 Dual Encoder 구조를 사용한다.

Dual encoder: 두 입력 문장을 동일한 인코더(sharing weight)로 각각 고정 길이 벡터로 인코딩하여 Projection layer와 L2 정규화를 거친 후 유사도를 계산

Constrastive Learning(대조 학습)

데이터들 간의 관계를 통해서 학습을 진행하는 방식. 이를 이용하면 임베딩 공간에서의 uniformity가 개선되어 downstream task 성능이 향상된다.

Constrastive Loss

paired dataset: $\mathcal D = {(v_i, v_i^+)}$ 이때 $v_i$는 입력 문장, $v_i^+$는 의미적으로 유사한(Positive) 문장

학습시에는 $v_i$와 $v_i^+$와의 거리는 줄이면서, $v_i^-$와의 거리는 멀게 임베딩을 배치해야한다. 이를 in-batch sampled softmax로 나타내어

$\mathcal L = {{e^{sim(v_i, v_i^+)/\tau}}\over{\sum_{j \in \mathcal B} e^{sim(v_i, v_j^+)/\tau}}}$ 만약 Negative sample $v_i^-$가 있다면, $\mathcal L = {{e^{sim(v_i, v_i^+)/\tau}}\over{\sum_{j \in \mathcal B} e^{sim(v_i, v_j^+)/\tau} + e^{sim(v_i, v_j^-)/\tau}}}$

Two-stage Training

1단계

• 커뮤니티 QA 사이트에서 추출한 질문-답변 데이터로 학습
• 답변의 질문을 positive example로 사용

  2단계

• NLI 데이터로 fine-tuning
• entailment는 positive
• contradiction은 negative

3. Experiment Setup

• 학습 데이터: 1단계 웹 포럼 수집, 2단계 SNLI 데이터셋
• Evaluation: SentEval 벤치마크(transfer, STS 작업)
• Goal Q1. 3가지 방법 중 가장 좋은 임베딩 방법 찾기 Q2. 파인튜닝 없이 T5 sentence 임베딩이 downstream에서 잘 동작하는지 확인 Q3. NLI나 QA같은 constrastive 문장 임베딩 학습이 실제 임베딩 성능에 얼마나 영향을 주는지 확인 Q4. 모델 용량 확장의 영향 확인

4. Result

Q1. 가장 좋은 임베딩 방법

• Raw T5 문장 임베딩에 관해선 ST5-Enc mean방식이 가장 안정적
• STS에서는 embedding collapse로 인해서 오히려 성능이 낮아짐
• Encoder-only: Transfer task에서 강하고, 일반화 성능이 좋음
• Encoder-Decoder: STS task에서 더 좋은 결과, Decoder가 Additional Attention Pooling의 역할도 함

  Additional Attention Pooling: 토큰마다 weight를 달리해서 합치는 방법. 중요한 token을 더 많이 반영한다.

  Q2. Downstream

• Transfer task에선 잘 작동함
• STS 성능이 낮음

  Q3. Contrastive Learning

• NLI기반 contrastive fine-tuning: STS 성능 향상
• QA + NLI 2단계: transfer + STS 모두 개선
• Embedding Geometry 개선에도 효과적

  Q4. 모델의 크기의 영향

• 파라미터가 증가할수록 성능은 꾸준히 향상됨
• Transfer task와 semantic similarity 모두 개선됨
• 모델 크기를 늘리는 것 만으로는 embedding collapse가 완전히 해결되진 못함

1. Introduction

Transfer: 텍스트 이해를 위한 범용적인 지식을 매번 직접 학습시키는 것이 아닌, unlabeled data로 일반적인 언어 지식을 학습하고, 특정 과제를 fine-tuning해서 전이 이런 전이학습의 결과물들은 서로 다른 설정, 평가 기준, 모델 구조로 인해서 개별 기법들의 장단점 파악이 어렵다는 단점이 존재한다. 위의 문제를 해결하기 위해 모든 NLP과제를의 입출력을 모두 텍스트로 변환해서 text-to-text로 만든다.

2. Setup

Transformer 기반, text-to-text 프레임워크(자연어 텍스트 입력 -> 텍스트 출력)

Model

Self-Attention: 시퀀스의 각 원소를 해당 시퀀스 내 다른 원소들의 가중 평균으로 대체하여 시퀀스를 처리한다.

Encoder-Decoder Transformer: 처음 제안된 Transformer 형태를 따른다. Encoder: 입력 토큰 시퀀스 -> 임베딩 시퀀스 -> Encoder로 전달

• 여러개의 block으로 구성된 stack(Self-attention layer + FFD)
• Layer normalization이 각 하위 구성 요소의 입력으로 적용
• Additive bias가 없는 단순화된 형태의 Layer normalization
• Residual skip connection을 사용하여 각 하위 구성 요소의 입력을 출력에 더함
• Dropout: FFD, skip connection, Attention weight, 입/출력에 적용

Decoder: 구조적으로 Encoder와 유사함.

• Encoder의 출력을 참조하는 standard attention도 포함
• autoregressive적인 self-attention을 사용하여 과거의 출력에만 집중하도록 함
• 최종 디코더는 Dense layer로 전달됨. 해당 layer의 가중치는 입력 임베딩 행렬과 공유됨
• 모든 attention은 여러 개의 독립적인 head로 분할ㄹ.

Positional Embedding

• Self-attention은 순서에 독립적인 연산이기에 명시적인 위치 정보(position signal) 제공
• 상대적 위치 기반으로 벡터가 아닌 하나의 스칼라로 위치를 표현

The Colossal Clean Crawled Corpus(C4)

웹에서 수집된 Common Crawl을 활용한다. Common Crawl은 메뉴나 오류 메시지 등 무의미한 문자열이 많기에 전처리 과정을 거친다.

• 끝이 종결 부호인 경우에만 유지
• 문장이 3개 미만인 페이지와 단어 수가 5개 미만인 줄은 폐기
• list of Dirty, Naughty, Obscene or Otherwise Bad Words가 포함된 페이지 제거
• Javascript라는 단어가 포함된 줄 제거
• lorem ipsum과 같은 place-holder text가 존재하면 제거
• 중괄호 제거(코드가 포함된 문서 제거)
• 인용표시 제거
• boilerplate policy 문구 제거
• 세 문장으로 이루어진 연속 구간이 이미 데이터셋 내에 존재할 경우 기존 데이터만 남기고 나머지 제거
• langdetect로 영어로 분류될 확률이 0.99인 페이지만 유지

Downstream Tasks

일반적인 언어 학습 능력을 측정한다.

• 텍스트 분류: GLUE, SuperGLUE
• 요약: CNN/Daily Mail
• 질의응답: SQuAD
• 번역: WMT(이 경우 fine-tuning이 필요)
• 데이터는 Tensorflow Datasets

Input and Output Format

• 모든 태스크를 text-to-text로 변환
• 항상 maximum likelihood를 목표로 학습: teacher forcing 사용(prefix)
• 각 개별 task에 대해서 별도로 fine-tuning이 허용됨
• 명시적 질의 응답 대신 짧은 prefix 사용

Experiments

굉장히 다양한 실험 결과가 나열되어있기에 일부만 본다.

Encoder-Decoder의 우위

Transfer Learning에 있어선 Encoder-only의 BERT나 Decoder-only의 GPT보다 Encoder-Decoder 구조가 더 안정적이다. 이는 구조적으로 입력의 이해(encoder)와 출력 생성(decoder)의 분리가 유리하기 때문이다.

Parameter Sharing의 손해가 적다

파라미터 수가 절반이 되고 FLOPs는 유지되지만, 성능 자체는 거의 동일하게 유지된다.

Denoising > LM

모든 구조와 태스크에서 Denoising Objective가 LM objective보다 더 유리하다. 즉, 자연어 이해에 있어서 다음 토큰을 예측하는 방식보다 손상된 의미 구조를 복원하는 것이 더 일반적인 학습 신호이다.

Span corruption > Token masking

연속 span을 날리고 복원하는 방식이 BERT스타일의 token masking보다 더 유리하다.

prefix LM은 한계가 명확하다

요약, 번역, 복합 생성 task등에서 prefix LM(GPT + prefix attention)은 Encoder-Decoder구조의 대체제가 될 순 없다.

데이터의 양보다는 정제 방식이 중요하다

같은 Common Crawl에서 필터링 규칙에 따라서 성능 차이가 극명하게 난다.

자원이 많다면 전이 학습의 이득이 적다

WMT와 같은 대형 번역 테스크에선 pre-training의 유무 차이가 거의 없다.

Conclusion

일반적인 Transformer에서 모든 NLP문제를 text-to-text로 변환하고 기존에 있던 방식들을 정리하여 어떤 세팅이 더 유리한지를 보여준 논문이다.

1. Reciprocal rank fusion

비지도 학습에서 정보 검색(IR) 방법들의 결과를 결합하기 위한 방법이다.

RRF score

전체 문서들의 집합을 D, 순위 결과 집합을 R이라고 할 때 $RRFscore(d \in D) = \displaystyle\sum_{r\in R}{{1}\over{k+r(d)}}$ 여기서 k = 60이라는 고정된 값으로 설정한다.

해당 식의 목표는 단 하나의 시스템에서만 높은 것이 아닌 여러 시스템에서 상위에 등장하는 문서가 가장 큰 값을 갖도록 하는 것이 목적이다. 단, 이때 하위 랭크의 문서도 무시되지 않도록 해야한다.

따라서 순위가 작을수록(높을수록) 커지도록 역수를 사용하며, 이때 등수간 너무 큰 차이가 발생하지 않도록 유리함수를 사용한다. 동시에 k값을 활용해 한 시스템에서의 결과가 전체 결과를 지배하는 것을 막는다.

최종적으로 각 시스템을 동등하게 단순 합으로 계산하여 최종 순위를 매긴다.

CMNZscore

Cordorcet Fuse는 두 문서 쌍 $(d_1, d_2)$에 대해서 $r(d_1) < r(d_2)$라는 쌍대 관계(pairwise relation)를 기준으로 문서를 정렬한다.

CombMNZ score: 각 순위 r에 대해 대응하는 점수 $s_r: D \rightarrow \mathbb R$와 cutoff rank c에 따라 $CMNZscore(d \in D) = |{r \in R | r(d) \le c }| \cdot \displaystyle\sum_{{r|r(d)\le c}}s_r(d)$ 즉, (해당 문서를 좋게 본 시스템의 수) * (점수의 합) 으로 계산한다.

비교

훨씬 단순해보이는 RRF가 모든 경우에서 Condorcet Fuse보다 더 높은 점수를 보였고, 한 경우를 제외하곤 CombMNZ보다도 더 높았다. 또한 RRF는 인간이 직접 개입한 경우를 제외하고 개별 ranking들보다 더 좋은 성능을 보였다.

2.Discussion

단순히 RRF가 성능이 좋았다를 넘어 voting 알고리즘도 없고, 전역적인 정보를 저장할 필요도 없고, 시스템마다 순위들이 계산되어 누적되면 되기에 순위 결과를 저장할 필요도 없다.

RRF가 가장 성능이 좋았던 이유는 한 두개의 특정 시스템이 특정 문서를 높게 랭킹되는 것만으로도 순위가 크게 상승되기에 다양성이 보장되기 때문이라고 한다.

결론

최근에 사용하는 reranking들에 비해 성능은 떨어지지만 학습따윈 존재하지 않는 정말 간단한 휴리스틱 알고리즘이기에 경량화가 필요한 경우 활용할 수 있다.

1. Introduction

Pre-training Language Model이 등장하며 텍스트 임베딩(텍스트 데이터 -> latent space로 인코딩)의 품질이 향상 되었고, 이런 임베딩 중 특히 Dense retrieval로 Query와 Document간 유사도를 검색하는 IR에 필수적으로 사용된다. 이 외에도 Multi-vector retrieval이나 Sparse retreival에도 사용된다. 하지만 이런 임베딩 모델들은 영어에만 특화되어있으며, 하나의 검색 기능만을 목표로 학습된다. 또한 학습 비용이 크기 때문에 짧은 입력만을 지원한다. M3-Embedding은 다양한 언어를 지원하고, 서로 다른 언어 간 교차 언어 검색도 가능하다. 동시에 Dense + Sparse + multi-vector 까지 지원하는 범용 임베딩이 가능하다. 또한 문장의 길이에도 비교적 자유롭다. 이를 위해 사용한 기법으로

1) self-knowledge distillation

2. Related Work

Pre-training Language Model의 대중화로 데이터를 효과적으로 인코딩 하는 것에 더 유리해졌고, Contrastive Learning의 Negative Sampling의 개선과 Knowledge Distillation의 활용이 두드러진 성과를 보였다. 이를 주로 응용하는 분야는 Neural Retrieval로, 임베딩을 통해 의미적 관계를 측정해서 임베딩 유사도를 기반으로 문서를 검색할 수 있다.

1) Dense reetrieval: text encoder의 출력값을 [CLS] token 또는 average pooling과 같은 방식으로 집계해서 임베딩 유사도를 계산한다. 2) Multi-vector retrieval: text encoder의 출력 임베딩간의 상호작용으로 유사도를 계산한다. 3) Sparse retrieval: 각 단어의 가중치로 변화하여 단어를 기준으로 유사도를 계산한다. 이 모든 기능을 하나로 통합한 모델은 존재하지 않고, 대부분 영어에만 초점이 맞추어져있다.

3. M3-Embedding

다양한 언어를 지원하고, 서로 다른 세분성의 데이터를 처리하며, 일반적 검색 기능을 통합하기 위해서 만들어졌디. $D^y: d^y \leftarrow fn^{}(q^x,D^y)$, 이때 $fn^(\cdot)$의 경우에는 위의 dense, sprase, multi-vector 중 어느 하나이다.

Data Curation

라벨이 없는 corpus에서 추출한 비지도 데이터, 라벨이 부여된 corpus에서 수집한 fine-tuning 데이터, 합성을 통해 생성된 fine-tuning data 3개를 사용한다.

1) 비지도 데이터: 다언어 corpus에 내재된 의미 구조를 추출함으로써 구축한다.. MTP에서 제공되는 정제된 데이터와 교차언어를 위한 NLLB의 병렬 문장도 데이터로 활용했다. 2) 고품질 fine-tuning 데이터: 영어는 8개, 중국어는 7개의 데이터셋을 통합했고, 그 외의 기타 언어는 MIRACL의 학습 데이터를 활용했다고 한다. 3) 합성 fine-tuning 데이터: 긴 문서를 샘플링한 후 임의의 단락을 선택하고, GPT-3.5를 사용하여 해당 단락을 기반으로 질문을 생성했다고 한다.

Hybrid Retrieval

Dense, Sparse, multi-vector retrieval을 한개로 통합한다.

Dense Retrieval

Query q는 text encoder을 통해서 hidden state $H_q$로 변환된다. 특수토큰 [CLS]의 hidden state를 Query 표현으로 사용한다. $e_q = norm(H_q[0])$. 여기서 norm은 보통 L2 정규화라고 보면 된다. Document passage p에 대해서도 동일하게 한다. $e_p = norm(H_p[0])$ 관련도 점수는 내적으로 계산한다. $s_{dense} \leftarrow <e_p, e_q>$

Lexical Retrieval(Sparse)

Query 내 각 토큰 t에 대해 토큰의 가중치는 $w_{qt} \leftarrow ReLU(W^{\top}{lex}H_q[i])$ 이때, $W{qt}\in R^{d \times 1}$은 d차원의 hidden state를 한개의 실수값으로 매핑해주는 행렬이다. 만약 t가 여러 번 등장하는 경우에는 최대 가중치만을 유지한다. 이를 바탕으로 Query와 Passage에 공통으로 등장하는 단어 집합 $q \cap p$의 중요도는 $s_{lex} \leftarrow \displaystyle\sum_{t \in q \cap p}(w_{qt} \cdot w_{pt})$

Multi-Vector Retrieval

Dense Retrieval의 확장으로 Query, Passage를 표현하기 위해 전체 출력 임베딩을 활용한다. Dense와 달리 벡터 하나로 요약하는 것이 아닌, 토큰마다 벡터를 가진채로 Query 토큰과 가장 잘 맞는 문서 토큰을 대응해 점수를 매긴다. $E_q = norm(W^{\top}{mul}H_q)$, $E_p = norm(W^{\top}{mul}H_p)$ $W_{mul} \in R^{d \times d}$는 학습 가능한 projection matrix이다. 여기서 이전에 배운 ColBERT의 Late Interaction을 적용한다.

$s_{mul} \leftarrow {{1}\over{N}}\displaystyle\sum^N_{i=1} \displaystyle \max^M_{j=1} E_q[i] \cdot E^{\top}_p[j]$, (N, M은 Query와 Passage의 길이)

통합

각 방법을 사용하여 개별적으로 후보 결과를 검색하고, 해당 후보들을 합집합한 후 최종 검색 결과는 통합 관련도 점수를 기반으로 reranking한다. 이 과정에 multi-vector은 너무 무겁기에 후보 결과 검색땐 제외하고, reranking시에만 사용한다. $s_{rank} \leftarrow w_1 \cdot s_{dense} + w_2 \cdot s_{lex} + w_3 \cdot s_{mul}$ 이때 $w_1, w_2, w_3$는 hyperparameter인데 task에 따라서 다르게 배분했다고 한다.

Self-Knowledge Distillation

Knowledge Distillation(지식 증류)

우선 Knowledge Distillation부터 알아보자. 간단하게 잘 아는 모델(Teacher)가 잘 모르는 모델(Student)에게 단순한 답이 아닌 발상까지 넘겨주는 방법이다.

우선 기존의 학습 방식은 정답이 A라고 주어지면 모델은 A일 확률이 0.80이라고 예측하는 방식이었다. 지식 증류에선 Teacher가 어떤 데이터에 대해 A(0.80), B(0.15), C(0.05)라고 예측했다면, 이 분포를 Student에게 그대로 넘겨준다. 즉, 단순한 답 뿐만 아니라 Decision Boundary를 전달해주는 방식이다. 해당 방식을 통해 일반화(Generalization) 성능이 증가하고, 작은 모델을 더 효율적으로 학습시킬 수 있다.

이를 수식적으로 나타내보면 교사 분포 P, 학생 분포 Q에 대해 loss는 $\mathcal L_{KD} = CrossEntropy(P, Q) = - \displaystyle \sum_i P_i log Q_i$

CrossEntropy: 두 분포가 가장 비슷해지도록(확률 분포가 가까워 지도록) 설정하기 위해 KL Divergence의 최소화를 정리한 식을 Cross-Entropy라고 한다.

Embedding Loss

임베딩 모델에선 Positive sample과 Negative sample을 구분하도록 학습된다. 각 검색에 대해서 Positive가 더 높은 점수를 받도록 부여되도록 학습한다. 수식으로 보면 $\mathcal L_s(\cdot) = -\log{{\exp(s(q,p^{}))/\tau}\over{\sum_{p\in{p^{},P'}}exp(s(q,p))/\tau}}$ 이때 p*와 P'는 Query q에 대한 positive sample과 negative sample이고, $s(\cdot)$은 {$s_{dense}(\cdot), s_{lex}(\cdot), s_{mul}(\cdot)$} 중 하나이다.

Self-Knowledge

지금 해당 모델이 전체적으로 dense, lexical, multi-vector 3가지를 섞어쓰는만큼 학습 목표가 충돌할 가능성도 존재한다. 이를 해결하기 위해서 Self-Knowledge distillation을 기반으로 학습 과정을 통합한다. Priciple of ensemble learning에 따라서 여러 검색 방식의 예측 결과를 결합하면 더 정확한 관련성 점수를 얻을 수 있다.

Principle of ensemble learning: 작은 모델 여러 개를 묶으면 개별 모델보다 안정적이다. Bagging: 같은 알고리즘, 랜덤 샘플링으로 과적합 방지 Boosting: 이전 모델이 틀린 문제에 가중치를 둬서 다음 모델이 집중 학습 Voting: 서로 다른 알고리즘을 섞어서 투표

Distillation을 적용시키지 않은 모델(단순 weight sum)을 보면, $s_{inter} \leftarrow w_1 \cdot s_{dense} + w_2 \cdot s_{lex} + w_3 \cdot s_{mul}$ $\mathcal L \leftarrow (\lambda_1 \cdot \mathcal L_{dense} + \lambda_2 \cdot \mathcal L_{lex} + \lambda_3 \cdot \mathcal L_{mul} + \mathcal L_{inter}) / 4$

여기서 다른 ranking모델에게 얻은 soft label을 활용해서 distillation을 할 수 있다. 해당 모델에선 $s_{inter}$을 teacher로 사용하며, 이때 loss는 (p($\cdot$)은 softmax) $\mathcal L'{*} \leftarrow -p(s{inter}) * \log p(s_)$ 식을 살펴보면, Query가 주어졌을 때 후보 문서들 중 어느 문서가 정답인지를 맞추는 문제이기에 이를 softmax()==p($\cdot$)으로 바꾼다. 이를 위에서 본 식에 대입하면 $P_i = p(s_{inter, i})$, $Q_i = p(s_{, i})$ $\mathcal L'* = -\displaystyle\sum_i p(s{inter, i})\log p(s_{*, i})$ 해당 논문에선 한 Query에 한 번만 사용하니 위의 식이 나오게 된다.

이제 위의 식을 바탕으로 loss function을 재정의하면 $\mathcal L' \leftarrow (\lambda_1 \cdot \mathcal L'{dense} + \lambda_2 \cdot \mathcal L'{lex} + \lambda_3 \cdot \mathcal L'_{mul}) / 3$

최종적으로 Self-Knowledge Distillation의 최종 loss function은 L, L'의 선형결합인 $\mathcal L_{final} \leftarrow (\mathcal L + \mathcal L')/2$ 로 한다.

Learning

1. text encoder은 비지도 데이터로 pre-training되며 dense retrieval만 학습한다.
2. 다른 지도학습 데이터까지 전부 활용하여 Self-Knowledge Distillatio을 통해 임베딩 모델이 세 가지 검색 기능을 모두 학습하도록 fine-tuning된다.

• 이때 $W_{lex}$는 무작위 초기화로 인해 학습 초기는 $s_{lex}$ 정확도가 낮고, $L_{lex}$가 크다. 이를 방지하기 위해 학습에선 $w_1 = 1, w_2 = 0.3, w_3 = 1$로, $\lambda_1 = 1, \lambda_2 = 0.1, \lambda_3 = 1$로 설정한다.
• 또한 각 Query에 대해 ANCE에 따라 Hard-Negative Sample을 도입한다.
• ANCE는 BM25와 달리 현재 모델을 기준으로 정답은 아니지만 점수가 높은 문서를 negative를 사용한다.

Efficient Batching

다양한 언어에 대한 일반적인 의미를 얻기 위해서 대규모의 다양한 다국어 데이터를 학습해야 한다. 추가로 텍스트 임베딩의 판별력을 보장하기 위해서 가능한 큰 batch size를 유지해야 한다. 하지만 GPU 메모리의 제약이 있기에 batch size를 유지하기 위해선 입력 데이터를 짧은 시퀀스로 잘라 학습해야 한다. M3-Embedding의 경우는 서로 다른 입력 단위의 데이터를 효과적으로 처리하기 위해서 짧은 시퀀스에서 긴 시퀀스까지 전부 다 학습해야하기 때문에 batching 전략을 최적화해야 한다.

우선 입력 데이터를 시퀀스 길이에 따라서 그룹화한다. mini-batch를 적용할 경우 동일한 길이의 그룹에서만 추출한다. 이로인해 padding으로 인한 낭비가 줄어들며, GPU 자원을 효율적으로 활용할 수 있다. 또한 서로 다른 GPU에서 데이터를 샘플링할 때 동일한 시드를 사용하여 부하 균형을 유지한다. 즉, 비슷한 난이도의 샘플을 뽑게 만드는 과정을 통해 가장 느린 GPU 한개로 전부가 멈추는 현상을 방지한다.

또한 긴 시퀀스의 경우 mini-batch를 더 작은 sub-batch로 분할한다. gradient checkpoint를 활용해 sub-batch를 순차적으로 인코딩하고, 모든 임베딩을 수집한다. 이후 서로 다른 GPU에서 생성된 임베딩을 broadcasting한다. 이를 통해 in-batch negative sample의 규모가 확장된다.

gradient checkpoint: 중간 activation을 저장하지 않고 backward때 다시 계산하는 방식. 연산을 조금 포기하고 메모리를 줄임.

추가로 논문에선 현실적인 방안으로 MCLS도 제안하고 있다.

MCLS: 긴 문서를 일정 간격으로 쪼개 [CLS]를 추가로 삽입하여 [CLS] embedding들을 평균내는 방식

4. Experiment

1) Multi-Lingual Retrieval: 18개 언어에서 각 언어 내에서 QA를 검색했을 때 Dense, Sparse, Multi-vector 각각의 성능도 뛰어나고 세개를 합쳤을 때 최고 성능이 나왔다. 2) Cross-Lingual Retrieval: Query는 비영어, 문서는 영어일 때 언어 전반에서 일관성을 보이고 특히 저자원 언어에서 M3가 훨씬 안정적이었다. 3) Multilingual Long-Document Retrieval: Sparse가 훨씬 유리했고 Dense+ Sparse, All조합이 최고성능이었다. 4) Ablation Study a) Distillation 제거시 Sparse가 폭락 b) 대규모 비지도 사전학습이 다국어 안정성의 기반

5. Conclusion

이전의 SPLADE도 Sparse 방식과 Dense방식을 합치려고 했다. 하지만 M3의 경우 단순히 dense를 보조로 사용한 것이 아닌 각각 다른 학습 방식을 같이 학습시킨다. 뭔가 난이도가 엄청 높은 느낌은 아니였지만 내용이 좀 많은 느낌이었다. 천천히 정리해보고 이젠 조금 더 실용적인? 쪽의 논문과 함께 해보고싶던걸 해볼 생각이다.

1. Introduction

정보 검색은 수작업 -> 임베딩 변환 후 내용 간 관계 모델링 -> LM(BERT)을 순위화 작업에 파인튜닝으로 발전해왔지만, 계산 비용이 100배에서 1000배까지 증가한다는 단점이 존재했다. 이를 해결하기 위해 BM25와 같은 전통적인 모델을 섞어서 사용하기도 하지만, 이 경우 검색 정확도가 크게 감소한다. 이를 해결하기 위해 BERT 기반 문맥화된 late interaction에 기반한 ranking model이 ColBERT이다.

late interaction: Query와 Document를 각각 두 개의 문맥화된 embedding 집합으로 독립적으로 인코딩하고, 이 두 집합 사이의 계산을 통해 관련성을 평가한다.

상호작용 중심 모델은 IR에서 성능이 좋고, 표현 중심 모델은 Query와 Document의 게산을 분리하면서 Document쪽을 오프라인에서 계산할 수 있게 해주어 Query당 계산 비용이 크게 줄어든다. ColBERT에선 두 장점을 합치는 것이 목적이다.

상호작용 중심 모델: Query의 token과 Document의 token을 하나하나 직접 비교하고 결과를 모아서 점수를 계산 표현 중심 모델: Query와 Document를 각각 하나의 벡터로 요약해서 유사도를 게산

2. Related Work

Neural Matching Models

KNRM: 상호작용 행렬로부터 matching signal을 추출하기 위해서 미분 가능한 kernel-polling techinique를 이용한다.

kernel-polling: query-document간 유사도 행렬을 kernel이라는 구간으로 나눠서 사용. 즉, 비슷한 단어들도 점수에 반영

Duet: exact-match-based와 embedding-based 유사도를 결합해서 순위를 매긴다.

Exact-match: BM25와 같이 질의 단어가 그대로 등장하는가 Embedding: 단어는 다르지만 의미가 비슷한가

ConvKNRM: Query-document간 n-gram을 매칭한다

n-gram: 단어 하나가 아닌 단어 덩어리를 맞춰본다. 이때 kernel-pooling으로 처리한다.

fastTest+ConvKNRM: sub-word token embedding을 활용해 rare word문제를 해결한다.

단어를 n-gram으로 쪼개서 단어 벡터를 구성한다.

SNRM: Query와 Document를 한 latent term로 인코딩하고 inverted index를 활용해 문서를 표현할 수 있다. 이를 통해 end-to-end 검색이 가능해진다. 하지만 최신 기법들에 비해서는 성능이 뒤쳐진다.

Language Model Pretraining

BERT와 같이 비지도 학습으로 pre-training을 진행하고, 이를 downstream 작업에 fine-tuning한다. 일반적으로 Query-Document 쌍을 BERT에 입력하고, [CLS]토큰에 MLP를 얹어서 점수를 계산하는 방식이다.

BERT Optimizations

LM기반 Ranking model들은 계산 비용이 매우 높다. 이를 줄이기 위해 distilling, compressing, pruning 등 연구가 있지만, 대부분 범용적인 BERT를 위한 것이기에 IR에서는 제한적인 향상만 보인다.

Efficient NLU-based Models

doc2query처럼 비용이 큰 NLU연산을 오프라인에서 수행하는 방향이다.

3. ColBERT

Query-Document간 상호작용을 지연시키는 것으로 비용을 감소시키고 , 벡터 유사도 검색을 통한 end-to-end 검색을 지원할 수 있다.

Architecture

구성 요소

1) Query Encoder $f_Q$ 2) Document Encoder $f_D$ 3) Late interaction

Query q와 Document d가 주어지면 $f_Q$는 q를 fixed dimension 임베딩 집합인 $E_q$로 인코딩하고, $f_D$는 d를 또 다른 임베딩 집합인 $E_d$로 인코딩한다. 이때 두 임베딩은 다른 토큰들에 의해 contextualized 된다. 이후 $E_q$와 $E_d$를 바탕으로 late interaction을 통해 Query와 Document간 score를 계산한다. 여기서 late interaction은 MaxSim의 합이다.

Maxsim: 가장 큰 값 한개만 남기는 방법 여기선 $E_q$에 속한 각 벡터 v에 대해 $E_d$의 벡터들과 cosine 유사도 중 최대값을 구하고, 이 값들을 모두 합산하는 방식이다. (suqared L2를 이용하기도 한다)

이런 토큰별 매칭 점수를 합산해서 문제의 전체 관련성을 추정한다. MaxSim이기에 계산 비용이 매우 낮고, top-k개의 검색을 위한 pruning이 가능하다.

Query & Document Encoder

Query와 Document 모두 하나의 BERT 모델을 공유하고, 특수토큰 [Q]와 [D]를 통해서 구분한다.

Query Encoder

$E_q:= Normalize(CNN(BERT("[Q]q_0q_1...q_l##...#")))$

Query q가 주어지면 BERT의 WordPiece tokenizer를 사용해 $q_1, q_2,...,q_l$ 형태의 토큰 시퀀스로 변환한다. 이후에 앞에는 [Q]토큰을 추가한다.([CLS] 바로 뒤) 사전에 정의된 길이에 비해 Query token이 작은 경우 [MASK]토큰(#)을 사용해 길이가 $N_q$가 되도록 padding하고, 길이가 초과할 경우 처음 $N_q$개만 남긴다. 이를 query augmentation이라고 부른다. 이를 통해 Query-based emgedding을 생성하도록 유도하며, 새로운 용어를 추가하거나 기존 용어의 중요도를 재조정하는 메커니즘으로 작동한다. BERT가 생성한 각 토큰의 문맥화된 표현은 activation function이 없는 Linear layer을 통과하면서 출력 차원 m의 임베딩을 생성한다. 이는 Query encoder의 효율에는 관련이 없지만, document embedding의 저장 공간에 영향을 주고, GPU로 전송하는 비용에도 영향을 주어 Query를 처리하는 시간에도 영향을 끼친다. 마지막으로 L2 norm이 1이 되도록 정규화된다.

Document Encoder

$E_d := Filter(Normalize(Linear(BERT("[D]d_0d_1...d_n"))))$

$d_1, d_2,...,d_n$ 형태의 토큰 시퀀스로 분할되며 맨 앞에는 [CLS], 그 뒤에는 [D]를 추가한다. 단, [MASK]는 추가하지 않는다. 해당 입력은 BERT, Linear layer을 통과하고, 구두점(제거해야하는 문장부호 등)에 해당하는 토큰 임베딩을 제거한다.

Late Interaction

$S_{q,d} := \displaystyle\sum_{i\in||E_q||} \displaystyle \max_{j \in ||E_d||} E_{q_i} \cdot E_{d_j}$

Query embedding $E_q$에 포함된 각 벡터에 대해 Document embedding $E_d$에 포함된 모든 벡터들과 유사도 중 최대값을 구하고, 이를 Query token전체에 대해 합산한다. 이때 보통 유사도 척도로는 cosine 또는 quarted L2 distance를 사용한다. ColBERT는 미분 가능하고 end-to-end로 학습된다. 또한 BERT는 추가로 도입된 선형 계층이나 특수 토큰 [Q], [D]의 임베딩도 학습한다. 이때 late interaction자체는 학습 가능한 파라미터가 존재하지 않는다. 만약 Query q, positive document $d^+$, negative document $d^-$로 구성된 <q, $d^+$, $d^-$>가 주어지면 각 문서에 대해 독립적으로 점수를 계산하고 $d^+$와 $d^-$에 대해 pairwise softmax cross-entropy loss를 사용해 학습한다.

Offline Indexing

Query와 Document 계산을 분리시키며 Document를 오프라인에서 계산할 수 있다. 문서 컬렉션을 batch 단위로 순회하면서 각 batch에 대해 document encoder $f_D$를 실행하고, 이런 임베딩을 저장한다. 이를 최적화하기 위해 GPU를 활용해 문서 batch를 병렬로 인코딩하고, batch 처리시에는 가장 긴 문서 길이의 문서에 맞춰서 모든 문서를 패딩한다. WordPiece tokenization에서 많은 시간이 소비되지만 이런 전처리는 문서들 간에 독립적이기에 이를 CPU 코어에서 병렬로 수행한다.

Top-k Re-ranking

비교적 작은 수 k(1000)에선 각 후보 문서를 전부 다 평가한다. Query serving sub-system은 사전에 계산되어 저장된 document embedding을 메모리로 불러오고, 각 문서는 임베딩 형태로 표현된다 Query q에 대해 Query embedding $E_q$를 계산하고, 동시에 k개의 문서에 대한 임베딩을 모아 3차원 텐서D를 구성한다. 이후 D는 GPU로 이동한다. GPU에서는 $E_q$와 D 사이의 batch dot-product 연산을 수행하고(이때 min-batch로 나누어질수도 있다) 결과적으로 query와 각 document간 교차 매칭 행렬로 구성된 3차원 텐서가 생성된다. 이후 document token 차원에 대해 합산을 하고, k개의 문서를 이 점수에 따라서 정렬한다. $l = |q|$인 단일 시퀀스만 BERT에 입력하기에 계산이 매우 저비용이고, k가 증가할수록 더 잘 확장이된다.

End-to-End Top-k

후보가 너무 많아 전부 다 평가하기엔 문제가 존재하는 경우, MaxSim의 Pruning에 유리하다는 점을 이용한다. Query 임베딩 하나와 Document 모두에 대해 Maxsim을 계산하는 대신, fast vector-similarity data structure를 이용한다. 여기선 FAISS를 사용한다.

Document

Offline indexing 이후에 각 임베딩과 document간 매핑을 유지하고 FAISS index에 삽입한다. 이때 문서별이 아닌 전체 컬렉션의 모든 토큰 임베딩을 한 바구니에 모은다.(V = {$v_1,v_2...,v_M$}

Query

ColBERT 점수 계산 방식을 기반으로 한 두 단계로 처리한다. 각 단계를 조금 더 나눠서 보면

1) 근사치 뽑아내기

• a) Query q를 $f_Q$로 처리해서 $E_q$ = {$q_1, q_2, ...,q_{N_q}$}
• b) 전체 V에서 $q_i$와 가장 유사한 임베딩 상위 k'(=k/2)개를 가져온다. 즉, Query token = 32개일때 $N_q \times k'$개의 document token을 가져온다.
• c) 해당 token들을 document로 변환한다. 이때 document 개수 $K \le N_q \times k'$

2) 정밀

• 후보 문서 K개에 대해서 ColBERT대로 $S_{q, d} = \displaystyle\sum_i \displaystyle\max_j sim(q_i, d_j)$를 계산해서 최종 top-k개를 만든다.

4. Experiment

BERT급 성능을 내면서 동시에 BERT보다 170배 이상 빠르고, FLOPs는 14000배 감소한 결과가 나온다. 또한 End-to-End Retrieval에서도 기존의 BM25, doc2query보다 성능이 우수하며 동시에 BM25없이도 직접 검색이 가능했다. Ablation Study결과 하나라도 빠질경우 성능이 떨어지는 결과가 나왔고, Average similarity보단 MaxSim의 성능이 더 좋았다.

5. Conclusion

기존의 BERT를 채용하되, Query와 Document를 Late Interaction을 통해 분리시켰다. 이를 통해 Offline Indexing이 가능해졌다. 동시에 GPU 연산에 더 유리한 연산들만 남기면서 속도도 더 높인 구조이다. 이때 Dense과 달리 단어 토큰 자체는 그대로 유지한다. 때문에 용량을 많이 차지한다는 단점이 존재한다.

1. Introduction

기존의 BERT등 Large Pre-training Language Model + Fine tuning만으로 여러 과제에서 적용이 가능했지만, 주로 Re-Ranker에서만 사용되고 1단계 Retrieval은 inverted index에 의존하는 BoW모델이 주로 사용된다. 하지만 다른 논문들에서도 언급했듯 용어가 관련 문서에 없는 경우 어휘 불일치 문제가 발생한다. 이전에 살펴본 DPR처럼 Dense retriever은 용어의 명시적 매칭을 모델링할 수 없다. 이를 해결하기 위해 sparse representation을 학습하려는 연구가 많아지고 있다. 해당 모델은 Logarithmic Activation과 Sparse Normalization을 기반으로 하여 복잡한 학습 파이프라인을 이용하는 dense 모델들과 비교해도 충분히 경쟁력 있는 성능을 보인다고 한다.

2. Related Works

BERT Siamese model을 기반으로 하는 Dense retreival이 QA와 정보 검색의 후보 생성의 표준이고, 여기서 더 나아간 ColBERT로 1단계 검색에서도 효율적으로 적용이 가능하다

ColBERT: 문서를 토큰 단위로 쪼개고, 질의 토큰 각각을 문서 토큰 전부랑 유사도를 계산해서 효율성을 더 높인 방식 정교하지만 저장공간이 많이 필요하고, 대규모로 갈수록 문제가 발생함.

Dense index대신 Term-based index방식(SNRM)도 존재하지만 효율성이 확실하지 않다.

SNRM: inverted-index를 유지하면서 가중치 게산만 신경망에 맡기는 방식 L1 Regularization을 사용하기에 희소성을 만들어내지만, 효율성이 떨어짐

생성 기반 접근(doc2query, docTTTTTquery)을 통한 document expansion방식도 존재하지만 학습 목표가 간접적이고, 통제가 불가능하다.

doc2query: 문서 자신을 설명하는 query를 생성모델(T5 등)을 활용해 미리 만들어서 넣어둠

Document, Query의 각 토큰이 전체 어휘의 각 토큰에 대해 갖는 중요도를 추정(SparTerm, EPIC, SPARTA)하는 방식도 존재한다. 하지만 EPIC, SPARTA와 같이 최대값을 이용하는 경우에는 희소성 확보를 위한 추가적인 방식이 필요하고, SparTerm과 같은 합을 이용한 방식은 명시적인 희소성 정규화가 없기에 성능이 제한된다.(자세한건 다음 챕터에서 다룬다)

3. SPARSE Lexical Representation

SparTerm

BERT WordPiece상에서의 단어 중요도를 Masked Language Model(MLM)의 logit을 기반으로 예측한다. 조금 더 자세히 보면

• $t = (t_1, t_2, ..., t_n)$: WordPiece를 Tokenization한 이후 문서 시퀀스 or Query
• $(h_1, h_2, ..., h_n)$: $t$에 대응하는 BERT embedding
• $i$: 입력 시퀀스 토큰
• $j$: $i$에 대한 어휘 토큰 -> 나중에 검색 후보가 될 문서들
• $E_j$: $j$에 대한 BERT embedding
• $b_j$: 토큰 수준의 bias. 문맥과 관련없이 잘 나오는 단어들(the 등)에 대한 baseline을 높여줌
• $transform(.)$: GELU Activation + LayerNorm

중요도 $w_{ij} = trasnform(h_i)^{\top}E_j+b_j;;;;; j\in{1,...,|V|}$

이후 최종 표현은 가중치가 양수가 되도록 ReLU를 적용하고, 중요도 예측값을 합산하여 $w_j = g_j + \displaystyle\sum_{i \in t}ReLU(w_{ij})$ ($g_j$는 이후에 설명) 이는 각 어휘에 대해서 새로운 가중치 $w_j$를 예측하기에 query/document expansion으로 볼 수 있다.

여기서 SparTerm은 두가지 희소화 방식을 추가한다

1) Lexical-only: 토큰 j가 입력 시퀀스 t에 등장하면 $g_j$ = 1, 아니면 0 2) Expansion-aware: 확장을 고려한 binary gating. $g_j$는 Lexical-only로 1인 경우를 제외하고 학습됨.

이제 얻은 질의 표현 q와 문서 표현 d의 내적을 통해 얻은 점수가 $s(q,d)$이고, Query $q_i$, positive document $d_i^+$, Negative document $d_i^-$에 대해서 $L_{rank} = -log {{e^{s(q_i,d_i^+)}}\over{e^{s(q_i,d_i^+)}+e^{s(q_i,d_i^-)}}}$를 최소화하도록 학습한다.

한계점

• expansion_aware이 복잡하고, end-to-end로 학습할 수 없다.
• $L_{rank}$로 모델을 fine-tuning하는 동안 gating이 업데이트 되지 않기에 과제에 최적화된 전략을 학습하는 것이 불가능하다.
• Lexical과 expansion-aware방식이 거의 동일한 성능을 보인다는 의문점이 존재한다.

SPLADE

Model: Log-Saturation

$w_j = \displaystyle \sum_{i \in t} \log(1+ReLU(w_{ij}))$ 일부 단어가 과도하게 지배하는 것을 막고, 표현의 희소성을 자연스럽게 유도하는 방식이다. log 그래프를 보면 일정 이상부터는 매우 천천히 증가하는 형태가 된다. 이로 인해 특정 단어가 수백번 반복되더라도 값이 엄청나게 차이가 나지 않는다. 동시에 실험적으로 log-saturation을 이용하면 정규화 없이도 sparse solution을 구할 수 있다고 한다.

Ranking Loss: in-batch negatives(IBN)

학습 단계에서 하나의 Query $q_i$를 하나의 정답 문서 $d_i^+$에 대응시킨다. 이때 기존처럼 batch 내의 Query $q_i$, Positive document $d_i^+$, BM25등으로 얻은 하나의 negative document $d_i^-$에 추가로 batch 내에서 $d_i^+$ 외의 문서들을 ${d_{i, j}^-}j$라고 한다. 우리는 DPR 에서 동일한 방식을 봤다. DPR에서도 데이터셋에 명시적으로 제시되는 Positive와 BM25에서 정답을 포함하지는 않지만 Query와 매칭되는 의미적으로 헷갈리기 쉬운 passage인 Negative, 그리고 다른 Query와 짝지어진 Gold를 사용했었다. 이를 종합해서 $L{rank-IBN} = -\log{{e^{s(q_i,d_i^+)}\over{e^{s(q_i,d_i^+)}+e^{s(q_i,d_i^-)}+\sum_j e^{s(q_i,d_{i,j}^-)}}}}$

Learning Sparse Representation: FLOPS 정규화

위의 과정을 거쳐도 특정 용어가 너무 많은 문서와 연결될 수 있다. 이때 해당 후보 문서를 전부 훑으면서 연산량이 터지는 문제가 발생한다(Zipf's law: $빈도 \propto {{1}\over{순위}}$). 이를 해결하기 위해서 L1정규화만으로는 부족하기에, 새로운 정규화 과정을 진행한다.

term j에 대한 batch 내 문서들의 평균 가중치

• N: batch 문서 수
• $w_j^{(d_i)}$: i번째 문서의 term j의 가중치 가중치 $a_j = {{1}\over{N}} \displaystyle \sum ^{N}_{i=1} w_j^{(d_i)}$

  FLOPS 정규화

  $l_{FLOPS} = \displaystyle \sum_{j \in V}a^2_j = \displaystyle \sum_{j \in V}({{1}\over{N}}\displaystyle \sum^N_{i=1}w_j^{(d_i)})^2$
• 특정 term j가 많은 문서에서 동시에 켜질 경우 평균 $a_j$가 제곱으로 커지면서 L1 정규화보다 더 큰 loss가 발생한다.

이런 방식을 통해서 희소성은 더 높아지고, 연산량은 더 줄어든다.

Overall loss

$L = L_{rank_IBN} + \lambda_qL_q^{reg} + \lambda_dL_d^{reg}$ 즉, 더 희소 + 확장도 가능한 표현을 end-to-end로 결합시킬 수 있게된다.

4. Experiment

기존의 sparse 모델을 압도하고 Dense retriever과 거의 동급의 성능을 보인다. 동시에 FLOPS는 BM25수준까지 낮춘다. 이를 통해 효율-성능 간 tradeoff가 L1보다 우수하다. 정규화 강도를 조절하는 것으로 성능/속도 중심을 선택할 수 있고, 극단적으로 sparse하더라도 여전히 경쟁력이 존재한다.

5. Conclusion

결론적으로 BM25의 단점인 확장 불가를 dense의 기법으로 해결하면서, dense의 단점인 명시적 매칭 불가를 희소성을 이용한 기법으로 해결한다. 성능도 dense의 성능이 나오면서 효율성은 BM25급의 효율이 나온다. 한계점으론 Expansion이 존재하기에 BM25에 비해서 index가 커질 수 밖에 없고, 하이퍼파라미터의 튜닝 정도에 따라서 크게 흔들린다. 또한 여전히 WordPiece Tokenizer에 의존한다.

잡소리: 원래 BM-25한 김에 SPLADE를 할려고 했는데, 아무리봐도 이게 먼전거같아서 RAG를 우선 읽어본다. 굳이 순서를 따지자면 연도순대로 BM25 - DPR - RAG - SPLADE나 아니면 DPR - RAG - BM25 - SPLADE로 읽도록 하자

1. Introduction

Pre-training Neural Language Model 자체도 외부 메모리에 대한 접근 없이 지식 베이스로서 동작이 가능하다. 하지만 메모리에 대한 확장이나 수정이 어렵고, 직관적인 설명도 제공할 수 없고, 환각(hallucination)을 생성하기도 한다.

이를 해결하기 위해 parametric memory와 non-parametric memory (retreival-based)를 결합한 하이브리드 모델이 나온다.(REALM, ORQA) 하지만 부분 Open domain에서 추출해오는 QA형 과제에만 적용되기에, 해당 모델에서는 Seq2Seq 모델로 확장시킨다. 즉, parametric memory로 pre-training된 Seq2Seq transformer를 사용하고, non-parametric memory로 위키피디아 Dense vector index를 활용하여 검색하는 DPR를 활용해 두 요소를 하나의 확률 모델로 결합하여 end-to-end로 학습한다.

기존이 연구들과 다른 점은 parametric memory와 non-parametric memory 모두 pre-training되어있기에 추가적인 학습 없이도 지식에 접근할 수 있는 능력이 존재한다.

2. Method

입력 시퀀스 $x$, 사용할 텍스트 문서 $z$, 목표 시퀀스(생성) $y$ 이때 DPR논문에서처럼

1) Retriever: Question $x$에 대해 text passage에 대한 분포를 변환하는 매개변수 $\eta$를 갖는 $p_{\eta}(z|x)$ 2) Generator: 이전 i-1개 토큰 $y_{1:i-1}$, 입력 $x$, 검색된 passage $z$에 따라 현재 토큰을 생성하는 매개변수 $\theta$를 가지는 $p_{\theta}(y_i | x, z, y_{1:i-1})$

end-to-end를 위해 Retriever을 통과한 문서를 잠재 변수(latent variable) 취급하고, 주변화하여 생성된 텍스트에 대한 분포를 generate하는 두 가지 모델이 존재한다.

Model

RAG-Sequence

동일한 문서를 사용하여 전체 목표 시퀀스 예측하는 것이 목표이다. 검색된 문서를 하나의 잠재 변수로 취급하고, 이를 marginalize하여 Seq2Seq 확률 $p(y|x)$를 구한다. $p_{RAG_Sequence}(y|x) \approx \displaystyle\sum_{z \in top-k(p(\cdot|x))}{p_\eta (z|x)p_\theta(y|x, z)} = \displaystyle\sum_{z \in top-k(p(\cdot| x))}p_\eta (z|x) \prod_{i}^{N} p_\theta(y_i|x,z,y_{1:i-1})$

식이 길고 지루하고 현학적이다. 장난이고 조금 더 자세히 설명해보면 $p_\eta(z|x)$는 질문 x에 대해 문서 z가 관련 있을 확률, $p_\theta(y|x,z)$는 x, z를 참고해 y라는 답변이 나올 확률이다. 즉, (해당 문서를 썼을 확률) $\times$ (이 문서를 통해 그런 답이 나올 가능성)으로 해석하면 된다. 가장 오른쪽 식은 $p_\theta(y|x, z) = \displaystyle\prod_{i}^{N}p_\theta(y_i|x,z,y_{1:i-1})$ 즉, 전체 문장이 나올 확률은 각 토큰이 나올 확률의 곱과 같다는 점만 생각해보면 된다.

RAG-Token

각 목표 토큰마다 서로 다른 잠재 문서를 선택해서 이에 대해 marginalize한다. 즉, retriever이 가져온 top-K 문서에서 선택적으로 활용하여 각 문서에 대해 다음 출력 토큰의 분포를 계산하고 marginalize한다. $P_{RAG_Token}(y|x) \approx \displaystyle\prod_{1}^N \displaystyle\sum_{z\in top-k(p(p\cdot|x))} p_\eta(z|x)p_\theta(y_i|x,z,,y_{1:i-1})$

DPR

자세한건 이전 글을 참고하자 $p_\eta(z|x) \propto exp(d(z)^{\top}q(x))$, 이때 $d(z) = BERT_d(z)$, $q(x)=BERT_q(x)$

즉, d(z)는 문서의 dense, q(x)는 Query 표현 이다. $top-k(p_\eta(\cdot |x))$를 계산하는 것은 MIPS문제로, 이는 sub-linear time동안 근사적으로 계산이 가능하다. (대충 그럴듯한 top-K를 찾는다고 보면 된다)

해당 문서 인덱스를 non-parametric memory라고 부른다.

Generator: BART

Generator은 다른 encoder-decoder모델로도 가능하지만, 해당 논문에서는 pre-training된 seq2seq인 BART-large를 사용한다.

출력 생성시에 입력 x와 Retriever를 통과한 문서 z를 단순하게 concat한다. BART generator의 매개변수 $\theta$는 parametric memory라고 부른다.

Training

Retriever와 Generator를 공동으로 학습하며, 입력-출력 쌍에 대해 negative marginal log-likelihood를 Adam을 사용하는 확률적 경사 하강법으로 최소화한다. 이때 학습 중에 $BERT_d$를 업데이트하는 것은 비용은 크지만 필수적이지 않기에 문서 encoder와 문서 인덱스를 고정하고 $BERT_q$와 BART Generator만 fine-tuning한다.

Decoding(토큰을 선택해서 문장 만들기)

테스트 시점에서 $argmax_y p(y|x)$를 근사하기 위한두 모델 간 차이가 약간 존재한다. 우선 Standard beam decoder을 위해선 다음 토큰의 확률이 이전 토큰들에만 의존하고, 토큰마다 하나의 명확한 확률 분포가 필요하다.

beam search: decoding 단계에서 확률이 제일 큰 토큰만 고르면 문장이 망할 수 있기에 여러 가능성을 고려하는 전략

RAG-Token

$P_{RAG_Token}(y|x) \approx \displaystyle\prod_{1}^N \displaystyle\sum_{z\in top-k(p(p\cdot|x))} p_\eta(z|x)p_\theta(y_i|x,z,,y_{1:i-1})$

해당 식에서 $p'\theta(y_i|x,y{1,i-1}) = \displaystyle\sum_{z \in top-k(p(\cdot|x))}p_\eta(z|x)p_\theta(y_i|x,z,y_{1:i-1})$로 사실상 표준 Seq2Seq generator로 볼 수 있기에 Standard beam decoder을 그대로 사용 가능하다.

RAG-Sequence

$p(y|x)$가 토큰 단위의 likelihood로 분해되지 않는다. 따라서 문서 z에 대해 개별적으로 beam search를 수행한다. 이때, 가설 y에 대해 $p_\theta(y_i | x,z,y_{1:i-1})$를 사용해 점수화된다. 이 과정에서 가설 집합 Y를 얻는데, 일부는 beam 결과에 나타나지 안흔ㄴ다.

가설 y의 확률을 추정하기 위해 해당 가설이 beam에 등장하지 않은 문서 z 각각에 대해 추가적인 forward pass를 수행하고 Generator 확률에 $p_\eta (z|x)$를 곱한 후 모든 문서에 대해 확률을 합산하여 margin 확률을 계산한다.

만약 출력 시퀀스가 길어서 |Y|가 너무 커진다면 beam search 과정에서 생성되지 않은 y에 대해 $p_\theta(y|x, z_i) \approx 0$이라고 근사해버린다.

3. Experiments & Result

DPR때처럼 위키피디아를 100단어 chunk기준 passage로 분할해서 2100만개 문서를 만들고, FAISS를 이용한 HNSW 기반의 MIPS를 구축하여 검색을 수행한다.

Open-domain Question

QA의 질문-답변을 입력-출력 쌍으로 취급하고 negative log-likelihood를 최소화하도록 RAG를 학습한다. 널리 사용되는 NQ, TQA, WQ, CT 데이터셋을 이용한다.

결과에서 크게 중요한 것은

1) 특수한 salient span masking pre-training없이도 강력한 성능을 보인다.

salient span masking: 아무 span이나 가리는 것이 아닌 핵심 정보만 가리는 방법

2) Re-ranker나 extractive reader 없이도 경쟁력있는 성능을 보이며 둘이 필수가 아님을 보여준다.

Abstractive Question Answering

Knowledge-intensive setting에서 자연어 생성 능력을 평가하기 위해 MSMARCO NLG v2.1을 이용했다. 이 중 특정 지역의 날씨와 같이 정답 문서에 접근하지 않으면 일치하는 답변을 할 수 없는 질문도 포함되어있다.

결과를 요약해보면 기존의 BART에 비해서 환각이 적고 사실적으로 올바른 텍스트를 생성하는 비율이 높다

Jeopardy Question

QA가 아닌 설정에서의 RAG의 생성 능력을 평가한다. Jeopardy는 어떤 개체에 대한 사실이 주어지고, 사실에 해당하는 답을 맞히는 형식이다. 정확하고 사실적인 진술이 요구되기에 answer entities를 조건으로 질문을 생성하는 것은 매우 어렵다고 한다.

결과를 요약해보면

1) Jeopardy 질문은 서로 다른 정보 조각을 포함하는 경우가 있는데, RAG-Token은 여러 문서에서 가져온 내용을 결합할 수 있기에 더 나은 성능을 보였다.

2) 제목의 일정 부분까지 보고난 후 모든 문서의 확률이 유사해지는 현상이 나온다. 즉, 검색 문서 없이도 다음 토큰을 예측한다. 이는 특정 지식이 parameter에 저장된다는 것을 알 수 있다. 즉, non-parametric memory가 생성을 유도하고 parametric-memory가 완성하는 방식으로 두 메모리가 협력함을 보인다.

Fact Verfication

FEVER task는 자연어 주장에 대해 위키피디아에 의해 지지/반박/정보 부족인지를 분류하는 태스크이다. 즉, 생성이 아닌 분류가 가능한지를 판단한다.

결과는 gold evidence(정답이 있는 파일)을 잘 가져오고, 설령 gold evidence가 없더라도 정확도가 크게 흔들리지 않았다.

추가 결과

1) BART보다 더 다양한 출력을 생성한다(RAG-Sequence > RAG-Token)

2) 제거 실험을 진행한 결과 학습된 검색이 성능을 향상시킴을 증명했고, FEVER에선 BM25와 같은 inverted-index를 retriever로 사용하는게 더 좋았지만, 그 외의 다른 태스크에선 differentiable retrievall이 더 효과적이었다.

3) Non-parametric memory의 교체만으로 모델의 지식을 업데이트 할 수 있었다.

4) 학습시 k에 따른 성능 차이는 크지 않았고, test-time에는 상황에 따라 다른 k값을사용할 수 있다(Retriever가 확률분포이기에). RAG-Sequence의 경우 k가 늘어날수록 성능이 단조증가했고(50근처까지), RAG-Token은 k가 일정 이상 증가하면 오히려 noise가 심해지며 성능이 감소했다.

Conclusion

생각보다 Titans때 처럼 엄청 식이 복잡하거나 그런 것은 거의없었다. 뭔가 이해가 될랑말랑해서 내가 쓴 정리본좀 다시 몇번 읽어봐야 될 것 같다. 그리고 아마 다음 SPLADE가 단점을 또 보완한 부분이 있을거라 이 방식의 어떤 단점이 있었을지도 고민해봐야겠다.

잡소리 간단하게만 보고 SPLADE까지 같이 정리하려고 했다. 어쩌다보니 생각보다 길어져서 그냥 분리시켰다.

BM25

무려 1994년 논문이시다. 하지만 DPR이 나온 아직까지도 병행해서 같이 쓰이곤 한다. 참고로 tf는 term-frequency, 한 문서 내에 특정 용어가 몇 번 등장했는가이다.

Introduction

1. Formal Model: 모델이 정확한 수식을 명시하는 방식.

• 검색 성능에 영향을 끼치는 매우 많은 변수들을 고려하기가 종종 어려움

2. Ad-hoc: 그럴듯해보이는 수식을 시도해본다.

• 특정 변수를 어떻게 다루는지에 대한 지침이 없음
• 수식을 추측해서 만들어보고 실험해야함

3. Regression: 독립 변수와 함수는 ad-hoc, 검색에서 가치를 평가하고 일부를 선택하여 가중치를 부여하는 Formal한 방식도 채용

Basic Probabilistic Weighting Model

(1) $w(\underline x) = log{{P(\underline x|R)P(\underline 0|\overline R)}\over{P(\underline x|\overline R)P(\underline 0|R)}}$

이때 $\underline x$는 문서에 대한 정보를 담은 벡터, \underline 0는 가중치가 0이되는 기준벡터, $R$과 $\overline R$은 관련/비관련 즉, $\underline 0$는 질문 용어가 문서 내에 전혀 등장하지 않은 문서 벡터로, 해당 값보다 $\underline x$가 큰지/작은지로 성능을 평가한다. $R$과 $\overline R$은 현재 질의를 기준으로 정답인가/아닌가를 구분한다.

간단하게 식을 요약해보자면 문서 x가 관련 문서(R)에 나타날 가능성이 비관련($\overline R$)에 나타날 가능성보다 얼마나 큰가를 설명한다. 이때 비교대상은 $\underline 0$이다. (2) $W = log{{p(1-q)}\over{q(1-p)}}$ = $log{{(r+0.5)/(R-r+0.5}\over{(n-r+0.5)/(N-n-R+r+0.5)}}$ 여기서 p = P(term present | R) q = P(term present | $\overline R$) N은 index된 전체 문서 수, n은 해당 용어를 포함하는 문서 수, R은 알려진 관련 문서 수, r은 R에서 해당 용어를 포함하는 문서 수 (3) 만약 R과 $\overline R$과 같은 관련성 정보가 없다면 ICF에 근사한다.(이를 W(1)로 정의한다) 문서 내 용어 빈도만 다룬다면 $W = log({{p_{tf}}\over{q_{tf}}}\cdot{{q_0}\over{p_0}})$ $p_{tf}$ = P(term present with frequency tf | R) $q_{tf}$ = P(term present with frequency tf | $\overline R)$ $p_0$, $q_0$는 용어가 전혀 나타나지 않을 확률

Eliteness

문서에서 특정 용어가 등장하는 현상이 무작위적이거나 확률적인 요소를 포함한다고 가정한다. 그리고 해당 확률은 그 문서가 실제로 그 개념에 관한 것인가/아닌가라는 구분을 반영한다. 이 중 이 문서가 그 개념에 관한 경우를 elite하다고 부른다.

용어 빈도는 eliteness에만 의존하지만, eliteness는 반드시 용어 빈도에 의존하지는 않는다. 기존의 term-independence assumption은 서로 다른 용어들에 대한 eliteness가 서로 독립적이라는 가정으로 대체한다.

2-푸아송 모델(2-Poisson Model)

자 일단 푸아송모델이 기억이 안날 수도 있다. 푸아송 모델은 $P(tf = k) = {{\lambda ^k e^{-k}}\over{k!}}$로 어떤 사건이 특정 구간(여기선 문서 전체)에 걸쳐 $\lambda$번 무작위로 발생한다는 확률 분포 모델이다. 일반 푸아송이 아닌 2-푸아송이 나온 이유는 elite문서와 non-elite문서를 구분하기 위함이다. 논문에 나오는 복잡한 식을 간략하게 줄여보면 $P(tf = k) = p(elite)Pois(\lambda)+P(non_elite)Pois(\mu)$이다. 즉, tf가 크다는 것을 문서의 eliteness로 엮어낸 것이다.

변형

2-푸아송 모델 자체는 식이 너무 복잡하기에 그대로 사용하기엔 문제가 많다. 따라서 2-푸아송 모델의 성질은 유지하되 더 간단한 식으로 변형한다.

2-푸아송 모델의 성질은

1) tf = 0일때 값이 0이다 2) tf가 증가하면 모델 값도 단조증가한다. 3) 그러나 어떤 점근적 최대값에 수렴한다. 4) 3)의 점근적 최대값은 eliteness를 직접적으로 나타내는 지표에 부여되는 Robertson/Sparck Jones 가중치에 근사한다.

이를 만족시키면서 훨씬 단순화시킨 식은 ${tf}\over{constant + tf}$로, (constant == 1인 경우다) 이때 tf>=0 이기 때문에, 시작은 0이면서 1에 점근한다. 또한 constant가 큰 상수값이 되면 거의 선형에 가깝게 증가한다. (constant == 50인 경우)

여기에 점근적 최대값을 변경하기 위해 적절한 가중치를 곱해줄 필요가 있다. 이를 Robertson/Sparck Jones 가중치(위의 Basic Probabilistic Weighting Model의 (2)번 식)에서 통상적인 추정치 W(1)을 사용하면 $w = {{tf}\over{k_1+tf}}\cdot W(1)$ ($k_1$은 1~2범위의 값이 적절해 보인다고 한다.) 다만 이 식은 Fuhr나 Cooper의 방법에 사용하기 적합하지 않다. 또한 핵심적으로 원래의 2-푸아송 그래프와 달리 좌상단 볼록(convex)하다는 차이점도 존재한다.

문서 길이

문서 길이가 다른 이유에 대해서도 두가지 가설이나 존재한다. 첫번째 Scope hypothesis는 긴 문서는 짧은 문서들보다 더 다루는 내용이 많다는 가설이고, Verbosity hypothesis는 긴 문서와 짧은 문서 모두 다루는 범위는 유사하지만 단지 더 많은 단어를 사용한다는 가설이다. 해당 논문에서는 Verbosity hypothesis를 가정한다. Verbosity hypothesis로 가정하면서 eliteness는 문서 길이와 독립적이지만, 특정 항에 대해 eliteness가 주어졌을 때 해당 항의 tf는 문서 길이에 의존한다.

이에 따라 문서가 길어질수록 tf는 커지지만, 단순히 "길어서 많이 나온 것"과 "내용상 중요해서 많이 나온 것"을 구분하려는 시도를 한다.

이제 식이 쭉 나오는데 핵심만 잡아보면

1) 가장 직관적인 모델(문서길이 d, $k_1$의 평균 길이 A) - term 기준 $w = {{tf}\over{k_1 \cdot {{d}\over{A}}+tf}} \cdot W(1)$ 즉, tf는 문서 길이로 나눠야한다. 2) 문서 자체가 긴지 짧은지를 고려해야 한다 - 문서 전체 기준 $k_2 \cdot n_q \cdot {{A-d}\over{A+d}}$ 즉 문서 길이는 평균 길이를 기준으로 평균보다 길면 패널티를 주고, 평균보다 짧으면 패널티가 덜해진다.

Query Term Frequency

이번엔 문서에서 tf가 아닌 유저가 한 질문에서 tf를 고려하는 방법이다. 문서는 위에서 말한 것 처럼 문서 길이에 따라서 나누고 패널티를 주고 복잡한 과정을 한다. 하지만 유저의 Query에 있는 tf는 어떤 노이즈도 아닌 '의도적'인 값에 가깝다. 그렇기에 그냥 $w = qtf \cdot W(1)$라는 아주 간단한 식이 완성된다.

결론

기존의 TF-IDF에서 TF가 클수록 중요한 단어라는 베이스를 바탕으로 진행한다. 여기에 BM25는 TF가 단순한 선형이 아닌 어느 순간부터 의미가 멈추도록(점근적 한계치) 설정하고, eliteness를 이용해 단순히 TF가 크다고 핵심이라고 끝낸 것에서 벗어나고, 문서의 길이를 노이즈로 구분해서 줄이는 방법을 채택했다.

개념적인 부분은 이전 글을 참조하자

Encoder 코드 살펴보기

에서 우리가 볼 것은 dpr/models/biencoder.py의 Biencoder 클래스와 dpr/models/hf_models.py이다. 코드를 다 적기엔 길기에 옆에 띄워두고 같이 보도록 하자

Biencoder.py

init

모델이 question(Question encoder)과 ctx(passage encoder) 두개가 있는 것을 확인할 수 있다. fix_...는 해당 모델을 고정할지 여부이다.

get_representation

Encoder 1개를 호출하고, 해당 encoder에 대한 출력값 3개를 반환한다. 이때 만약 fix된 모델이라면 no_grad로 가중치를 업데이트하지 않는다. 다만 encoder이후의 projection layer, loss는 requires_grad_를 통해서 학습되도록 한다. return

• sequence_output(토큰별 hidden state. [batch, seq_len, hidden]
• pooled_output(문장 요약 벡터([CLS] 토큰의 hidden state), [batch_size, hidden_dim]
• hidden_states(레이어별 히든 스택)

forward

Question vector와 Passage vector를 각각 뽑아내서 반환 단, 이때 q_pooled_out과 ctx_pooled_out만을 사용하면서 문장 전체에 대한 대표 벡터만 리턴

create_biencoder_input

해당 함수가 핵심 Objective: 학습용 batch 만들기. Question tensor은 [B, Lq], Passage tensor은 [Nctx, Lc], 각 질문의 positive 문서가 ctx_tensor의 어디 위치인지도 기록하는것이 목표

1) sample마다 positivie 1개 + negative 뽑기

positive_ctx = sample.positivie_passages[0] # shuffle_positives면 랜덤으로 1개 추출
neg_ctxs = sample_negative
hard_neg_ctxs = sample.hard_negative_passages

이때 만약 hard_neg_fallback = True이고 hard negative가 부족하다면 일반 negative에서 hard_negative 개수만큼 대체함

if hard_neg_fallback and len(hard_neg_ctxs) == 0
    hard_neg_ctxs = neg_ctxs[0:num_hard_negatives]

이를 원하는 개수(num_other_negatives or num_neg_ctxs)만큼 잘라서 최종 문서에선 all_ctxs = [positivie_ctx] + neg_ctxs + hard_neg_ctxs

2) ctx를 텐서로 바꾸고 정답 인덱스 기록

sample_ctxs_tensors = [
    tensorizer.text_to_tensor(ctx.text, title=ctx.title if (insert_title and ctx.title) else None)
    for ctx in all_ctxs
]

ctx_tensors.extend(sample_ctxs_tensors)

tokenization + padding으로 fixed-size tensor을 만들어줌

     positive_ctx_indices.append(current_ctxs_len)
    hard_neg_ctx_indices.append(...)

positive_ctx가 전체 ctx_tensor에서 시작한 위치가 positive index. 추가로 hard negative의 index도 기록

3) Question(Query) tensor화 query_token 옵션이 있다면 질문 앞에 특정 token을 붙임 특수 토큰([START_ENT])인 경우 span을 골라 넣음(특수 목적용 -> if query_token == "[START_END]"부분은 무시해도 된다.)

if query_token:
    # TODO: tmp workaround for EL, remove or revise
    if query_token == "[START_ENT]":
        query_span = _select_span_with_token(question, tensorizer, token_str=query_token)
        question_tensors.append(query_span)
    else:
        question_tensors.append(tensorizer.text_to_tensor(" ".join([query_token, question])))
else:
    question_tensors.append(tensorizer.text_to_tensor(question))

4) 최종 텐서 구성 passage tensor

ctxs_tensor = torch.cat([ctx.view(1, -1) for ctx in ctx_tensors], dim=0)

Question tensor

questions_tensor = torch.cat([q.view(1, -1) for q in question_tensors], dim=0)

class BiEncoderNllLoss

얘는 위의 함수들처럼 Biencoder class가 아닌 다른클래스인데, 간단하게 question 벡터와 ctx벡터에서 score를 계산한다. 이후 해당 score를 softmax하고, 이런 softmax_score를 바탕으로 loss를 계산한다는 softmax cross entrophy의 교과서적인 방법이다. 약간의 차이점으로는 class가 일반적인 라벨이 아닌 배치 내의 passage라는 것이다. 이 부분이 논문에서 말한 in-batch negatives이다.

hf_models.py

get_bert_biencoder_components

DPR의 Dual Encoder를 조립하는 부분이다.

question_encoder = HFBertEncoder.init_encoder(
    cfg.encoder.pretrained_model_cfg,
    projection_dim=cfg.encoder.projection_dim,
    dropout=dropout,
    pretrained=cfg.encoder.pretrained,
    **kwargs
    )
ctx_encoder = HFBertEncoder.init_encoder(
    cfg.encoder.pretrained_model_cfg,
    projection_dim=cfg.encoder.projection_dim,
    dropout=dropout,
    pretrained=cfg.encoder.pretrained,
    **kwargs
)

크게 특별한 부분은 없지만, question과 ctx가 모두 BERT 기반이지만 가중치가 별도로 구성되어있다. 그리고 이런 두 BERT를

biencoder = BiEncoder(question_encoder, ctx_encoder, fix_ctx_encoder=fix_ctx_encoder)

합친다. 이 부분이 DPR에서 말한 아이디어처럼 $E_q$와 $E_p$를 묶은 것이다. 그 밑에는 평범하게 AdamW로 optimizer하고, tokenziar하고 한다.

get_bert_reader_components

DPR의 두 단계에서 말한 부분 중 reader에 속하는 부분으로, 위에서 뽑은 top-k를 바탕으로 가져온 passage 중 정답 span을 뽑는 부분이다.

encoder = HFBertEncoder.init_encoder(
    cfg.encoder.pretrained_model_cfg,
    projection_dim=cfg.encoder.projection_dim,
    dropout=dropout,
    pretrained=cfg.encoder.pretrained,
    **kwargs
)

BERT encoder를 만들고

reader = Reader(encoder, hidden_size)

Span 추출용 헤드를 얹는다. 이후 위와 동일하게 optimizer하고 tensorizer한다.

Class HFBertEncoder

Huggingface에 있는 BertModel을 inherited(상속)하고 DPR에 필요한 형태로 적절하게 변형한다.(wrap한다고 한다.)

__init__

BERT model 초기화, projection차원 설정, 가중치 초기화

init_encoder

cfg = BertConfig.from_pretrained(cfg_name if cfg_name else "bert-base-uncased")

pre-training된 BERT를 그대로 가져오고

if dropout != 0:
    cfg.attention_probs_dropout_prob = dropout
    cfg.hidden_dropout_prob = dropout

그 위에 원하는 dropout을 덮어쓴다.

if pretrained:
    return cls.from_pretrained(cfg_name, config=cfg, project_dim=projection_dim, **kwargs)
else:
    return HFBertEncoder(cfg, project_dim=projection_dim)

이부분은 pre-training BERT와 scratch BERT 둘 다 실험하기 위한 부분이다.

forward

out = super().forward(
    input_ids=input_ids,
    token_type_ids=token_type_ids,
    attention_mask=attention_mask,
)

순수 BERT를 그대로 실행하고

sequence_output = out.last_hidden_state
pooled_output = None
hidden_states = out.hidden_states

if문은 Hugging face의 버전에 따라서 맞춰둔 설정이라 넘어가면, 해당 부분에서 sequence_output을 얻고

if isinstance(representation_token_pos, int):
    pooled_output = sequence_output[:, representation_token_pos, :]

위의 sequence_output을 바탕으로 DPR식으로 pooled representation을 생성한다. 해당 pooled_output을 바탕으로 유사도 계산을 진행한다.

결론

DRP의 코드를 아주 간단하게만 살펴보았다. 아마 아예 배경지식이 없는 경우(BERT도 모른다거나) 이해가 안될텐데, 그건 다른 글들을 찾아보길 바란다. 결론적으로 설명된 것 처럼 batch를 만들어내고, $E_q$와 $E_p$ 두 BERT를 묶어서 둘 다 학습시킨다는 점이 중요하다.