无需训练,动态优化大模型注意力头选择以提升短上下文效率
BOSCH: Black-Box Binary Optimization for Short-Context Attention-Head Selection in LLMs
- 通过黑箱探测识别层重要性,自适应分配滑动窗口比例
- 在4个模型上比现有方法性能更优,尤其在高比例时优势明显
- 适合需要高效推理的长文本处理场景,如持续预训练
后训练阶段将大语言模型(LLMs)的二次复杂度自注意力替换为滑动窗口注意力(SWA),可降低KV缓存占用并改善延迟。现有混合方案通常在层级别(如交替)或头部级别采用静态排名,但层级方案忽略了同一层内头对局部与全局依赖的路由差异,而静态头部排名存在耦合问题:混合后头的行为可能改变。本文提出BOSCH——一种无需训练的黑箱二值优化方法,将问题建模为大规模邻域搜索,并分解为三个子问题:(i) 用小预算黑箱探针检测层重要性;(ii) 基于敏感度自适应分配每层的SWA比例;(iii) 在比例桶内进行分组头部优化。在4个参数量从1.7B到30B的模型上,跨4种SWA比例的实验表明,BOSCH持续优于层级启发式和6种强静态头部方法,且在更高比例下增益更大。在持续预训练中,BOSCH能更快、更全面恢复原始长上下文性能。对所选头部的分析显示,不同SWA比例下头部选择有显著变化,证明应针对目标比例动态进行头部筛选,而非依赖固定局部性排名。
原文摘要 · Abstract (English)
Post-training hybridization of large language models (LLMs) often replaces quadratic self-attention with sliding-window attention (SWA) to reduce KV cache usage and improve latency. Existing hybridization schemes are typically defined either at the layer level (e.g., interleaving) or at the head level via static rankings from local to global. Layer-level schemes ignore that local and global dependencies are routed through heads within the same layer, while static head-level rankings suffer from entanglement: a head's local/global behavior can change after hybridization. We propose BOSCH, Black-box Binary Optimization for Short-context Head Selection, a training-free method that formulates the problem as a Large Neighborhood Search and decomposes it into three subproblems: (i) layer-importance detection via small-budget black-box probes, (ii) adaptive per-layer SWA-ratio assignment based on these sensitivities, and (iii) grouped head-level optimization within ratio buckets. Extensive experiments on 4 LLMs ranging from 1.7B to 30B parameters, across 4 SWA ratios, show that BOSCH consistently outperforms layer-level heuristics and 6 strong static head-level methods, with larger gains at higher SWA ratios. Under continual pretraining, BOSCH recover original long-context performance faster and to a higher level. Analysis of the selected heads reveals substantial turnover for BOSCH across different SWA ratios, underscoring the importance of performing head-level selection for each target ratio rather than relying on fixed locality rankings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。