기술 요약: One for All: 비선형 트랜스포머가 문맥 내 강화학습을 위한 교차 도메인 일반화를 가능하게 함

문제 제기

강화학습 (RL) 의 핵심적인 도전 과제는 학습된 특정 작업을 넘어 일반화되는 모델을 개발하는 것이다. 전통적인 다중 작업 및 메타-RL 접근법은 이를 해결하려 하지만, 최근 연구에서는 명시적 매개변수 업데이트 없이 문맥을 통해 새로운 작업에 적응하는 문맥 내 학습을 위해 트랜스포머 아키텍처를 활용하는 방안을 탐구해 왔다. 그러나 트랜스포머가 어떻게 그리고 왜 작업 간, 특히 서로 다른 도메인 간에 일반화되는지에 대한 원칙적인 이해는 여전히 제한적이다. 왕 (Wang) 등 (2025) 의 기존 이론적 프레임워크와 같은 것들은 강한 선형성 가정 (선형 마르코프 보상 과정 및 선형 트랜스포머) 에 의존하여, 가치 함수가 비선형인 복잡하고 현실적인 설정에서의 적용 가능성을 제한한다. 본 논문은 문맥 내 RL 의 메커니즘, 특히 정책 평가를 이해하는 데 있어 존재하는 간극을 해소하여, 고정된 가중치를 가진 단일 트랜스포머가 여러 도메인 간에 일반화할 수 있는지 규명하는 데 초점을 맞춘다.

방법론

저자들은 커널 기반 관점에서 문제에 접근하여 비선형 트랜스포머와 커널 기반 시간차 (TD) 학습 간의 연결을 확립한다.

이론적 프레임워크

1. 함수 연산자로서의 트랜스포머: 본 논문은 트랜스포머를 문맥 (전이 시퀀스) 을 작업별 함수로 매핑하는 함수 연산자로 해석한다.
2. RKHS 해석: 비선형 어텐션 메커니즘이 커널 함수를 유도한다는 관찰 (청 (Cheng) 등, 2024) 에 기반하여, 저자들은 비선형 트랜스포머가 이 커널에 의해 유도된 재생 커널 힐베르트 공간 (RKHS) 에서 회귀를 수행한다고 가정한다.
3. 커널 TD 업데이트 구성: 핵심 이론적 기여는 고정된 가중치를 가진 깊이 $L$ 의 비선형 트랜스포머가 순전파 (forward pass) 만을 통해 커널화된 반경사 (semi-gradient) TD 업데이트의 $L$ 회 반복을 구현할 수 있음을 명시적으로 보여주는 구성이다.
   • 입력 표현: 입력 행렬 $Z_0$는 문맥 전이 (상태, 보상) 와 쿼리 상태를 인코딩한다.
   • 메커니즘: 트랜스포머는 문맥 상태를 커널 중심 (kernel centers) 으로 사용한다. 각 트랜스포머 레이어는 TD 업데이트의 한 단계에 해당한다. 구체적으로, 어텐션 메커니즘은 커널 함수 $\kappa(s_j, s_i)$로 가중된 TD 잔차 ($\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$) 의 가중 합을 계산한다.
   • 가중치 구성: 저자들은 두 개의 헤드 어텐션 블록에 대한 키 ($K$), 쿼리 ($Q$), 값 ($V$) 투영 행렬의 명시적 구성을 제공한다. 한 헤드는 현재 가치 추정의 차이를 계산하고, 다른 헤드는 부트스트랩된 타겟을 계산하여, 실제로 다음 업데이트를 수행한다:
     $$v^{(i)}_{k+1} = v^{(i)}_k + \alpha_k \sum_{j=0}^{n-1} \delta^{(j)}_k \kappa(s^{(j)}, s^{(i)})$$
   • 출력: 쿼리 상태의 값은 $L$ 레이어 후 출력 행렬의 오른쪽 아래 항목에서 읽어낸다.

한계 및 조건

본 논문은 이 구성에 대한 특정 조건과 한계를 식별한다:

• 공유 RKHS: 교차 도메인 일반화는 이론적으로 서로 다른 도메인의 가치 함수가 트랜스포머의 고정된 커널에 의해 유도된 동일한 RKHS 내에 있을 때만 보장된다. 만약 도메인들이 근본적으로 다른 커널 특성 (예: 매끄러운 대역 대 날카로운 가치 지형에 대한 서로 다른 대역폭) 을 요구한다면, 단일 고정 커널은 둘 모두를 잘 표현할 수 없다.
• 상수 오프셋: 이 구성은 절대적 가치 추정치에 상태와 무관한 상수 편향을 도입한다. 그러나 저자들은 이 편향이 상대적 비교 및 오차 계산에서 상쇄되므로 정책 개선 (예: 정책 경사) 이나 TD 오차 수렴에 영향을 미치지 않는다고 주장한다.
• 고정 커널 매개변수: 커널 매개변수 (예: 온도 $\delta$) 는 고정되어 있으며 문맥 내에서 적응되지 않는다. 이는 다양한 커널 특성을 요구하는 작업을 처리하는 모델의 능력을 제한한다.

주요 기여

1. 문맥 내 정책 평가에 대한 커널 관점: 본 논문은 커널 기반 RL 의 렌즈를 통해 문맥 내 정책 평가를 프레임하며, 트랜스포머 추론을 커널화된 TD 업데이트와 연관시킨다.
2. 명시적 가중치 구성: 고정된 가중치를 가진 비선형 트랜스포머가 문맥 상태를 커널 중심으로 사용하여 반복적 커널 TD 업데이트를 구현할 수 있음을 보여주는 이론적 구성을 제공한다.
3. 교차 도메인 재사용을 위한 충분 조건: 작업들의 가치 함수가 공유된 RKHS 내에 있을 때 교차 도메인 일반화가 가능함을 확립한다.
4. 실증적 검증: 저자들은 이 구성을 실제 세계 작업에 구현하여 이론적 예측을 검증했다.

실험 결과

저자들은 다양한 조작 도메인 (예: Pick-Place, Shelf-Place, Button-Press) 으로 구성된 벤치마크인 MetaWorld(유 (Yu) 등, 2021) 에서 방법을 평가했다.

• 단일 도메인 일반화: 고정된 가중치를 가진 트랜스포머 (한 도메인 내 고잡음 작업에서만 학습됨) 는 동일한 도메인 내의 중간 및 저잡음 작업과 다른 도메인 (Pick-Place, Pick-Place-Wall, Shelf-Place, Plate-Slide) 에서 TD 손실에 대해 성공적으로 수렴했다.
• 교차 도메인 일반화: 한 도메인 (예: Pick-Place-v3) 에서 학습된 모델은 미세 조정 없이 다른 도메인에 적용되었다. 이 모델은 Pick-Place, Pick-Place-Wall, Shelf-Place, Plate-Slide 간에 성공적으로 일반화되었다. 그러나 Button-Press-v3로의 전이는 실패했다.
  • 분석: 이 실패는 공유 RKHS 에 관한 이론적 한계와 일치한다. 하이퍼파라미터 분석 결과, Button-Press 는 나머지 네 가지 도메인과 비교하여 현저히 다른 커널 온도 (활성화 매개변수) 를 요구함이 밝혀졌으며, 이는 그 가치 함수가 다른 RKHS 에 있음을 시사한다.
• 선형 대 비선형: 왕 등 (2025) 이 제안한 선형 트랜스포머는 MetaWorld 작업에서 수렴에 실패하여, 복잡한 가치 함수 구조를 포착하려면 비선형성이 필요함을 확인했다.
• 합성 환경: 유도된 RKHS 내에 실제 가치 함수가 명시적으로 구성되도록 설계된 통제된 합성 MRP 에서 트랜스포머는 가치 함수의 모양을 정확하게 복원하여 커널 회귀 메커니즘을 검증했다.

중요성 및 주장

본 논문은 비선형 트랜스포머를 사용한 문맥 내 RL 이 도메인 간에 일반화될 수 있음을 보여주는 최초의 이론적 구성을 제공한다고 주장한다. 트랜스포머를 커널 회귀 수행으로 해석함으로써, 저자들은 단일 가중치 세트가 여러 도메인에서 학습할 수 있는 이유에 대한 원칙적인 설명을 제공한다: 작업들이 공통된 RKHS 를 공유할 때 가능하다.

이 연구는 보다 보편적인 정책을 개발하기 위한 기초적인 단계로 기능한다. 고정된 가중치를 가진 트랜스포머가 일반화할 수 있지만, 이 능력은 유도된 커널의 표현력에 의해 제한된다는 점을 명확히 한다. 저자들은 정책 평가가 필수적인 기초 단계에 초점을 맞추고 있으며, 정책 개선과 적응형 커널 구조는 향후 연구 과제로 남긴다고 명시적으로 밝혔다. 그들은 모든 RL 작업에 대한 일반화 문제를 해결했다고 주장하는 것이 아니라, 그러한 일반화가 가능한 이론적 조건을 확립했다고 주장한다.