기술 요약: 프로그램 합성으로서의 무손실 텐서 압축 (Brevis)

문제 정의

대규모 언어 모델(LLM) 및 기타 딥러닝 아키텍처의 급격한 확산은 모델 체크포인트의 수와 크기의 기하급급적인 증가를 초래했습니다. 이러한 자산을 저장하고 전송하며 배포하는 데 드는 비용이 점점 더 높아지고 있습니다. gzip, zstd와 같은 범용 압축기들이 널리 사용되고 있지만, 이들은 텐서 가중치를 일반적인 바이트 시퀀스로 취급하며 텐서의 특정 구조적 특성(예: 데이터 타입, 형상, 부동 소수점 레이아웃, 요소 간 관계)을 무시합니다. 반대로, 기존의 텐서 특화 압축기들은 고정된 포맷별 파이프라인이나 사전 정의된 변환에 의존하는 경우가 많으며, 이는 반복되는 값, 재발하는 부분 시퀀스, 또는 요소들 사이의 단순한 관계와 같은 다양한 텐서 특유의 패턴을 포착하는 데 실패할 수 있습니다.

핵심 과제는 고정된, 사전 정의된 방식에 의존하지 않고 각 텐서의 고유한 구조에 적응하는 무손실(bit-exact) 압축을 달성하는 것입니다.

방법론: Brevis

저자들은 Brevis를 통해 무손실 텐서 압축을 프로그램 합성(program synthesis) 문제로 정식화하여 제시합니다. Brevis는 고정된 코덱을 선택하는 대신, 각 텐서에 대해 압축된 독립형 프로그램을 합성하며, 이 프로그램이 실행되면 원래의 텐서를 비트 단위로 정확하게 재구성합니다.

1. 타입 지정 도메인 특화 언어 (DSL)

Brevis는 반복되는 텐서 구조를 포착하기 위해 가역적 연산자를 갖춘 타입 지정 DSL을 사용합니다:

• Lit: 다양한 물리적 인코딩(raw words, bit packing, Huffman, rANS)을 사용하여 임의의 워드 스트림을 저장함으로써 가장 작은 표현을 선택하는 범용 폴백(fallback)입니다.
• Const: 단일 워드와 그 반복 횟수를 저장합니다.
• Concat: 인접한 영역에 대한 프로그램을 결합합니다.
• Repeat: 비어 있지 않은 자식 프로그램을 여러 번 확장하여 반복되는 부분 시퀀스를 포착합니다.
• Map: 폭을 보존하는 단사 함수(예: XOR, 모듈로 덧셈, ZigZag, Gray coding)를 적용하여 단순한 문자열 형태로는 보이지 않는 관계를 드러냅니다.
• Scan: 초기 워드와 그 이후의 업데이트(XOR 또는 모듈로 덧셈)를 저장하여 인접한 상관관계를 드러냅니다.
• Merge: 부모의 폭과 합계가 같고 길이가 동일한 두 자식을 결합하며, 이는 부동 소수점 필드(부호, 지수, 가수)나 비트 평면(bit planes)을 분할하는 데 유용합니다.

2. 제한된 A* 탐색을 통한 목표 지향적 합성

합성 과정은 $Exec(P) = Bits(X)$를 만족하며 직렬화된 크기$L(P)$를 최소화하는 프로그램$P$를 찾습니다.

• 역방향 탐색 (Backward Search): 프로그램을 나열하는 대신, Brevis는 타겟 텐서로부터 역방향으로 탐색합니다. 부분 프로그램의 모든 "구멍(hole)"은 반드시 생성해야 하는 정확한 스트림과 쌍을 이룹니다. 연산자는 대상 스트림의 분해(decomposition)가 유효할 때만 적용됩니다 (즉, 자식들의 구성이 타겟을 정확히 재구성해야 함).
• 체크포인트별 생성 우선순위 (Checkpoint-Specific Production Prior): 탐색을 효율적으로 유도하기 위해, Brevis는 체크포인트 내의 텐서 샘플 일부로부터 DSL 생성에 대한 확률 분포를 학습합니다. 이 우선순위는 모델의 특정 구조적 경향(예: 특정 레이어에서 Repeat 또는 Merge를 빈번하게 사용하는 경우)을 포착합니다.
• 제한된 A 알고리즘: 탐색은 $g(s) + h(s)$에 의해 우선순위가 결정되는 제한된 A 알고리즘을 사용합니다.
  • $g(s)$는 학습된 우선순위에 기반한 누적 비용입니다.
  • $h(s)$는 남은 완료 비용에 대한 허용 가능한 휴리스틱(admissible heuristic)입니다.
  • 바이트 하한선($LB_L$)은 현재의 최적해를 개선할 수 없는 상태를 제거(pruning)합니다.
  • 탐색은 확장 예산(예: 메인 평가 시 텐서당 1회 확장)과 노드/깊이 제한을 가집니다.

3. 아카이브 포맷 및 재구성

결과물인 아카이브는 원래의 safetensors 헤더와 레이아웃을 보존합니다. 각 텐서 레코드는 이름, dtype, shape 및 합성된 프로그램을 포함합니다. 역압축은 프로그램을 파싱하고, 타입 규칙을 검증하며, 탐색이나 학습된 우선순위 없이 프로그램을 직접 실행하여 비트 단위로 정확한 재구성을 보장합니다.

주요 기여

1. 정식화: 본 논문은 비트 단위의 체크포인트 압축을 프로그램 합성으로 정식화하여, 각 텐서를 압축된 독립형 DSL 프로그램으로 표현합니다.
2. DSL 설계: 값, 시퀀스, 필드 및 평면 등 다양한 텐서 구조를 포착하면서도 비트 단위의 정확한 재구성을 보장하는 가역적 연산자를 갖춘 타입 지정 도메인 특화 언어를 설계했습니다.
3. 합성 알고리즘: 소량의 샘플로부터 학습된 체크포인트 특화 생성 우선순위에 의해 유도되는 제한된 A* 합성 알고리즘을 통해, 실용적인 예산 내에서 효율적인 압축 프로그램 발견을 가능하게 합니다.
4. 평가: 2.13 TB의 데이터를 포함하는 10개의 공개 체크포인트(언어, 오디오, 이미지)에 대해 종합적인 평가를 수행했습니다.

결과

저자들은 언어, 오디오, 이미지 생성 모델을 아우르는 10개의 공개 체크포인트(2.13 TB의 원본 데이터)에 대해 Brevis를 평가했습니다.

• 저장 공간 절감: Brevis는 2.13 TB의 전체 데이터 크기를 1.41 TB로 줄여, 33.93%의 저장 공간 절감을 달상했습니다.
• 범용 압축기와의 비교: Brevis는 10개의 모든 체크포인트에서 zstd, gzip, LZ4, Snappy보다 더 작은 아카이브를 생성했습니다. 베이스라인 대비 평균 절감률은 zstd 대비 12.94%에서 Snappy 대비 30.87% 사이였습니다.
• 텐서 특화 압축기와의 비교:
  • ZipNN과 비교했을 때, Brevis는 10개의 모든 체크포인트에서 더 작은 아카이브를 생성하여 총 10.21 GB를 추가로 절감했습니다. 통합 절감률은 **0.72%**입니다 (체크포인트당 평균 절감률은 0.71%).
  • DFloat11(검증된 네이티브 결과가 존재하는 Llama-3.1-8B 체크포인트 기준)과 비교했을 때, Brevis는 2.90% 더 작았습니다 (316.27 MB 절감).
• 처리량 (Throughput): 실용적인 병렬 처리 구성 하에서, Brevers는 3.60 GB/s의 압축 속도와 6.61 GB/s의 역압축 속도를 달성했습니다. 이는 더 나은 저장 공간 절감을 제공하면서도 경쟁력 있는 처리량을 유지하는 파레토 최적 경계(Pareto frontier)에 위치합니다.
• 통계적 유의성: 모든 쌍을 이룬 베이스라인과의 개선 사항은 Holm 보정 후에도 통계적으로 유의미했습니다 ($p_H = 0.0098$).

의의 및 주장

본 논문은 압축을 프로그램 합성으로 취급함으로써 Brevis가 고정된 파이프라인의 한계를 극복한다고 주장합니다. 이는 자기 완결적인 프로그램이 범용 코덱이 놓치는 특정 구조적 중복성에 적응하여, 매우 효율적인 비트 단위의 텐서 표현이 될 수 있음을 보여줍니다.

저자들은 다음과 같은 점을 강조합니다:

• 모든 원본 바이트를 보존합니다 (비트 단위 정확성). 이는 양자화나 가지치기(pruning)와 같은 손실 방식과 달리, 아카이브 및 정확한 체크포인트 전송에 적합합니다.
• 하드웨어 특화 디코더나 사전 정의된 코덱 패밀리를 요구하지 않습니다.
• 다양한 모델 유형 및 데이터 형식(BF16, FP32, FP8)에 대해 상당한 저장 공간 절감을 달성합니다.

저자들은 현재 구현이 텐서를 독립적으로 합성한다는 점(교차 텐서 합성은 지원하지 않음)과 평가 범위가 모든 텐서 워크로드가 아닌 공개 체크포인트에 국한된다는 점을 언급하며 한계를 겸허히 인정합니다. 저자들은 Brevis를 성장하는 대규모 모델 생태계를 위한 더 적응적이고 구조를 인식하는 무손실 압축을 향한 단계로 포지셔닝합니다.