BlockRank는 희소 어텐션과 지도학습한 어텐션 출력으로 후보 문서의 순위를 매긴다. 각 문서는 공통 접두부와 자신을 읽고, 마지막 질의 토큰이 모든 후보를 비교하므로 답변 토큰 디코딩 없이 순위를 계산할 수 있다.
후보 배치
1단계 검색기가 \(D_q=(d_1,\ldots,d_N)\)을 제공한다. 순위 모델은 문서 앞에 지시문, 뒤에 마지막 질의 구간을 배치한다.
\[P=I\;\Vert\;d_1\;\Vert\;\cdots\;\Vert\;d_N\;\Vert\;q.\]접두부와 마지막 질의의 역할은 다르다. 접두부에 이미 질문이 들어갈 수 있다. 문서는 뒤의 질의 구간을 읽지 않아도 질의를 조건으로 표현되며, 마지막 질의 토큰은 점수를 읽어 내는 위치가 된다.
안정적인 식별자로 후보를 구분하고 점수를 반환 문서에 연결한다. 재순위화는 빠진 문서를 복구하지 못하므로 1단계 재현율이 전체 결과를 제한한다. Scalable In-context Ranking with Generative Models.
방향이 있는 어텐션
질의·키 행렬 \(Q,K\)와 값 행렬 \(V\)를 사용하는 하나의 헤드에서 어텐션은 다음과 같다.
\[\operatorname{Attn}(Q,K,V;M) =\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V.\]허용 쌍의 마스크는 0, 나머지는 \(-\infty\)다. 행이 열을 읽으며 구간은 지시문–문서–질의 순서다.
블록 마스크는 다음과 같다.
\[M=\begin{bmatrix} C_I&-\infty&-\infty&\cdots&-\infty\\ 0&C_{d_1}&-\infty&\cdots&-\infty\\ 0&-\infty&C_{d_2}&\cdots&-\infty\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 0&0&0&\cdots&C_q \end{bmatrix},\]대각 블록 \(C\)는 구간 내부의 인과성을 유지하고, 0과 \(-\infty\)는 크기가 맞는 블록 전체다. 지시문은 앞선 지시문을, 문서는 접두부와 자기 과거를 읽는다. 마지막 질의 토큰은 접두부·모든 문서·자기 과거를 읽는다.
문서는 다른 문서나 마지막 질의를 읽지 않는다. 이 비대칭성이 정보 흐름을 결정한다.
복잡도
접두부 길이 \(p\), 마지막 질의 길이 \(q\), 문서 길이 \(\ell_1,\ldots,\ell_N\)에서 전체 길이는 \(L=p+q+\sum_i\ell_i\)다. 특징 차원을 제외한 밀집 어텐션의 층별 토큰 쌍 연산량은 \(O(L^2)\)이다.
블록 구조에서 허용하는 토큰 쌍의 수는 다음 차수이다.
\[p^2+\sum_{i=1}^{N}(\ell_i^2+p\ell_i) +q\left(p+q+\sum_{i=1}^{N}\ell_i\right).\]문서 길이가 \(\ell\)로 제한되고 \(p,q\)가 고정이면 다음과 같다.
\[O\big(N\ell^2+Np\ell+Nq\ell+p^2+q^2+pq\big),\]이 조건 아래에서만 \(N\)에 대해 선형이다. 개별 문서의 셀프 어텐션은 문서 길이에 여전히 제곱으로 증가하므로 길이를 두 배로 해도 선형 비용은 유지되지 않는다.
\(p=q=32\), \(N=100\), \(\ell=128\)이면 단순한 전체 블록 계산에서 밀집 어텐션은 약 \((12{,}864)^2\approx165.48\)백만 쌍, 허용 쌍은 \(2.46\)백만이다. 67배의 속도 향상을 예측하는 값은 아니다. 투영·순전파 층·패딩·커널·메모리 접근·실행 호출 비용이 남는다.
계산을 줄이려면 희소·청크 커널로 금지 블록을 만들지 않아야 한다. 밀집 행렬을 계산한 뒤 마스킹해도 쌍 연산은 줄지 않는다.
위치 인코딩
BlockRank는 지시문 다음의 문서 내부 위치를 여러 블록이 공유하고 마지막 질의 영역은 따로 두어 목록 위치 편향을 줄인다.
식별자·잘라내기·패딩·구현·수치적 동점은 순열 불변성에 영향을 줄 수 있다. 접두부에 질의가 있으면 문서 표현도 질의에 의존하므로 모든 질의를 위한 공통 캐시를 만들 수 없다.
어텐션 지도학습
층 \(\ell^*\)에서 신호 전달 질의 토큰 \(S\)가 문서 토큰의 점수를 계산한다. 전달 토큰 \(i\)와 문서 토큰 \(j\)에 대해 문서 토큰만으로 정규화한다.
\[z_{ij}=\frac{Q_iK_j^\top}{\sqrt{d_k}},\qquad \alpha'_{ij}=\frac{e^{z_{ij}}}{\sum_{j'\in T_D}e^{z_{ij'}}}.\]문서의 질량을 다음과 같이 정의한다.
\[s_k=\sum_{i\in S}\sum_{j\in T_{d_k}}\alpha'_{ij}.\]전달 토큰의 평균을 쓰면 배율과 실질적인 손실 온도가 바뀐다. 합 정의에서는 \(\sum_k s_k=\vert S\vert\)이다. 질량은 후보 사이의 상대적인 순위를 매긴다. 정답 관련성이 같은 값만큼 높다는 뜻이나 보정된 관련성 확률인 것은 아니다.
정답 양성 문서 \(d^+\)가 주어졌을 때 대조 손실은 다음과 같다.
\[\mathcal L_{\mathrm{aux}} =-\log\frac{\exp(s_+/\tau)}{\sum_k\exp(s_k/\tau)}, \qquad\tau>0.\]학습에서는 답변의 다음 토큰 예측 손실과 이 보조 손실을 결합한다.
\[\mathcal L=\mathcal L_{\mathrm{NTP}}+\lambda\mathcal L_{\mathrm{aux}},\qquad\lambda\ge0.\]보조 손실은 정답 양성을 선호하도록 출력을 학습하고 다음 토큰 예측은 답변 출력을 학습한다. 이 지도학습이 어텐션에 순위화 역할을 부여하며, 임의의 사전학습 어텐션에는 같은 보장이 없다. 토큰 정규화와 문서 집계에는 길이 효과가 생길 수 있으므로 길이 분포와 잘라내기가 영향을 준다.
직접 출력
선택한 층까지만 추론해 질량을 계산하고 \(\arg\max_k s_k\)나 정렬 목록을 반환한다. 자기회귀 순위 토큰과 이후 층을 생략할 수 있다.
유지한 모든 후보의 문서 처리와 질의→문서 어텐션은 필요하다. 일찍 점수를 읽으면 뒤의 층과 디코딩을 줄이지만 후보 목록을 읽는 비용은 남는다. 코퍼스 규모 검색에는 1단계 인덱스가 따로 필요하다.
점수 \(N\)개의 전체 정렬은 \(O(N\log N)\), 최댓값 선택은 \(O(N)\)이다. 모델의 선형 점수 계산과 출력 정렬은 별도 비용이다.
보고된 성능
저자는 문서 100개의 MSMARCO에서 \(4.7\times\) 빠른 추론, 약 100K 문맥 토큰인 문서 500개를 1초 이내에 처리한 결과를 보고했다. 특정 모델·하드웨어·형식·기준선의 측정이며 여기서 재현하지 않았다. 실험 설정.
밀집 미세조정, 지도학습 없는 희소 어텐션, 지도학습한 어텐션 출력, 자기회귀 출력은 다른 기여를 분리한다. 후보·길이·모델·학습 데이터·정밀도·하드웨어·출력을 맞춰 nDCG/MRR, p50/p95 지연, 최대 메모리, 1단계 재현율을 비교할 수 있다.
지식 그래프 블록과 모바일 검색에는 새로운 표현·커널·대응 실험이 필요하다. 보고된 텍스트 결과가 이런 확장을 입증하지는 않는다.
의사결정 점수
BlockRank는 타입이 정해진 의사결정 모델처럼 산문 생성 없이 점수를 추출한다. 마지막 질의로 문서를 비교하지만, 여러 질문을 처리하는 서비스는 공유 상태의 질문들을 격리할 수 있다. 출력이 비슷해도 정보 흐름이 같지는 않다.
희소 어텐션은 통신과 계산을 정한다. 보정과 후속 임계값에는 별도의 학습·평가가 필요하다.