MIT의 Han Lab은 단 256KB의 메모리만으로도 학습이 가능한 프레임워크인 Tiny Training을 공개했습니다.
이 글에서는 해당 기술의 핵심 개념과 구조를 스토리지나 IoT 디바이스 등 제한된 환경에서 어떻게 응용할 수 있는지까지 살펴봅니다.
왜 학습이 어려운가?
대부분의 AI 학습은 수백 MB ~ 수 GB에 달하는 메모리와 GPU 연산을 요구합니다. 특히 학습 중에는:
- 모든 activation을 저장해야 하고
- 모든 gradient를 계산해야 하며
- 모든 weight를 update할 수 있어야 합니다.
그러나 IoT 디바이스나 MCU(Microcontroller Unit)에는 SRAM이 256KB, Flash는 1MB 정도로 극히 제한적입니다. 이 환경에선 기존 방식으로 학습이 불가능합니다.
Tiny Training의 핵심 기술 요약
1. Quantization-Aware Scaling (QAS)
- 8-bit quantized gradient는 불안정하지만, layer마다 scale factor를 조정하면 학습 안정성이 크게 향상됩니다.
- 덕분에 32bit float 없이도 학습 가능, 메모리 대폭 절감.
2. Sparse Update (기여도 기반 선택적 업데이트)
- 모든 weight를 학습하지 않고, 기여도가 높은 일부 block만 학습합니다.
- 기여도(importance) = |activation × gradient|
- 오프라인에서 Δaccuracy를 기준으로 Top-k block만 선택, 그 외는 update 생략 → 20× 이상 메모리/연산 절감.
3. Compile-Time AutoDiff + Tiny Training Engine (TTE)
- PyTorch처럼 동적으로 연산하지 않고, 정적(Static) 그래프를 컴파일 타임에 생성합니다.
- 불필요한 연산/gradient/activation은 그래프에서 제거(prune)하고, 최적의 순서로 재배열 → 메모리 효율 최상화
- TTE는 MCU나 SSD에 올릴 수 있는 초경량 학습 엔진
Top-k 방식, 그런데 전체 gradient를 계산해야 하지 않나?
좋은 의문입니다. 그러나 Tiny Training은 다음과 같이 해결합니다:
- Top-k는 오프라인에서 한 번만 탐색합니다.
- 이때는 GPU 등 자원이 충분한 환경에서 전체 기여도를 계산해보고, 업데이트할 weight block을 미리 정합니다.
- 이 정보를 기반으로 컴파일 시점에 backward 그래프를 잘라냅니다.
- 즉, 런타임에서는 계산할 필요도, 저장할 필요도 없는 부분은 아예 존재하지 않습니다.
결론: 실제 디바이스에서는 "이미 정해진 block만 학습"하게 되어 있고, 전체 gradient를 매번 계산하지 않습니다.
왜 작은 모델에도 학습이 필요할까?
Tiny LLM이 이미 학습된 모델이라도, 현장에서 다음과 같은 상황이 발생합니다:
- 사용자 음성/행동/센서 패턴에 맞춰 개인화
- 환경 적응 (예: 특정 공장 상황에 맞춘 센서 calibration)
- 클라우드 연결 불가/지연/보안 이슈
따라서 경량 모델도 디바이스 내에서 미세조정 학습(fine-tuning) 이 필요합니다.
Tiny Training은 이 목적에 최적화된 툴입니다.
정리 및 활용 가능성
| 기술 | 효과 |
| QAS | 8-bit 양자화로도 안정적 학습 가능 |
| Sparse Update | 불필요한 gradient 생략 → 메모리/연산 절감 |
| Compile-Time AutoDiff | 저장과 계산이 필요한 부분만 남긴 최적의 학습 그래프 생성 |
| TTE | MCU, SSD 등 초저자원 환경에서도 학습 실행 가능 |
📌 적용 아이디어: Storage Controller, IoT Edge Device, Wearable AI, 산업용 센서, Automotive 시스템 등
마무리
MIT의 Tiny Training은 극한의 환경에서도 학습 가능한 방법을 보여줬습니다.
리소스가 제한된 MCU 기반 Edge Device 에서도 AI 모델을 효율적으로 학습하고 활용할 수 있는 가능성을 제시한 사례입니다.
