본문으로 건너뛰기

[llm-compressor] Logarithmic Equalization: 로그 스케일 채널 균등화

들어가며

SmoothQuant는 활성화 outlier를 가중치로 이동시킨다. **Logarithmic Equalization(LogEq)**은 비슷한 철학이지만 약간 다른 공식을 쓴다. 채널 간 가중치 스케일 비율을 로그 공간에서 균등화해 "매우 작은 채널과 매우 큰 채널"의 차이를 줄인다. 이는 DFQ(Data-Free Quantization) 계열 논문에서 처음 제안된 기법으로, 활성화 통계가 필요 없다는 장점이 있다. llm-compressor의 src/llmcompressor/modifiers/logarithmic_equalization/base.py를 분석한다.

핵심 구조/코드 분석

LogEqualizationModifier 생성자 파라미터

class LogEqualizationModifier(Modifier):
    """
    Equalize per-channel weight scales in log-space between consecutive
    Linear layers to improve quantization friendliness.
    """
    mappings: list[LogEqMapping] | None = None   # smooth-balance 매핑 (SmoothQuant 와 유사)
    ignore: list[str] = field(default_factory=list)

이 Modifier는 SmoothQuant보다 파라미터가 적다. smoothing_strength 같은 $\alpha$가 없다. LogEq는 항상 "기하 평균 균등화"를 적용하기 때문이다.

수학적 아이디어

두 개의 연속된 Linear 레이어 $y = W_2 W_1 x$가 있다고 하자. 중간 축(즉 $W_1$의 출력 채널 = $W_2$의 입력 채널)에 대각 스케일 $S$를 도입하면

$$ y = W_2 W_1 x = W_2 S S^{-1} W_1 x = (W_2 S) (S^{-1} W_1) x = \hat{W_2} \hat{W_1} x $$

수학적 등가성이 유지된다. LogEq는 $S$를 각 채널 $c$의 $W_1$ 출력 채널 범위와 $W_2$ 입력 채널 범위의 기하 평균으로 설정한다.

$$ s_c = \sqrt{\frac{r_1(c)}{r_2(c)}} $$

여기서 $r_1(c) = \max(|W_1[c, :]|)$, $r_2(c) = \max(|W_2[:, c]|)$. 이 선택은 변환 후 두 행렬의 max-abs가 로그 공간에서 균등해지게 만든다.

on_initialize: 매핑 기반 변환 적용

def on_initialize(self, state: State, **kwargs) -> bool:
    self.resolved_mappings = self._resolve_mappings(state.model)

    for mapping in self.resolved_mappings:
        W1 = state.model.get_submodule(mapping.prev_layer).weight.data
        W2_list = [
            state.model.get_submodule(name).weight.data
            for name in mapping.next_layers
        ]

        # W1 의 출력 채널별 최대 절댓값
        r1 = W1.abs().amax(dim=1)                          # shape (out_features_of_W1,)

        # W2 들의 입력 채널별 최대 절댓값 (여러 W2 중 최대)
        r2_stack = torch.stack([w.abs().amax(dim=0) for w in W2_list])
        r2 = r2_stack.amax(dim=0)                          # shape (in_features_of_W2,)

        # 기하 평균 기반 스케일
        s = torch.sqrt(r1 / r2.clamp(min=1e-8))

        # 변환 적용
        self._apply_equalization(mapping, s, state)

    return True


def _apply_equalization(self, mapping, s, state):
    # W1 의 출력 채널을 1/s 로 축소
    prev_module = state.model.get_submodule(mapping.prev_layer)
    prev_module.weight.data /= s.unsqueeze(-1)
    if prev_module.bias is not None:
        prev_module.bias.data /= s

    # W2 의 입력 채널을 s 로 확대
    for name in mapping.next_layers:
        module = state.model.get_submodule(name)
        module.weight.data *= s.unsqueeze(0)

핵심은 s = sqrt(r1 / r2) 한 줄이다. 이 값으로 $W_1$의 출력 채널을 축소하고 $W_2$의 입력 채널을 확대하면, 두 행렬의 max-abs가 로그 공간에서 같아진다. 변환 전후 모델 출력은 동일하다.

LogEqMapping: 연속 Linear 쌍 정의

@dataclass
class LogEqMapping:
    prev_layer: str                # 이전 Linear 또는 RMSNorm
    next_layers: list[str]         # 다음 Linear 들 (공유 입력 가능)

SmoothQuant의 SmoothMapping과 구조가 같다. 차이는 이름뿐이다. 실제로 LogEq의 매핑 후보는 "RMSNorm → projection" 같은 쌍보다 "fc1 → fc2" 같은 Linear-Linear 쌍에 더 적합하다. 연속 Linear 사이에는 activation function(GELU 등)이 있지만, 스케일 변환은 element-wise이므로 activation을 통과해도 유효하다.

실제 코드에서는 다음과 같은 매핑이 많이 쓰인다.

# MLP 의 gate_proj/up_proj → down_proj
LogEqMapping(
    prev_layer="mlp.gate_proj",      # 또는 up_proj
    next_layers=["mlp.down_proj"],
)

왜 이 설계인가

1. Data-free 작동. LogEq는 활성화 통계를 보지 않는다. 가중치만 보고 계산한다. 이는 캘리브레이션 데이터가 없거나 부적절한 시나리오에서 유용하다. 빠르고 단순하다.

2. 기하 평균의 수학적 아름다움. sqrt(r1/r2)는 "로그 공간에서 두 범위의 중점"이다. 변환 후 두 행렬의 max-abs 범위가 동일해지므로 "채널별 양자화 해상도 편차"가 최소화된다.

3. 별도 파라미터 없음. SmoothQuant의 $\alpha$ 같은 튜닝 대상이 없다. 사용자는 그냥 적용하면 된다. 이는 "알고리즘이 자기 결정을 완전히 담당"한다는 철학이다.

4. 비파괴적 변환. W_1 ← W_1/s, W_2 ← W_2*s 두 변환이 수학적 등가성을 유지한다. 양자화 전까지는 모델 출력이 정확히 보존된다.

5. 연속 Linear 쌍에 집중. LogEq는 Linear-Linear 쌍에서 가장 효과적이다. RMSNorm-projection 쌍은 SmoothQuant가 더 낫다. 두 Modifier를 조합해 쓰는 것이 가능하다.

마무리

Logarithmic Equalization은 DFQ 시절의 고전 기법을 LLM에 적용한 Modifier다. 데이터가 없을 때 SmoothQuant 대체재로 유용하다. 이로써 Quantization Modifier 섹션이 끝났다. 다음 글부터는 Pruning 섹션이다.

참고 자료

댓글

관련 포스트

llm-compressor 의 다른글