FlexAttention은 PyTorch 팀이 2024년 7월에 공개한 새로운 API로, 몇 줄의 일반적인 PyTorch 코드로 많은 어텐션 변형을 구현할 수 있는 유연한 인터페이스를 제공합니다. torch.compile 이는 FlashAttention 커널을 융합하여 성능을 저하시키지 않고도 유연성을 제공합니다. 관련 논문 결과는 다음과 같습니다.효율적인 고해상도 시각-언어 모델을 위한 FlexAttention", ECCV 2024에 의해 승인되었습니다.

FlexAttention은 고해상도 시각 언어 모델의 효율성을 개선하도록 설계된 유연한 주의 메커니즘입니다. 제안된 메커니즘은 고해상도 및 저해상도 이미지 레이블을 모두 인코딩하고 저해상도 레이블과 선택된 몇 개의 고해상도 레이블만을 사용하여 어텐션 맵을 계산함으로써 계산 비용을 크게 줄입니다. 고해상도 랜드마크 선택은 고해상도 선택 모듈을 통해 수행되며, 이 모듈은 입력된 주의 맵을 기반으로 관련 지역의 랜드마크를 검색할 수 있습니다. 선택된 고해상도 토큰은 저해상도 토큰과 텍스트 토큰과 함께 계층적 자기 주의 계층에 입력되고, 이 계층에서 생성된 주의 맵은 다음 단계에서 고해상도 토큰 선택에 사용됩니다. 이 과정은 각 주의 계층에서 반복됩니다. 실험 결과 FlexAttention은 다중 모드 벤치마크에서 기존의 고해상도 시각-언어 모델보다 성능이 뛰어나고, 계산 비용을 약 40%만큼 크게 절감하는 것으로 나타났습니다.