가입 후 초대 링크를 공유하면 동영상 재생 및 초대 보상을 받을 수 있습니다.

Larry Dial
@classiclarryd
Technical Staff at Open Athena, working on Marin
가입 May 2024
48 팔로잉 중    2.2K
Question for the LLM Research Community: Is anyone aware of fully reproducible experimental results showing that MLA beats GQA under the same KV cache? I ask for two reasons: (1) I find it surprising that there is still a divide between Chinese labs using MLA, and Western labs using GQA + sliding window, when certainly many ablations have been run by many parties. (2) At Marin we are considering MLA for our next large scale run, but in early (!) smaller scale ablations it appears worse than heavily tuned feature-rich GQA, even after controlling for KV cache. I would like to run better experiments here. Below I cover my thoughts on general reproducibility, then specifics on MLA. Every empirical result in ML is only contextually true. Conditioned on the data distribution, optimizer settings, model width, model depth, finer architecture details, hardware, kernel engineering, initialization, token count, tokenizer, context length, and evaluation protocol, one can reach different conclusions. Contextual results are still useful. Typically if I see a promising method, I will first attempt a full 'context jump', where I apply it to my own context, hoping results transfer. Sometimes they do. If they don't, I can try 2 things: modify the implementation of the method, or modify the context. Ideally I have access to the full context of the original result. Then I can perform a 'context bridge', where I ablate one aspect of the context at a time, isolating exactly why a method performs differently. This lets me make an informed decision to either update my context to let the method shine, or stick with my context and leave the method out. MLA is tricky to assess at small scale. A core aspect of MLA is compressing hidden_dim->latent_dim. Then for each head, latent_dim->head_dim. Typically head_dim is fixed at 128, partially for hardware reasons, and partially for learning dynamics (head_dim of 8 wouldn't have sufficient representational capacity). To get MLA dynamics, you want hidden_dim>>latent_dim, and latent_dim>head_dim. This window closes at small scale. The degree of tuning can unfairly alter the scales. In GQA we have partial RoPE, QK Norm, Gated Attention, attention sharpening, sliding window, and other techniques that give a 30%+ training boost. They don't seem to give the same boost to MLA. On one hand, you want to compare techniques apples:apples with equal tuning. On the other hand, there is a finite amount of future tuning you can do, so prior tuning influences which approach is most pragmatic. Creating controlled tests between MLA and GQA is tricky. Several factors: kv_cache, quadratic attention flops, attention projection flops. kv_cache is controlled by scaling down kv_heads to match MLA, or scaling up kv_latent to match kv_heads. quadratic attention flops are controlled by scaling up GQA's query head count to match MLA head count, or scaling down MLA head count. Also scaling up GQA head_dim 128->192, or scaling down MLA head_dim to 192->128. In general, it's informative to context match to both option A's preferred context and option B's preferred context. Sliding window is another confounder. MLA is theoretically elegant, if we ignore RoPE. It replaces the 'replicate' op of kv_heads in GQA with a 'mix' op (pic below). Since the 'mix' can learn to 'replicate' if it wants, MLA is purely more expressive, and the cost of 'mix' is hidden at inference with absorb trick. Yet in practice, I find that at small scale this 'mix' op doesn't add much value and interacts poorly with the optimizer dynamics. And the change to RoPE hurts. My current plan is to first tune and ablate our model features around MLA, then run 3 scaling ladders: MLA, GQA with 2 kv_heads, and GQA with higher kv_heads. For each ladder, fit a loss vs compute projection. If MLA performs worse at our target compute compared to both GQA options, drop it. If MLA beats 2 kv_heads but loses to higher KV_heads, then it becomes a kv_cache tradeoff. Early results indicate MLA will perform worse than both feature-rich GQA ladders, but we will see. Any positive external reproducible results for MLA would help make sure I give it the best chance possible.
더 보기