Controlling changes to attention logits

• URL: http://arxiv.org/abs/2511.21377v1
• Date: Wed, 26 Nov 2025 13:24:35 GMT
• Title: Controlling changes to attention logits
• Abstract summary: Stability of neural network weights is critical when training transformer models.<n>Applying normalization to queries and keys, known as QK norm', fixes stability issues in practice, but is not always applicable.<n>We show that our cheap intervention allows us to increase the base learning rate of the network, outperform other methods in the MLA setting, and achieve performance competitive with QK norm.
• Score: 21.620742229590885
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Stability of neural network weights is critical when training transformer models. The query and key weights are particularly problematic, as they tend to grow large without any intervention. Applying normalization to queries and keys, known as `QK norm', fixes stability issues in practice, but is not always applicable. For example, QK norm is not compatible with Multi Latent Attention (MLA) because QK norm requires full materialization of queries and keys during inference, which is not done in MLA. In this paper we suggest that controlling the changes to logits is important for stability. We show that these changes are controllable by assigning parameter-dependent learning rates to the query and key weights. We find that our cheap intervention allows us to increase the base learning rate of the network, outperform other methods in the MLA setting, and achieve performance competitive with QK norm when using Multi-head Attention.

Related papers

• QK-Normed MLA: QK normalization without full key caching [65.51862319225098]
  We show that QK normalization is a practical stabilization option for MLA models without requiring full-key caching.<n>In our 400M runs trained for up to 100B tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping.
  arXiv  Detail & Related papers  (2026-06-15T07:16:30Z)
• JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates [9.03449876943648]
  Quantization-aware training (QAT) is widely deployed but typically relies on the Straight-Through Estimator (STE), which passes through non-differentiable quantizers by fiat.<n>This often makes training aligned near bin and weakly with actual low-precision model.<n>We introduce JacQuant, a framework that learns a lightweight surrogate model's local sensitivity to parameter changes.
  arXiv  Detail & Related papers  (2026-05-25T06:19:49Z)
• Trust the Batch, On- or Off-Policy: Adaptive Policy Optimization for RL Post-Training [50.86545293331458]
  Reinforcement learning is structurally harder than supervised learning because the policy changes the data distribution it learns from.<n>We present a simple yet effective batch-adaptive objective that replaces fixed clipping with the normalized effective sample size of the policy ratios.
  arXiv  Detail & Related papers  (2026-05-12T16:44:47Z)
• QUEST: A robust attention formulation using query-modulated spherical attention [12.908688634483253]
  We explore the role played by norms of the queries and keys, which can cause training instabilities when they arbitrarily increase.<n>We propose a new attention formulation, QUEry-modulated Spherical aTtention (QUEST), that constrains the keys to a hyperspherical latent space.<n>QUEST can be easily used as a drop-in replacement for standard attention.
  arXiv  Detail & Related papers  (2026-03-31T20:04:17Z)
• ReasonCACHE: Teaching LLMs To Reason Without Weight Updates [75.2707292367514]
  We show that large language models (LLMs) can learn to reason without overloading the context window and without any weight updates.<n>We introduce ReasonCACHE, an instantiation of this mechanism that distills demonstrations into a fixed key-value cache.<n> Empirically, ReasonCACHE outperforms standard ICL and matches or surpasses IWL approaches.
  arXiv  Detail & Related papers  (2026-02-02T17:24:23Z)
• Enhancing Post-Training Quantization via Future Activation Awareness [84.76726857601753]
  Post-training quantization (PTQ) is a widely used method to compress large language models (LLMs) without fine-tuning.<n>We propose Future-Aware Quantization (FAQ), which leverages future-layer activations to guide quantization.<n>FAQ consistently outperforms prior methods with negligible extra cost, requiring no backward passes, data reconstruction, or tuning.
  arXiv  Detail & Related papers  (2026-01-28T12:03:30Z)
• Quantum Machine Learning for Secondary Frequency Control [0.0]
  This paper introduces a novel approach using a pure variational quantum circuit (VQC) for real-time secondary frequency control in diesel generators.<n>The proposed VQC operates independently during execution, eliminating latency from classical-quantum data exchange.<n>The VQC achieves high prediction accuracy (over 90%) with sufficient quantum measurement shots and generalizes well across diverse test events.
  arXiv  Detail & Related papers  (2025-11-29T15:03:52Z)
• GateRA: Token-Aware Modulation for Parameter-Efficient Fine-Tuning [51.79350934271497]
  GateRA is a unified framework that introduces token-aware modulation to dynamically adjust the strength of PEFT updates.<n>By incorporating adaptive gating into standard PEFT branches, GateRA enables selective, token-level adaptation.<n> Experiments on multiple commonsense reasoning benchmarks demonstrate that GateRA consistently outperforms or matches prior PEFT methods.
  arXiv  Detail & Related papers  (2025-11-15T17:55:47Z)
• Causal Attention with Lookahead Keys [52.63961482746826]
  In standard causal attention, each token's query, key, and value (QKV) are static and encode only preceding context.<n>We introduce CAuSal aTtention with Lookahead kEys (CASTLE), an attention mechanism that continually updates each token's keys as the context unfolds.
  arXiv  Detail & Related papers  (2025-09-09T00:15:23Z)
• Quantize What Counts: More for Keys, Less for Values [63.51476878610841]
  Large Language Models (LLMs) suffer inference-time memory bottlenecks dominated by the attention Key-Value ( KV) cache.<n>This paper proposes two theorems that anchor mixed-precision KV quantization in the intrinsic geometry of Transformer models.
  arXiv  Detail & Related papers  (2025-02-20T22:24:27Z)
• Beyond Uniform Query Distribution: Key-Driven Grouped Query Attention [3.3457276841127315]
  Transformer architecture has revolutionized deep learning through its Self-Attention mechanism. Grouped Query Attention (GQA) addresses this issue by grouping queries and mean-pooling the corresponding key-value heads. We introduce enhancements to GQA, focusing on two novel approaches that deviate from the static nature of grouping.
  arXiv  Detail & Related papers  (2024-08-15T23:34:04Z)
• Calibration in Machine Learning Uncertainty Quantification: beyond consistency to target adaptivity [0.0]
  This article aims to show that consistency and adaptivity are complementary validation targets, and that a good consistency does not imply a good adaptivity. Adapted validation methods are proposed and illustrated on a representative example.
  arXiv  Detail & Related papers  (2023-09-12T13:58:04Z)
• On the Sample Complexity of Imitation Learning for Smoothed Model Predictive Control [27.609098229134]
  We show how a smoothed expert can be designed for a general class of systems. We prove on the optimality gap of the analytic center associated with a convex Lipschitz function.
  arXiv  Detail & Related papers  (2023-06-02T20:43:38Z)
• Grounding Representation Similarity with Statistical Testing [8.296135566684065]
  We argue that measures should have sensitivity to changes that affect functional behavior, and specificity against changes that do not. We quantify this through a variety of functional behaviors including probing accuracy and robustness to distribution shift. We find that current metrics exhibit different weaknesses, note that a classical baseline performs surprisingly well, and highlight settings where all metrics appear to fail.
  arXiv  Detail & Related papers  (2021-08-03T17:58:16Z)
• Accelerating Quadratic Optimization with Reinforcement Learning [39.64039435793601]
  We show how Reinforcement Learning can learn a policy to tune parameters to accelerate convergence. Our policy, RLQP, significantly outperforms state-of-the-art QP solvers by up to 3x. RLQP generalizes surprisingly well to previously unseen problems with varying dimension and structure from different applications.
  arXiv  Detail & Related papers  (2021-07-22T17:59:10Z)
• Causally-motivated Shortcut Removal Using Auxiliary Labels [63.686580185674195]
  Key challenge to learning such risk-invariant predictors is shortcut learning. We propose a flexible, causally-motivated approach to address this challenge. We show both theoretically and empirically that this causally-motivated regularization scheme yields robust predictors.
  arXiv  Detail & Related papers  (2021-05-13T16:58:45Z)
• Meta-Learning with Adaptive Hyperparameters [55.182841228303225]
  We focus on a complementary factor in MAML framework, inner-loop optimization (or fast adaptation) We propose a new weight update rule that greatly enhances the fast adaptation process.
  arXiv  Detail & Related papers  (2020-10-31T08:05:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.