The UNDO Flip-Flop: A Controlled Probe for Reversible Semantic State Management in State Space Model

1. 연구 배경 및 문제 제기 (Problem)

• 이론적 표현력과 학습 가능성의 괴리: 최근 연구 (Sarrof et al., 2024) 에 따르면, 2 층 구조의 상태 공간 모델 (SSM) 은 이론적으로 '별 없는 (star-free)' 순차적 태스크와 유계 계층 구조를 모델링할 수 있는 표현력 (Expressivity) 을 갖는 것으로 증명되었습니다. 그러나 이론적으로 표현 가능한 해법이 존재한다고 해서, 경사 하강법 (Gradient Descent) 기반의 최적화 알고리즘이 반드시 그 해법을 찾아내는 것은 아닙니다.
• 기존 벤치마크의 한계:
  • Flip-Flop 태스크: 단조로운 (monotonic) 상태 추적 및 덮어쓰기 메모리를 평가하지만, 이전 상태를 되돌리는 '가역적 (reversible)' 상태 복구는 평가하지 못합니다.
  • Dyck 언어: 구조적 중첩 (nesting) 을 평가하지만, 상태의 의미적 내용 (semantic payload) 을 복원하는 능력은 평가하지 않습니다.
• 핵심 문제: 자연어 처리 (수정, 대화, 추론 등) 에서는 이전 업데이트를 취소하고 과거 상태를 복구하는 가역적 의미 상태 검색 (Reversible Semantic State Retrieval) 능력이 필수적입니다. 기존 벤치마크는 이 능력을 격리하여 평가하지 못하므로, SSM 이 실제로 이러한 복잡한 상태 관리 능력을 학습하는지 여부가 불분명했습니다.

2. 방법론 (Methodology)

2.1 제안된 태스크: UNDO Flip-Flop

기존 Flip-Flop 태스크에 [UNDO] 연산자를 추가하여 가역적 상태 업데이트를 평가하는 새로운 태스크를 제안했습니다.

• 작동 원리: 모델은 암시적인 유계 스택 (bounded stack) 을 유지해야 합니다.
  • Write (Push): 스택에 값을 푸시하고 현재 상태를 업데이트.
  • Ignore (No-op): 상태 변경 없음.
  • [UNDO] (Pop): 최근 쓰기 작업을 취소하고 스택에서顶层 (top) 을 제거하여 이전 상태로 복원.
  • Read (Peek): 현재 스택의 최상단 값을 출력.
• 복잡도: 이 태스크는 정규 언어를 넘어 Dyck 언어와 유사한 괄호 규칙을 따르며, 동적 취소 하에서 의미적 상태를 검색해야 하므로 계산 복잡도가 증가합니다.

2.2 실험 설정

• 모델: Mamba-2 (1 층 및 2 층 버전) 를 사용. Mamba-2 의 스칼라 - 항등 (scalar-identity) 구조가 이론적 표현력에 제약을 줄 수 있는지 확인.
• 데이터 생성:
  • In-Distribution (ID): 1~50 토큰 (학습 및 평가).
  • Out-of-Distribution (OOD): 51~100 토큰 (일반화 능력 평가).
  • 데이터 분포: Write 와 UNDO 가 드물게 발생하고 Ignore 가 우세하도록 설정하여 학습 가능성을 확보하되, UNDO 가 충분히 학습되도록 함.
• 메커니즘 프로빙 (Mechanistic Probing):
  • Aggressive UNDO Pressure Test: 학습 길이 범위 (50 토큰) 내에서 구조적 복잡성을 극대화 (깊은 스택 쌓기 후 연속적인 취소) 하여 모델의 내부 표현을 분석.
  • 인과적 절단 (Causal Ablation): Deep History Loss Rate(과거 상태 접근 실패) 와 Local Toggle Heuristic Rate(현재 비트 반전 의존) 를 측정하여 모델의 실제 작동 방식을 규명.

3. 주요 결과 (Results)

3.1 성능 저하

• 단조로운 태스크 (기존 Flip-Flop): 1 층 및 2 층 Mamba-2 모두 100% 에 가까운 정확도로 빠르게 수렴하며, OOD 일반화도 성공적임.
• UNDO Flip-Flop 태스크:
  • 학습 한계: 1 층 (95.95%) 과 2 층 (98.55%) 모두 ID 데이터에서 높은 정확도를 보였으나, 완벽하게 수렴하지 못함.
  • 일반화 실패: OOD 데이터에서 성능이 급격히 하락 (2 층 모델 87.35%, 1 층 모델 79.10%). 이는 기존 Flip-Flop 태스크의 일반화 실패 (1% 미만) 와 대조적임.
  • 적대적 테스트 붕괴: 'Aggressive UNDO Pressure Test'에서 2 층 모델의 정확도는 **41.10%**로 떨어졌으며, 이는 무작위 추측 (50%) 보다도 낮은 수치입니다. 이는 모델이 단순히 추측하는 것이 아니라 체계적으로 오답을 내고 있음을 의미합니다.

3.2 메커니즘 분석 (원인 규명)

• 저장 (Storage) vs 검색 (Retrieval): Deep History Loss Rate가 매우 낮아 (2.17%) 정보 저장 능력 자체는 문제가 아님을 확인.
• 단순화 학습 (Shortcut Learning): 모델은 진정한 스택 기반 롤백 메커니즘을 학습하지 못했습니다. 대신 **로컬 토글 휴리스틱 (Local Toggle Heuristic)**을 학습했습니다.
  • 휴리스틱 작동 방식: [UNDO] 명령을 "현재 상태 비트를 반전 (Invert)"하는 것으로 오해하여 처리.
  • 실패 원인: 연속된 동일한 비트나 깊은 롤백이 발생할 경우, 단순 비트 반전 논리는 과거 상태를 복구하지 못해 실패합니다.
  • 통계적 우연: 희소한 학습 데이터 분포에서는 비트 반전이 스택 팝과 동일한 결과를 내므로 모델이 이 '단순한 길 (shortcut)'을 선택하게 됨.

4. 주요 기여 (Key Contributions)

1. UNDO Flip-Flop 태스크 제안: 가역적 의미 상태 검색 능력을 격리하여 평가할 수 있는 새로운 벤치마크를 제시.
2. 이론적 표현력과 학습 가능성의 분리 증명: Mamba-2 가 이론적으로는 스택 기반 롤백을 표현할 수 있음 (Sarrof et al. 의 정리 6 에 의거) 에도 불구하고, 경사 하강법 최적화는 이를 학습하지 못하고 더 단순한 휴리스틱에 수렴함을 실증적으로 보임.
3. SSM 의 한계 규명: SSM 이 비단조적 (non-monotonic) 상태 관리 태스크에서 체계적인 실패를 겪으며, 이는 저장 능력의 부족이 아닌 검색 (retrieval) 및 알고리즘적 일반화 실패임을 규명.

5. 의의 및 시사점 (Significance)

• 이론과 실전의 간극: "무엇을 표현할 수 있는가 (Expressibility)"와 "무엇을 reliably 학습하는가 (Learnability)"는 별개의 문제임을 명확히 함. 이론적 분석만으로는 최적화 알고리즘이 실패하는 패턴을 예측할 수 없음.
• 실제 NLP 태스크에 대한 경고:
  • 음성 수정 및 대화: 화자가 실수를 수정하거나 이전 상태로 돌아가는 상황 (Speech repair, dialogue correction) 에서 모델이 상태를 '반전'시키는 오류를 범할 수 있음.
  • 추론 (Reasoning): 체인 - 오브 - 씽킹 (Chain-of-thought) 과정에서 잘못된 중간 결론을 버리고 이전 단계로 돌아가야 하는 경우, 모델이 체계적으로 실패할 가능성이 높음.
• 향후 연구 방향:
  • Mamba-2 의 스칼라 - 항등 구조가 표현력 자체를 제한하는지, 아니면 학습 동역학의 문제인지에 대한 형식적 분석 필요.
  • 학습 분포를 변경 (취소 빈도 증가 등) 하여 휴리스틱을 억제하고 진정한 상태 관리 메커니즘을 학습시킬 수 있는지 탐구 필요.

결론

이 논문은 SSM 이 복잡한 계층적 구조를 이론적으로 다룰 수 있음에도 불구하고, 실제 학습 과정에서는 국소적인 비트 반전 휴리스틱에 의존하여 가역적 상태 복구에 실패함을 보여줍니다. 이는 모델의 아키텍처적 표현력과 실제 최적화 성능 사이의 중요한 간극을 드러내며, 자연어 처리에서 비단조적 추론이 필요한 태스크에 대한 SSM 의 신뢰성에 대해 재고찰을 요구합니다.