Skip to main content
누빈트
QUICK REVIEW

[논문 리뷰] TokenUnify: Scaling Up Autoregressive Pretraining for Neuron Segmentation

Yinda Chen, Haoyuan Shi|arXiv (Cornell University)|2024. 05. 27.
Advanced Image and Video Retrieval Techniques인용 수 4
한 줄 요약

TokenUnify는 시각 자동회귀 모델링에서 오류 누적을 줄이기 위해 무작위, 다음, 다음-모든 토큰 예측을 결합한 새로운 혼합 토큰 예측 사전학습 방법을 도입한다. Mamba 아키텍처와 12000만 개 이상의 애너테이션 벡셀을 포함한 초고해상도 3D 전자현미경 데이터셋을 활용하여, 이전 방법들에 비해 뉴런 세그먼테이션에서 45% 향상된 성능을 달성하면서도 확장 가능하고 효율적인 장시계열 학습을 가능하게 한다.

ABSTRACT

Neuron segmentation from electron microscopy (EM) volumes is crucial for understanding brain circuits, yet the complex neuronal structures in high-resolution EM images present significant challenges. EM data exhibits unique characteristics including high noise levels, anisotropic voxel dimensions, and ultra-long spatial dependencies that make traditional vision models inadequate. Inspired by autoregressive pretraining in language models, we propose TokenUnify, a hierarchical predictive coding framework that captures multi-scale dependencies through three complementary learning objectives. TokenUnify integrates random token prediction, next-token prediction, and next-all token prediction to create a comprehensive representational space with emergent properties. From an information-theoretic perspective, these three tasks are complementary and provide optimal coverage of visual data structure, with our approach reducing autoregressive error accumulation from O(K) to O(sqrt(K)) for sequences of length K. We also introduce a large-scale EM dataset with 1.2 billion annotated voxels, offering ideal long-sequence visual data with spatial continuity. Leveraging the Mamba architecture's linear-time sequence modeling capabilities, TokenUnify achieves a 44% performance improvement on downstream neuron segmentation and outperforms MAE by 25%. Our approach demonstrates superior scaling properties as model size increases, effectively bridging the gap between pretraining strategies for language and vision models.

연구 동기 및 목표

  • 시각 자동회귀 사전학습에서 누적 오류를 해결함으로써 시각 모델의 스케일링 법칙을 제한하는 문제를 해결한다.
  • 특히 초고해상도 3D 데이터에 대해 장시계열 이미지 모델링에서의 계산 비효율성을 해결한다.
  • 이미지 내에서 공간적으로 관련된 장시계열을 구성하여 효과적인 자동회귀 학습을 가능하게 한다.
  • 언어와 시각 사전학습 간 격차를 해소하기 위해 시각 분야에서 확장 가능하고 자동회귀적인 전략을 가능하게 한다.
  • 대규모 고해상도 EM 데이터셋을 사용하여 뉴런 세그먼테이션을 위한 통합 벤치마크를 구축한다.

제안 방법

  • TokenUnify는 전역적 맥락 모델링을 향상시키기 위해 무작위 토큰 예측, 다음 토큰 예측, 다음-모든 토큰 예측의 세 가지 사전학습 작업을 통합한다.
  • 자기회귀 계산 복잡도를 이차원에서 선형으로 감소시키는 Mamba 아키텍처를 사용하여 장시계열 처리를 효율적으로 가능하게 한다.
  • 초고해상도 전자현미경 이미지에서 겹치지 않는 3D 패치를 사용하며, 이를 순차적 표현으로 토크나이즈한다.
  • 여섯 개의 기능적 마우스 뇌 영역에 걸쳐 12000만 개 이상의 애너테이션 벡셀을 포함한 대규모 3D EM 데이터셋을 구축하여 장시계열 학습을 지원한다.
  • 모든 세 가지 예측 작업을 통합한 유일한 목표 함수를 사용해 엔드 투 엔드로 사전학습 프레임워크를 훈련한다.
  • 학습 가능한 마스크 토큰과 커플런스 어텐션 디코더를 사용해 Mamba 기반 백본(EMmamba)을 다운스트림 뉴런 세그먼테이션에 피지컬러닝한다.
Figure 1 : TokenUnify prediction paradigms divide the 3D EM image into non-overlapping patches, which are tokenized into a sequence. Three tasks are performed for rich 3D image representations: (a) random token prediction, (b) next token prediction, and (c) next-all token prediction.
Figure 1 : TokenUnify prediction paradigms divide the 3D EM image into non-overlapping patches, which are tokenized into a sequence. Three tasks are performed for rich 3D image representations: (a) random token prediction, (b) next token prediction, and (c) next-all token prediction.

실험 결과

연구 질문

  • RQ1혼합 토큰 예측 전략은 시각 자동회귀 사전학습에서 누적 오류를 줄일 수 있는가?
  • RQ2여러 예측 목표를 조합함으로써 시각 모델의 표현 학습이 어떻게 향상되는가?
  • RQ3Mamba 아키텍처는 선형 복잡도로 자동회귀 시각 모델을 효과적으로 확장할 수 있는가?
  • RQ4대규모 고해상도 공간 연속 3D EM 데이터셋은 뉴런 세그먼테이션을 위한 더 나은 장시계열 모델링을 가능하게 하는가?
  • RQ5TokenUnify는 마스킹된 자동에코딩과 표준 자동회귀 방법에 비해 다운스트림 세그먼테이션 작업에서 얼마나 뛰어난 성능을 보이는가?

주요 결과

  • TokenUnify는 기존 방법들에 비해 기존 방법들에 비해 기존 방법들에 비해 다운스트림 EM 뉴런 세그먼테이션 작업에서 45%의 상대적 성능 향상을 달성했다.
  • 더 적은 파rameter를 사용함에도 불구하고 MAE에 비해 사전학습 성능에서 21% 향상된 성과를 기록했다.
  • 제거 실험을 통해 무작위, 다음, 다음-모든 토큰 예측 작업을 모두 조합할 경우 최고의 성능을 내며, VOI와 ARAND 점수는 각각 0.9951과 0.0509을 기록했다.
  • Mamba 기반의 EMmamba 모델은 TokenUnify 사전학습 전략을 사용할 경우 Wafer4 데이터셋에서 최신 기술 수준의 성능을 달성했으며, VOI는 0.2124 ± 0.0172, ARAND는 0.0967 ± 0.0097를 기록했다.
  • Mamba 아키텍처는 뛰어난 확장성과 효율성을 입증했으며, 기존 트랜스포머의 이차 복잡도에서 선형 복잡도로 감소시켜 장시계열 학습이 가능하게 했다.
  • 12000만 개의 애너테이션 벡셀을 포함한 대규모 EM 데이터셋은 지금까지 가장 큰 통합 뉴런 세그먼테이션 벤치마크를 제공한다.
Figure 2 : illustrates the main pretraining workflow of TokenUnify. The image $\mathbf{X}$ is fed into the Tokenizer, transforming it into a long sequence of tokens $\mathbf{x}_{i}|_{i=1}^{K}$ . The predictions for the random token, next token, and next-all token are performed sequentially. The Perc
Figure 2 : illustrates the main pretraining workflow of TokenUnify. The image $\mathbf{X}$ is fed into the Tokenizer, transforming it into a long sequence of tokens $\mathbf{x}_{i}|_{i=1}^{K}$ . The predictions for the random token, next token, and next-all token are performed sequentially. The Perc

더 나은 연구,지금 바로 시작하세요

논문 읽기부터 검토까지, 연구 시간을 획기적으로 줄여보세요.

카드 등록 없음 · 무료 플랜 제공

이 리뷰는 AI가 만들고, 인간 에디터가 검토했습니다.