🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

원저자: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

게시일 2026-05-14

📖 4 분 읽기☕ 가벼운 읽기

원저자: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대하고 초지능적인 로봇(대형 언어 모델)에게 글쓰기, 대화, 추론을 가르치려 한다고 상상해 보세요. 이를 위해 당신은 방대한 도서관(학습 데이터)을 로봇에게 보여줍니다. 로봇은 실수를 하고, 작업을 점검하며, 더 나아지기 위해 내부의 '뇌 설정'(파라미터)을 조정함으로써 학습합니다.

이러한 조정을 수행하는 표준적인 방법은 **아담(Adam)**이라는 매우 꼼꼼한 회계사와 같습니다. 아담은 모든 뇌 설정에 대해 상세한 장부를 유지하며, 현재 실수뿐만 아니라 과거 실수의 기록까지 추적하여 무엇을 얼마나 변경할지 결정합니다.

문제: 회계사가 너무 무겁습니다
문제는 이 '회계사'(최적화기)가 극도로 무겁다는 점입니다. 수십억 개의 뇌 설정을 가진 로봇의 경우, 회계사의 장부는 로봇의 뇌 자체보다 두 배 더 많은 메모리를 차지합니다.

비유: 집을 옮기려 한다고 상상해 보세요. 가구(모델)는 크지만, 이동 트럭(최적화기 상태)은 그보다 더 큽니다. 만약 작은 트럭만 있다면(제한된 컴퓨터 메모리), 집이 아무리 훌륭하더라도 아예 집을 옮길 수 없습니다. 이것이 더 크고 똑똑한 로봇을 학습시키는 연구자들을 막는 '메모리 병목 현상'입니다.

구솔루션: 모서리를 잘라내기
이전에는 이동 트럭을 줄이려는 시도가 있었지만 결함이 있었습니다:

집의 일부 동결: 일부 가구는 움직이지 않고 작은 상자 몇 개만 조정합니다. 이는 공간을 절약하지만 집을 덜 유연하게 만듭니다(성능 저하).
흐릿한 사진 촬영: 실제 가구를 옮기는 대신 공간을 절약하기 위해 저해상도 사진을 찍습니다. 하지만 이러한 사진을 계산하는 데는 많은 시간과 에너지(계산 오버헤드)가 소모됩니다.
장부 공유: 서로 다른 방이 같은 노트를 공유하게 합니다. 이는 공간을 절약하지만 조정이 덜 정밀해집니다.

새로운 솔루션: FOAM(스마트 접기 방식)
이 논문은 FOAM(Folded Optimizer with Approximate Moment)을 소개합니다. FOAM은 중요한 세부 사항을 잃지 않고 거대한 이동 트럭을 작은 밴에 넣을 수 있는 영리한 '접기' 기술을 사용하는 마스터 패커라고 생각하세요.

다음은 FOAM이 작동하는 방식을 간단한 비유로 설명한 것입니다:

1. "블록 접기"(압축)

FOAM은 모든 뇌 설정을 개별적으로 추적하는 대신, 이를 작은 블록(8 개 또는 16 개 설정의 그리드 등)으로 그룹화합니다.

비유: 100 명의 사람들이 줄지어 서 있고 그들의 평균 키를 기억해야 한다고 가정해 보세요. 100 개의 숫자를 따로 적는 대신, 10 명씩 10 개 팀으로 그룹화합니다. 그리고 단 하나의 숫자만 적습니다: 각 팀의 평균 키입니다.
결과: 이로 인해 '장부'에 필요한 메모리가 엄청난 양(최대 90% 감소)으로 줄어듭니다.

2. "잔류 보정"(안전망)

평균만 사용하면 각 사람의 고유한 세부 사항을 잃게 됩니다. 팀 안에 매우 키가 큰 사람이 있고 또 다른 이는 매우 작은 사람일 수 있습니다. 평균은 이를 숨깁니다.

비유: FOAM은 똑똑합니다. '팀 평균'을 적은 후, 실제 사람들과 평균 사이의 차이(오차)를 빠르게 계산합니다. 이를 '잔류값(Residual)'이라고 부릅니다.
트릭: FOAM은 이 차이를 영원히 보관하지 않습니다. 계산하고, 그 순간의 업데이트를 수정하는 데 사용한 후 버립니다. 마치 요리사가 수프를 맛보고 소금을 조절했다가, 한 숟가락의 정확한 맛을 기록해 두는 대신 잊어버리는 것과 같습니다.
중요성: 이는 메모리가 압축되었음에도 로봇이 모든 뇌 설정의 고유한 세부 사항을 여전히 학습할 수 있도록 보장합니다.

3. "펼치기"(복원)

로봇의 뇌를 실제로 업데이트할 때가 되면 FOAM은 압축된 '팀 평균'을 가져와 전체 크기로 다시 확장하고, 그 위에 '잔류값' 보정을 추가합니다.

결과: 로봇은 무거운 회계사를 사용할 때와 마찬가지로 정밀한 업데이트를 받지만, 컴퓨터는 작은 접힌 버전만 운반하면 됩니다.

논문이 주장하는 것 (결과)
저자들은 6 천만 파라미터 모델부터 70 억 파라미터 모델까지 다양한 로봇 모델을 학습시켜 FOAM 을 테스트했습니다. 그들이 발견한 바는 다음과 같습니다:

막대한 메모리 절약: FOAM 은 최적화기에 필요한 메모리를 최대 **90%**까지 줄입니다. 전체적으로 학습에 필요한 메모리 사용량을 약 50% 감소시킵니다. 이는 동일한 하드웨어에서 더 큰 모델을 학습하거나, 동일한 모델을 훨씬 빠르게 학습할 수 있음을 의미합니다.
성능 손실 없음: 무거운 압축에도 불구하고 FOAM 으로 학습된 모델은 표준적인 무거운 회계사로 학습된 모델만큼 잘(때로는 더 잘) 수행합니다. 더 빠르게 학습하고 더 높은 정확도에 도달합니다.
속도: 무거운 짐을 운반할 필요가 없으므로 학습 과정이 더 빠릅니다.
다용도성: LLaMA, Qwen 등 다양한 유형의 로봇 아키텍처에서 작동하며 다른 메모리 절약 기술과도 결합할 수 있습니다.

한 줄 요약
FOAM 은 마법의 여행가방과 같습니다. 정보를 깔끔하게 접고 잃어버리지 않도록 빠른 '보정 메모'를 추가함으로써, 방대한 양의 정보(최적화기의 메모리)를 작은 공간에 담을 수 있게 해줍니다. 이를 통해 연구자들은 거대하고 비싼 컴퓨터 없이도 더 똑똑하고 큰 AI 모델을 구축할 수 있습니다.

기술 요약: FOAM – 메모리 효율적 LLM 학습을 위한 차단 상태 폴딩

1. 문제 제기

대규모 언어 모델 (LLM) 은 파라미터 수와 학습 데이터의 증가를 통해 탁월한 성능을 달성해 왔으나, 그 규모는 학습 과정에서 상당한 메모리 병목 현상을 초래합니다. 사실상의 최적화기인 Adam(및 AdamW) 은 모든 파라미터에 대해 1 차 및 2 차 모멘트 추정을 저장해야 하므로, 모델 크기의 두 배에 해당하는 최적화기 상태 메모리 오버헤드가 발생합니다. 이로 인해 7B(70 억) 이상의 파라미터를 가진 대규모 모델의 사전 학습 및 미세 조정은 메모리 제약에 묶이게 되며, 작은 배치 크기조차도 수백 기가바이트의 VRAM 을 필요로 하는 경우가 많습니다.

기존의 메모리 효율적 솔루션들은 다음과 같은 트레이드오프에 직면해 있습니다:

파라미터 고정 (예: LoRA): 저랭크 부분 공간으로 업데이트를 제한하여 모델의 표현력과 성능을 잠재적으로 제한할 수 있습니다.
저랭크 투영 (예: GaLore, APOLLO): 특이값 분해 (SVD) 또는 무작위 투영을 사용하여 최적화기 상태를 압축합니다. 효과적이기는 하나, SVD 는 상당한 계산 오버헤드를 초래하며, 무작위 투영은 엄격한 메모리 제약 하에서 성능 저하를 유발할 수 있습니다.
블록 공유 (예: Adam-Mini): 2 차 모멘트 추정을 블록 간에 공유하여 저장 공간을 줄이지만, 1 차 모멘트 오버헤드는 여전히 높거나 메모리 감소가 약 50% 로 제한되는 경우가 많습니다.

2. 방법론: FOAM

저자들은 최적화 충실도를 희생하지 않고 최적화기 상태를 압축하도록 설계된 구조화된 상태 폴딩 프레임워크인 FOAM(Folded Optimizer with Approximate Moment) 을 제안합니다. FOAM 은 차단 평균 (blocked averaging) 과 잔차 보정 (residual correction) 을 결합한 원리에 따라 작동합니다.

핵심 메커니즘

시간 단계 $t$ 에서의 기울기 행렬 $G_t \in \mathbb{R}^{m \times n}$ 과 폴딩 레벨 $l$ 이 주어졌을 때:

차단 평균 (Fold): 기울기는 $2^l$ 개의 연속된 요소로 구성된 블록으로 분할됩니다. 폴딩 연산자 $A^{(l)}$ 은 각 블록의 평균을 계산하여 기울기 차원을 $n$ 에서 $n/2^l$ 로 압축합니다. 1 차 및 2 차 모멘트 ( $\tilde{M}_t, \tilde{V}_t$ ) 는 이러한 압축된 기울기를 기반으로 업데이트됩니다.
언폴딩: 전체 파라미터 공간에 업데이트를 적용하기 위해, 폴딩 연산자 $E^{(l)}$ 이 압축된 엔트리를 원래 차원으로 복제합니다.
잔차 보정: 업데이트를 직접 복제하면 블록 내 모든 요소가 동일한 값을 공유하게 되어, 차별화는 초기화에만 의존하게 됩니다. 손실된 정보를 복구하기 위해 FOAM 은 원래 기울기와 폴딩 - 언폴딩된 버전 간의 오차를 나타내는 잔차 항 $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ 을 계산합니다.
- 1 차 모멘트는 다음과 같이 재구성됩니다: $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- 2 차 모멘트는 다음과 같이 재구성됩니다: $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- 중요하게도, $R_t$ 는 온더플라이 (on-the-fly) 로 계산되어 즉시 해제되며 메모리에 저장되지 않습니다.

하이브리드 적용

FOAM 은 MLP 및 어텐션 모듈에 구체적으로 적용되는 반면, 임베딩 레이어에는 표준 Adam 이 사용됩니다. 메모리 효율적 최적화기에서 일반적인 관행인 스케일링 계수 $\alpha$ 가 FOAM 과 Adam 모듈 간의 학습률 비율을 조정합니다.

3. 주요 기여

혁신적인 최적화기 설계: FOAM 은 Full-Adam 대비 최적화기 상태 메모리 오버헤드를 최대 90% 감소시키고 총 엔드투엔드 메모리 풋프린트를 약 50% 감소시킵니다. 투영 기반 방법과 달리 투영 행렬이 필요 없으며 전체 파라미터 최적화를 지원합니다.
이론적 보장: 표준 비볼록 최적화 가정 (L-부드러움, 유계 분산, 유계 기울기) 하에서, 저자들은 FOAM 이 바닐라 Adam 과 동일한 $O(\frac{\log T}{\sqrt{T}})$ 의 수렴 속도를 달성함을 증명합니다. 잔차 오차 비율 $\delta_l$ 은 1 로 제한되며 주요 수렴 속도에 부정적인 영향을 미치지 않습니다.
광범위한 평가: 이 방법은 모델 크기 (60M~7B), 아키텍처 (LLaMA, Qwen, GPT, DeBERTa), 작업 (C4 에 대한 사전 학습, GLUE 및 MMLU 에 대한 미세 조정) 전반에 걸쳐 검증되었습니다.

4. 실험 결과

본 논문은 다음과 같은 실증적 발견을 보고합니다:

사전 학습 성능: FOAM 은 메모리 효율적 베이스라인 (GaLore, APOLLO, Adam-Mini, Muon) 보다 검증 퍼플렉시티 (PPL) 측면에서 일관되게 우수한 성능을 보이며, 동시에 훨씬 낮은 메모리 사용량을 유지합니다. 예를 들어, LLaMA-1.3B 에서 FOAM-3 은 Adam 대비 총 메모리 소비를 약 50% 줄이면서 더 낮은 PPL 을 달성합니다.
수렴 속도: FOAM 은 다른 메모리 효율적 방법들보다 목표 PPL 에 도달하는 데 종종 더 적은 학습 반복 횟수가 필요하여 더 빠른 수렴을 보여줍니다.
확장성: FOAM 은 더 큰 모델 (LLaMA-3B, 7B) 과 긴 시퀀스 학습 환경 (최대 2048 토큰) 에서도 우수한 성능을 유지합니다.
미세 조정: GLUE 및 MMLU 벤치마크에서 FOAM 은 Full-Adam 및 기타 베이스라인과 비교할 수 있거나 더 나은 성능을 달성하여 사전 학습을 넘어선 유용성을 입증합니다.
강건성: 제거 실험 (Ablation studies) 은 잔차 항이 성능에 결정적임을 확인시켜 주며, 이를 제거하면 수렴이 느려지고 최종 PPL 이 높아집니다. FOAM 은 8 비트 양자화와 결합되거나 Adam-Mini 및 Muon 과 같은 다른 최적화기와 통합될 때도 강건성을 유지합니다.

5. 중요성 및 주장

본 논문은 FOAM 이 대규모 언어 모델 학습의 메모리 병목 현상에 대한 실용적이고 매우 효율적이며 최적화기 독립적인 솔루션을 제공한다고 주장합니다. 가중치 고정이나 비싼 투영 계산의 필요성을 없애함으로써, FOAM 은 메모리 오버헤드를 극적으로 줄이면서 전체 파라미터 최적화를 가능하게 합니다. 저자들은 FOAM 을 저랭크 투영 기술에 대한 보완적 대안으로 위치시키며, 이론적으로는 Adam 의 수렴 속도와 일치하면서도 경험적으로는 우수한 메모리 효율성과 종종 더 나은 학습 성능을 제공하는 방법을 제시합니다.

이 연구는 FOAM 이 기존 하드웨어에서 더 큰 모델의 학습을 가능하게 하거나 현재 최첨단 모델의 학습 비용을 줄여 고성능 LLM 학습을 더 접근 가능하게 만들 수 있음을 시사합니다.