[논문 리뷰] PRISM: Parallel Residual Iterative Sequence Model

Jie Jiang, Ke Cheng외 9명|arXiv (Cornell University)|2026. 02. 11.

Parallel Computing and Optimization Techniques인용 수 0

한 줄 요약

PRISM은 다단계 반복 보정을 모방하는 병렬화 가능 amortized 잔여 최적화 프레임워크를 도입하여 선형 주의에서 높은 표현력을 달성하고 명시적 최적화 기반 방법보다 최대 174배 높은 처리량을 달성합니다.

ABSTRACT

Generative sequence modeling faces a fundamental tension between the expressivity of Transformers and the efficiency of linear sequence models. Existing efficient architectures are theoretically bounded by shallow, single-step linear updates, while powerful iterative methods like Test-Time Training (TTT) break hardware parallelism due to state-dependent gradients. We propose PRISM (Parallel Residual Iterative Sequence Model) to resolve this tension. PRISM introduces a solver-inspired inductive bias that captures key structural properties of multi-step refinement in a parallelizable form. We employ a Write-Forget Decoupling strategy that isolates non-linearity within the injection operator. To bypass the serial dependency of explicit solvers, PRISM utilizes a two-stage proxy architecture: a short-convolution anchors the initial residual using local history energy, while a learned predictor estimates the refinement updates directly from the input. This design distills structural patterns associated with iterative correction into a parallelizable feedforward operator. Theoretically, we prove that this formulation achieves Rank-$L$ accumulation, structurally expanding the update manifold beyond the single-step Rank-$1$ bottleneck. Empirically, it achieves comparable performance to explicit optimization methods while achieving 174x higher throughput.

연구 동기 및 목표

• 긴 시퀀스에 대한 선형 주의의 표현력 병목 현상을 해결한다.
• 효율적인 선형 모델과 표현력이 있는 최적화 기반 방법 사이의 격차를 메운다.
• 병렬로 다단계 보정을 가능하게 하는 하드웨어 인식(hardware-aware) 아키텍처를 개발한다.
• Rank Accumulation과 Write-Forget Decoupling 원리를 이론적으로 특성화한다.
• 강력한 기준선과의 비교를 통해 장기 시퀀스 추천 벤치마크에서 PRISM을 실증적으로 검증한다.

제안 방법

• 잊힘 다이내믹스를 저랭크로 유지하고 고랭크 비선형 보정을 주입 항으로 밀어 넣기 위해 Write-Forget Decoupling을 제안한다.
• S_{t-1}k_t를 추정하는 ShortConv 앵커와 다단계 보정을 생성하는 학습 예측기를 두 단계 프록시로 가지는 Input-Anchored Loop Unrolling 아키텍처를 도입한다.
• 게이트드 잔차 업데이트를 사용한 반복 보정을 통해 L개의 직교 rank-1 구성요소의 합으로 고랭크 주입 B_t를 구성한다.
• 축적된 B_t를 순환 상태에 주입하는 동시에 병렬 스캔 효율성을 유지하기 위해 상태 독립적 망각 연산자 A_t를 유지한다.
• 스펙트럴 섭동 하에서 Rank Accumulation을 이론적으로 제시하고 망각 vs. 주입 구성요소의 안정성을 분석한다.
• PRISM을 Transformers, 선형 베이스라인, 및 최적화 기반 솔버와 실험적으로 비교하고 정확도와 학습 처리량을 보고한다.

Figure 1 : The PRISM Architecture. The framework operates in two phases to approximate the Ideal Non-Linear Solver within a parallelizable linear recurrence. Phase 1 (Input-Anchored Simulation): A ShortConv anchor captures the local pre-activation proxy ( $u_{t}\approx S_{t-1}k_{t}$ ). Parallel pred

실험 결과

연구 질문

• RQ1 아모타이즈된, 입력 고정 보정이 병렬성을 유지하면서 TTT와 같은 명시적 반복 솔버의 성능에 맞출 수 있는가?
• RQ2 PRISM이 장기 시퀀스에서 명시적 최적화 기반 방법보다 더 높은 처리량을 달성하되 모델링 충실도를 손상시키지 않는가?
• RQ3 성능에 필수적인 반복적 고랭크 주입인가, 그리고 그 구성요소들(앵커, 게인 예측기, 반복 깊이)이 실질적으로 기여하는가?

주요 결과

• PRISM은 도전적인 벤치마크에서 명시적 반복 솔버 및 딥 트랜스포머와 비슷한 성능을 달성한다.
• PRISM은 학습 처리량에서 명시적 최적화 기반 방법에 비해 최대 174배의 처리량을 달성한다.
• PRISM은 quadratic Transformers와의 격차를 좁혀 아모타이즈드 보정의 표현력이 상당함을 시사한다.
• 소거 연구(Ablation studies)에서 반복 깊이, 비선형성, 앵커링(anchor), 게이팅이 모두 성능에 의미 있게 기여하는 것으로 나타났다.
• 메카니즘적 탐색(mechanistic probing)은 제약된 설정에서 PRISM이 선형 베이스라인을 비선형 작업에서 능가할 수 있음을 보여준다.]
• table_headers:[]
• table_rows:[]

Figure 2 : Training throughput comparison of 0.13B models on a single H20 GPU.