이전 문제의 계산을 재활용해 LLM의 다음 문제 정확도를 높이는 STAIR
Can Computation from Earlier Problems Help LLMs Solve New Ones?
무엇인가
한 대화 안에서 여러 문제를 연달아 푸는 상황은 흔하다. 각 문제는 서로 독립적이고 자기 정보만 갖고 있지만, 모델 입장에서는 앞선 문제와 그 답을 생성하면서 남긴 내부 계산이 어텐션을 통해 계속 보인다. 이 논문은 그 잔여 계산이 다음 문제 풀이에 도움이 되는지 해가 되는지를 묻고, 도움이 되도록 읽는 방법을 학습하는 STAIR를 제안한다.
어떻게 동작하나
저자들은 먼저 같은 문제 집합을 4턴 세션의 다른 위치에 놓고 측정했다. 첫 턴(Vanilla)에는 앞선 문제가 없고, 이후 턴(Native)에는 이전 문제와 어시스턴트 응답이 문맥에 남는다. Qwen3-4B Instruct의 MATH-500에서 Avg@4는 Vanilla 95.00%에서 Native T4 93.75%로 떨어졌다. 반대로 Qwen3.5-4B의 GPQA-Diamond에서는 63.64%에서 Native T2·T3·T4 각각 75.25%로 올랐다. 이력 유지가 모델과 과제에 따라 정확도를 올리기도 내리기도 한다는 뜻이다.
무엇과 다른가
왜 그런지 보려고 저자들은 가중치를 고정한 채 저장된 대화를 재생(replay)했다. 현재 문제의 질의가 이전 어시스턴트 응답 위치에 할당하는 어텐션 비중은 3층에서 39.49%와 38.29%, 19층에서도 12.41%와 12.87%로 남아 있었다. 상태 변위를 문제 효과 a(x), 이력 효과 b(H), 둘의 상호작용 e(x,H)로 분해하면 19층에서 상호작용은 전체 변위 에너지의 0.69%와 0.66%에 불과하다. 그런데도 상호작용 응답 기준 최근접 이웃 8개의 정체성은 서로 다른 이력 사이에서 48.06%와 46.50%가 겹쳤고, 이웃 정체성을 섞은 대조군은 11.73%와 12.08%에 그쳤다. 문제별로 반복되는 구조가 이력과 무관하게 존재한다는 관찰이다. 같은 문제에 대한 Native와 Vanilla 답을 짝지어 보면 T4에서 오답에서 정답으로 바뀐 경우 35건, 정답에서 오답으로 바뀐 경우 60건이 나와 Avg@4가 1.25%포인트 하락했다.
어떻게 쓰나
STAIR(Stale-Token Attention for Inter-query Reuse)는 이미 끝난 문제의 어시스턴트 토큰, 즉 stale 토큰의 키와 값을 읽기 전용 뱅크에 저장한다. 키는 키 투영과 정규화 뒤 RoPE 적용 전 상태로 잡고, 값은 값 투영 출력을 쓴다. 뱅크 키에는 0부터 M-1까지의 보조 위치를, 현재 토큰에는 M+p 위치를 부여해 별도 어텐션 경로를 만든다. 핵심은 질의를 반사(reflection)하는 것이다. 제어할 층과 헤드마다 반사면의 법선 벡터 n을 학습해 Householder 반사 q~R = q~ − 2u(u^T q~)(u는 정규화된 n)를 적용한다. 질의의 노름은 보존되고, 반사로 생긴 변화량 c = q~R − q~가 뱅크 위의 확률분포를 다시 가중한다. STAIR는 반사된 읽기 r^R에서 원래 질의의 기준 읽기 r^ref를 빼서 차분 Δr만을 네이티브 어텐션 출력에 더한다. 차분 가중치는 합이 0이라 같은 뱅크 위에서 어텐션 재분배가 만든 내용 변화만 반영된다. 학습되는 것은 반사 법선뿐이며 총 12,288개 파라미터다. 백본은 고정되고, 보조 경로는 프롬프트 처리(prefill) 중에만 작동하며 이후 자기회귀 디코딩은 원래 경로를 따른다. 손실은 −log p_STAIR(y|s) + λ·KL(p_Native‖p_STAIR), λ=1로, 네이티브 예측에서 너무 멀어지지 않게 묶는다. 제어 층은 0-based [3, 11, 19]로 세 모델 구성에서 고정했다.
전제와 한계
실험은 Qwen3-4B Instruct, Qwen3.5-4B, Qwen3.5-9B(사고 모드) 세 모델과 MATH-500, AIME 2025, AMC23†, GPQA-Diamond 네 벤치마크에서 이뤄졌다. 컨트롤러는 각 백본이 DAPO-Math-17k(학습 14,806문제, 검증 128문제)에서 생성한 응답으로 1에폭 학습했다. 문제당 4개 응답을 샘플링해 Avg@4와 Pass@4를 측정했다. 가장 큰 개선은 AIME 2025에서 나왔다. T2–T4 평균 Avg@4가 Native 대비 Qwen3.5-4B +11.67%포인트, Qwen3.5-9B +6.11%포인트, Qwen3-4B Instruct +3.61%포인트다. 세 모델 모두 이 벤치마크에서 각 후속 턴의 Avg@4와 평균 Pass@4를 함께 올렸다. MATH-500에서는 Qwen3.5 두 모델이 두 평균 지표 모두 Vanilla를 넘었고, Qwen3.5-4B는 AMC23†에서 Avg@4 +6.37%포인트를 냈다. GPQA-Diamond에서는 Instruct가 두 지표를 개선했지만 Qwen3.5 모델들은 비슷한 이득을 보이지 않았다.
T1 이력을 공유한 AIME 2025 T2 조건(Qwen3.5-4B)에서 STAIR는 Avg@4 48.33%로 Native 21.67%를 26.67%포인트 앞섰고 Pass@4는 56.67%에서 80.00%로 올랐다. 뱅크를 쓰지 않고 같은 12,288개 파라미터로 학습한 Bank-free 컨트롤러는 36.67%에 머물렀고, 키는 그대로 두고 값만 헤드 안에서 뒤섞으면 33.06%(3시드)로 떨어졌다. 학습된 반사를 무작위 반사로 바꾸거나 차분 대신 전체 반사 읽기를 더하는 변형도 성능이 낮았다. LoRA와의 비교에서는 Qwen3-4B Instruct와 AIME 2025 기준 STAIR가 12,288개 파라미터로 +3.61%포인트를 냈고, LoRA는 336배 많은 파라미터를 학습한다. 둘을 함께 학습하면 12,288개를 더해 LoRA 대비 Avg@4 +3.89%포인트, Pass@4 +4.44%포인트를 얻었다. 뱅크 길이를 앞 2K, 8K, 32K 토큰으로 제한하면 T2 Avg@4가 35.00%, 43.33%, 39.17%로 전체 사용 시 48.33%보다 낮았다. 전체 뱅크 읽기는 Native 대비 중앙값 0.81초와 최대 메모리 6.60GiB를 추가로 쓴다.
실무적으로 이 논문은 멀티턴 에이전트나 대화형 추론 서비스에서 이전 턴을 그냥 남겨두는 것이 항상 이득이 아니라는 점을 수치로 보여준다. 같은 도메인 안에서도 정확도가 오르거나 내릴 수 있으므로 대화 이력 유지 정책을 정할 때 벤치마크별로 확인해야 한다. STAIR 자체는 백본을 건드리지 않고 12,288개 파라미터만 학습하는 어댑터라 기존 서빙 스택에 얹기 쉬운 형태지만, 프롬프트 처리 단계에서 별도 K/V 뱅크를 읽어야 하므로 지연과 메모리 비용(중앙값 0.81초, 6.60GiB)을 감수할지 판단해야 한다. 또 학습된 반사 방향이 수학 문제로만 맞춰져 있어 GPQA-Diamond 같은 과학 추론에서는 이득이 재현되지 않았다는 점을 배포 전에 확인해야 한다.
저자들이 밝힌 한계는 다음과 같다. 반사 방향은 수학 문제에 맞춰 학습됐고 그 효과는 현재 질의와 이력 키에 의존하므로, 도메인별 편차가 학습 분포의 질의-이력 관계 때문인지는 더 넓은 학습 혼합으로 검증해야 한다. 현재 근거는 Qwen 계열의 4턴 세션에 한정된다. 세션이 길어지면 완료된 턴마다 새 상태가 쌓여 뱅크 저장량과 prefill 읽기 비용이 증가하며, 선택적 보존으로 뱅크 증가를 제한하는 방향이 남아 있다. 다양한 과제 전환 데이터로 사전학습해 반사 방향을 백본과 함께 학습하는 것, 그리고 SFT나 강화학습으로 추론 행동 변화에 맞춰 읽기를 적응시키는 것도 향후 과제로 제시된다.