ICML 2026 Papers

International Conference on Machine Learning

Decision breakdown

6,341papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

Browse all papers

Structured 4D Latent Predictive Model for Robot Planning

Zhiyi Li, Peilin Wu, Xiaoshen Han, Ruojin Cai, Yilun Du

Expand abstract

Video predictive models are emerging as a powerful paradigm in robotics, offering a promising path toward task generalization, long-horizon planning, and flexible decision-making. However, prevailing approaches often operate on 2D video sequences, inherently lacking the 3D geometric understanding necessary for precise spatial reasoning and physical consistency. We introduce a **Structured 4D Latent Predictive Model**, which predicts the evolution of a scene’s 3D structure in a structured latent space conditioned on observations and textual instructions. Our representation encodes the scene holistically and can be decoded into diverse 3D formats, enabling a more complete and 3D consistent scene understanding. This structured 4D latent predictive model serves as a planner, generating future scenes that are translated into executable actions by a goal-conditioned inverse dynamics module. Experiments demonstrate that our model generates futures with strong visual quality, substantially better 3D consistency and multi-view coherence compared to state-of-the-art video-based planners. Consequently, our full planning pipeline achieves superior performance on complex manipulation tasks, exhibits robust generalization to novel visual conditions, and proves effective on real-world robotic platforms. Our website is available at https://structured-4d-model.github.io/.

Poster
Configurable Reward Model for Balanced Safety Alignment

Zhengping Jiang, Mehran Khodabandeh, Akash Bharadwaj, Manik Bhandari, Mayur Srungarapu, Anqi Liu, Benjamin Van Durme, Li Chen

Expand abstract

Aligning large language models (LLMs) to heterogeneous and rapidly evolving safety requirements remains a critical challenge. Existing instruction-tuned LLMs and standalone safety classifiers often fail to generalize to new safety configurations, motivating the need for Reward Models (RMs) that are explicitly configurable to changing specifications. We introduce the Configurable Safety Reward Model (CSRM), which is jointly optimized for calibrated safety compliance and reward modeling. Our approach is supported by configuration-targeted data augmentation that enforces instruction adherence while preserving relative severity structure. The resulting RM is sensitive to fine-grained safety configurations and conversational nuances, substantially improving generalization to previously unseen safety configurations. CSRM achieves state-of-the-art performance on recent configurable safety benchmarks, including CoSApien (94.6\% F1) and DynaBench (75.8\% F1), without requiring additional human annotation. When used for downstream safety alignment, CSRM yields LLMs with a significantly improved helpfulness–safety tradeoff compared to existing baselines.

Poster
Merge to Remember: Sharpness-Aware Isotropic Merging for Continual Learning

Qun Yang, Enneng Yang, Wei Chen, Li Shen, Long Lan

Expand abstract

Continual learning with large pre-trained models offers significant potential for cross-task knowledge accumulation, but faces critical challenges such as catastrophic forgetting and parameter interference, especially when historical data is unavailable. Existing approaches typically rely on sequential fine-tuning or model merging strategies, yet often overlook the impact of loss landscape sharpness and dominant singular value directions, which leads to subspace misalignment and severe knowledge forgetting. In this paper, we propose the Sharpness-Aware Isotropic Merging (SAIM) framework, which introduces targeted optimizations in both the fine-tuning and merging stages to address these issues. Specifically, SAIM consists of two synergistic modules: (1) a Sharpness-Aware Block Coordinate Descent (SA-BCD) optimizer that guides the model toward flatter minima and selectively updates the most task-sensitive parameters, thereby mitigating parameter interference and enhancing robustness; (2) an adaptive isotropic merging algorithm that dynamically balances the singular value spectrum across tasks, effectively preventing the model from overemphasizing any single task direction, maintaining balanced knowledge representation, and improving subspace alignment. Extensive experiments on vision and language benchmarks demonstrate that SAIM achieves 5-10\% higher accuracy than existing methods and maintains robust performance as the number of tasks increases. Our code are available at \url{https://github.com/Yangqun123456/SAIM}.

Poster
MultiLoReFT: Decoupling Shared and Modality-Specific Subspaces in Multimodal Learning via Low-Rank Representation Fine-Tuning

Sana Tonekaboni, Viktoria Schuster, Caroline Uhler

Expand abstract

Real-world perception and decision making are inherently multimodal, integrating complementary signals across modalities. However, training multimodal models faces two main obstacles. First, collecting large-scale, well-aligned paired multimodal datasets is often impractical, making end-to-end multimodal training difficult. Second, existing multimodal representations frequently entangle information shared across modalities with modality-specific information, hindering interpretability and control. We introduce MultiLoReFT, an efficient and scalable low-rank representation fine-tuning framework for multimodal learning with pretrained unimodal models. MultiLoReFT extends low-rank adaptation to the multimodal setting and learns interpretable projection subspaces that decouple shared and modality-specific information. Across simulated and real-world benchmarks, it produces representations that support multimodal prediction while explicitly revealing how shared and modality-specific information is distributed across modalities.

Poster
Beyond Euclidean Summaries: Online Change Point Detection for Distribution-Valued Data

Yingyan Zeng, Zipan Huang, Xiaoyu Chen

Expand abstract

Existing online change-point detection (CPD) methods rely on fixed-dimensional Euclidean summaries, implicitly assuming that distributional changes are well captured by moment-based or feature-based representations. They can obscure important changes in distributional shape or geometry. We propose an intrinsic distribution-valued CPD framework that treats streaming batch data as a stochastic process on the 2-Wasserstein space. Our method detects changes in the law of this process by mapping each empirical distribution to a tangent space relative to a pre-change Fréchet barycenter, yielding a reference-centered local linearization of 2-Wasserstein space. This representation enables sequential detectors by adapting classical multivariate monitoring statistics to tangent fields. We provide theoretical guarantees and demonstrate, via synthetic and real-world experiments, that our approach detects complex distributional shifts with reduced detection delay at matched $\mathrm{ARL}_0$ compared with moments-based and model-free baselines. The code is available at https://github.com/yyzeng43/IDD-icml.

Poster
LeakGFN: Robust Molecular Generation in Generative Flow Networks via Flow Decomposition

Hwanhee Kim, Seungyeon Choi, SangHyun Park

Expand abstract

Generative Flow Networks (GFlowNets) have emerged as a powerful framework for molecular generation, sampling diverse candidates proportionally to a reward function. However, the vast chemical space necessitates truncating trajectory length, forcing models to treat incomplete molecular fragments as terminal states alongside valid molecules. This conflation distorts the learned distribution by allocating probability mass to chemically meaningless states. We propose LeakGFN, a dual-head architecture that decomposes flow into two components: a chemical head modeling flow over the full chemical space, and a valid head estimating the fraction of flow reaching valid molecules within the truncation boundary. Through this decomposition, the valid head implicitly learns molecular reachability without explicit supervision. We prove that LeakGFN recovers the correct distribution over accessible molecules under mild assumptions. Experiments on five molecular optimization tasks demonstrate consistent improvements over flow matching baselines, achieving state-of-the-art performance on four out of five tasks. Our module integrates as a plug-and-play enhancement into existing frameworks, improving performance on both pocket-conditioned and multi-objective generation tasks.

Poster
SMM Transformer: Leveraging Spiking Neural Networks for Multimodal Tasks

Xiubo Liang, Jinxing Han, Yuke Li, Haoqi Zhu, Yu Zhao, Hongzhi Wang

Expand abstract

Spiking Neural Networks (SNNs) enable event-driven computation with sparse activations, but building multimodal Transformers on SNNs is hindered by unstable training in deep spiking stacks and a mismatch between dense softmax attention and spike-based communication. We propose SMM Transformer, an SNN-based multimodal Transformer framework that combines (i) a Parallel LIF with Multistage Learnable Parameters (PLMP) neuron and a tailored P-STBP algorithm to stabilize training, (ii) a spike-driven attention approximation (SMSA) with a lightweight self-compensation branch, and (iii) a spiking mixture-of-experts (SMoE) module for modality-aware fusion. Across visual and multimodal benchmarks, SMM Transformer achieves competitive accuracy compared to ANN baselines while reducing the estimated compute energy of the attention module by up to 97\% under a standard MAC/AC cost model.

Poster
AlignedNorm: Prompting Vision–Language Models via Coupled Prompt Field

Qi Ma, Chen-Yang Wang, Dehong Gao, Deng-Ping Fan

Expand abstract

Prompt learning for vision-language models (VLMs) primarily follows end-to-end or decoupled routes to balance base and new task performance, but suffers a fundamental bottleneck: sample-wise optimization within task-specific feature spaces traps models in local optima, hindering global optimality. To address this, we identify a key insight that VLMs can be prompted within a Coupled Prompt Field-a shared space where base and new tasks are mutually constrained-and present AlignedNorm, which enforces the field coupling. By dynamically aligning the norms of prompts to VLMs' native scale, our method enables joint optimization of both tasks. Without complex designs, our method matches leading decoupled approaches on 15 datasets across 4 experimental settings, offering both a new perspective and a practical solution to the local-optima dilemma in prompt learning.

Poster
OmniSIFT: Modality-Asymmetric Token Compression for Efficient Omni-modal Large Language Models

Yue Ding, Yiyan Ji, Jungang Li, Xuyang Liu, Xinlong Chen, Junfei Wu, Bozhou Li, Bohan Zeng, Yang Shi, Yushuo Guan, Yuanxing Zhang, Jiaheng Liu, Qiang Liu, Pengfei Wan, Liang Wang

Expand abstract

Omni-modal Large Language Models (Omni-LLMs) have demonstrated strong capabilities in audio-video understanding tasks. However, their reliance on long multimodal token sequences leads to substantial computational overhead. Despite this challenge, token compression methods designed for Omni-LLMs remain limited. To bridge this gap, we propose OmniSIFT (Omni-modal Spatio-temporal Informed Fine-grained Token compression), a modality-asymmetric token compression framework tailored for Omni-LLMs. Specifically, OmniSIFT adopts a two-stage compression strategy: (i) a spatio-temporal video pruning module that removes video redundancy arising from both intra-frame structure and inter-frame overlap, and (ii) a vision-guided audio selection module that filters audio tokens. The entire framework is optimized end-to-end via a differentiable straight-through estimator. Extensive experiments on five representative benchmarks verify the efficacy and robustness of OmniSIFT. Notably, for Qwen2.5-Omni-7B, OmniSIFT adds 4.85M parameters while still achieving lower latency than training-free baselines such as OmniZip. With only 25% of the original token context, OmniSIFT consistently outperforms all compression baselines and even surpasses the full-token model on several tasks.

Poster
DistFlow: A Fully Distributed RL Framework for Scalable and Efficient LLM Post-Training

Zhixin Wang, Jiaming Xu, Tianyi Zhou, Mingjun Zhang, Liming Liu, Jiarui Hu, Dian Yang, Tongyu Wang, Ping Zhang, Jinlong Hou, Siyuan Feng, Yuan Qi, Yuan Cheng

Expand abstract

Effectively scaling Reinforcement Learning (RL) is crucial for enhancing the reasoning and alignment of Large Language Models. The massive data and complex execution flows inherent in these tasks require a distributed architecture capable of efficient scaling. However, to simplify programming and dependency management, mainstream frameworks often rely on a centralized architecture where a single node dispatches both control and data. This inherent coupling creates significant communication bottlenecks, severely limiting system scalability and efficiency. We present DistFlow, a novel, fully distributed RL framework that adopts a multi-controller paradigm. By decoupling data transmission from control dispatch, DistFlow establishes a parallelism-aware, decentralized Data Coordinator that leverages local caching, load balancing, and asynchronous double buffer to minimize communication overhead and mitigate straggler effects. For control logic, it introduces a task scheduler built upon Directed Acyclic Graph (DAG) that facilitates fine-grained, independent execution. Experimental results demonstrate that DistFlow achieves near-linear scalability up to 512 GPUs and delivers up to a 2.63x throughput improvement over state-of-the-art (SOTA) frameworks.

Poster
IQA-Spider: Unifying Multi-Granularity Image Quality Assessment with Reasoning, Grounding and Referring

Xinge Peng, Yiting Lu, Xin Li, Zhibo Chen

Expand abstract

We present IQA-Spider, the first image quality assessment (IQA) framework that unifies reasoning, grounding, and referring into a single LMM-based framework for multi-granularity quality understanding. Existing LMM-based IQA methods typically support only partial perception dimensions, such as quality description and question answering~(\textit{i.e.}, reasoning) or pixel-level grounding. This limitation largely stems from the absence of (i) a unified task and data formulation and (ii) effective optimization paradigms for multi-granularity learning. To address these limitations, we formulate a rigorous four-task paradigm covering global and local quality description, pixel-level grounding, and region-level referring. Based on this formulation, we construct a corresponding IQA dataset with a scalable and automatic annotation pipeline, thereby providing a solid foundation for unified multi-granularity learning. To further enable unified perception, we adopt a conflict-free two-stage design that progressively extends text-level multi-granularity understanding to pixel-level grounding: (i) the first stage equips the model with fine-grained text-level reasoning across multiple IQA tasks, and (ii) the second stage introduces a training-free text-to-point grounding paradigm, which bridges textual semantics and pixel-level perception by mapping token logits to spatial coordinates. Based on these efforts, we achieve IQA-Spider with unified multi-granularity explainable image quality assessment. Extensive experiments across multiple benchmarks demonstrate strong performance, validating the effectiveness and versatility of the proposed formulation and framework. Our code and dataset will be released at: \href{https://github.com/Helen1p/IQA-Spider.git}{https://github.com/Helen1p/IQA-Spider.git}.

Poster
StableVLA: Towards Robust Vision-Language-Action Models without Extra Data

YIYANG FU, Chubin Zhang, Shukai Gong, Yufan Deng, Kaiwei Sun, Qiyang Min, Qibin Hou, Yansong Tang, Jianan Wang, Daquan Zhou

Expand abstract

It is infeasible to encompass all possible disturbances within the training dataset. This raises a critical question regarding the robustness of Vision-Language-Action (VLA) models when encountering unseen real-world visual disturbances, particularly under imperfect visual conditions. In this work, we conduct a systematic study based on recent state-of-the-art VLA models and reveal a significant performance drop when visual disturbances absent from the training data are introduced. To mitigate this issue, we propose a lightweight adapter module grounded in information theory, termed the Information Bottleneck Adapter (IB-Adapter), which selectively filters potential noise from visual inputs. Without requiring any extra data or augmentation strategies, IB-Adapter consistently improves over the baseline by an average of 30%, while adding fewer than 10M parameters, demonstrating notable efficiency and effectiveness. Furthermore, even with a 14x smaller backbone (0.5B parameters) and no pre-training on the Open X-Embodiment dataset, our model StableVLA achieves robustness competitive with 7B-scale state-of-the-art VLAs. With negligible parameter overhead (<10M), our approach maintains accuracy on long-horizon tasks and surpasses OpenPi under real-world visual disturbances. The code will be made publicly available.

Poster
QEDBench: Quantifying the Alignment Gap in Automated Evaluation of University-Level Mathematical Proofs

Santiago Gonzalez, Alireza Amiri Bavandpour, Peter Ye, Edward Zhang, Ruslans Aleksejevs, Todor Antić, Polina Baron, Sujeet Bhalerao, Shubhrajit Bhattacharya, Zachary Burton, John Byrne, Hyungjun Choi, Nujhat Ahmed Disha, Koppány István Encz, Yuchen Fang, Robert Joseph George, Ebrahim Ghorbani, Alan Goldfarb, Jing Guo, Meghal Gupta, Stefano Huber, Annika Kanckos, Minjung Kang, Hyun Jong Kim, Dino Lorenzini, Levi Lorenzo, Tianyi Mao, Giovanni Marzenta, Ariane M. Masuda, Lukas Mauth, Ana Mickovic, Andrés Miniguano-Trujillo, Antoine Moulin, Wenqi Ni, Tomos Parry, Kevin Ren, Hossein Roodbarani, Mathieu Rundström, Manjil Saikia, Detchat Samart, Rebecca Steiner, Connor Stewart, Dhara Thakkar, Jeffrey Tse, Vasiliki Velona, Yunhai Xiang, Sibel Yalçın, Jun Yan, Ji Zeng, Arman Cohan, Quanquan C. Liu

Expand abstract

As Large Language Models (LLMs) saturate elementary benchmarks, the research frontier has shifted from generation to the reliability of automated evaluation. We demonstrate that standard "LLM-as-a-Judge" protocols suffer from a systematic evaluation Alignment Gap when applied to upper-undergraduate to early graduate level mathematics. To quantify this, we introduce QEDBench, the first benchmark to systematically measure alignment with human experts on undergraduate-level math proofs by contrasting course-specific rubrics against expert common knowledge criteria. By deploying a dual-evaluation matrix ($7$ judges $\times$ $5$ solvers) against 1,000+ hours of human evaluation, we reveal that certain frontier evaluators like Claude 4.5 Opus exhibit significant positive bias (up to $+0.28$ mean score inflation), effectively "hallucinating rigor" in flawed proofs. Furthermore, we uncover a critical reasoning disparity: while Gemini 3.0 Pro achieves state-of-the-art performance (0.91 raw score), specialized reasoning models like o3-deep-research collapse in discrete domains, dropping to 42.1\% accuracy in Graph Theory. We release QEDBench as a public benchmark for evaluating and improving AI judges.

Poster
AgentWebBench: Benchmarking Multi-Agent Coordination in Agentic Web

Shanshan Zhong, Kate Shen, Chenyan Xiong

Expand abstract

Agentic Web is an emerging paradigm where autonomous agents help users use online information. As the paradigm develops, content providers are also deploying agents to manage their data and serve it through controlled interfaces. This shift moves information access from centralized retrieval to decentralized coordination. To study this setting, we introduce AgentWebBench, a benchmark that evaluates how well a user agent synthesizes answers by interacting with website-specific content agents. We evaluate four tasks that cover common web information needs, spanning ranked retrieval (web search, web recommendation) and open-ended synthesis (question answering, deep research). Across seven advanced LLMs and three coordination strategies, multi-agent coordination generally lags behind centralized retrieval as expected, because user agent cannot directly access the corpus, but the gap shrinks with model scale and can even outperform centralized retrieval on question answering. This benchmark also enables us to study properties of the emerging paradigm of the digital world. We find that decentralized access concentrates traffic toward a small set of websites, test time scaling improves both interaction reliability and task performance, and strong results require sufficient interactions guided by careful planning. Finally, our failure analysis suggests that user agents need better planning and answer synthesis, while content agents need more reliable retrieval and evidence quality. Code, data, and APIs are released on https://github.com/cxcscmu/AgentWebBench.

Poster
Riemannian stochastic optimization for sufficient dimension reduction

Thibault Pautrel, François Portier

Expand abstract

Sufficient dimension reduction (SDR) makes high-dimensional regression tractable by projecting the covariates onto a low-dimensional subspace that preserves the conditional mean of the response. Existing gradient-based estimators either operate in the ambient space and suffer from the curse of dimensionality, or localize in the reduced space at a per-outer-iteration cost at least quadratic in the sample size. We show that minimizers of the population Minimum Average Variance Estimation (MAVE) risk approximate the same Grassmannian target as the Outer Product of Gradients (OPG), and recast the empirical criterion as a smooth maximization on the Stiefel manifold with closed-form Riemannian gradient. The resulting algorithm, SMAVE, combines sparse projected-space nearest-neighbor localization with Riemannian stochastic gradient ascent. A simplified version comes with almost-sure convergence and a non-asymptotic rate matching the standard non-convex stochastic first-order scaling. Empirically, SMAVE matches or improves on RMAVE's synthetic subspace recovery at moderate-to-high ambient dimension, and on four real datasets it uniformly improves over OPG and is competitive with or outperforms RMAVE at orders of magnitude lower runtime.

Poster
DRIVE: Best Data Scheduling Practices for Reinforcement Learning with Verifiable Reward in Competitive Code Generation

Speed Zhu, Chuheng Zhang, Jianwei Cai, Guang Chen, Lulu Wu, Xiaolong Xu, Xuyun Zhang, Saiyong Yang, Wiggin Zhou

Expand abstract

Recent success of large reasoning models (such as OpenAI o1 and DeepSeek R1) have spurred a resurgence of interest in reinforcement learning from verifiable rewards (RLVR). However, progress is still largely driven by RL algorithm design, while data scheduling -- the data-side decisions that determine what the model trains on over time -- is critical but remains underexplored. Therefore, data scheduling becomes the focus of this paper, including how to curate data for supervised fine-tuning (SFT) and how to select prompts and collect rollouts for reinforcement learning (RL). We introduce a pipeline with careful designs on data scheduling, consisting of hardness-prioritized SFT and two-stage RL. Specifically, we first fine-tune the base model on supervision data that is curated to prioritize difficulty based on both arena learning and classification. Then, we introduce two-stage RL where a decreased max sequence length during rollout is used in the first stage to expand entropy and reduce repetition, and a large number rollouts per prompt and curriculum design are adopted in the second stage to encourage exploration for challenging problems. We implement this pipeline on Qwen2.5-32B and an internal 389B MoE model, and evaluate them on a wide range of benchmarks including challenging LeetCode and Codeforces weekly contests. The results not only indicate the effectiveness and scalability of our pipeline but also demonstrate our model achieve sota of 32B models in competitive code generation.

Poster
QiMeng-PerceptOS: Semantic-Aware Kernel Optimization for OS-Intensive Workloads via Hardware-Software Alignment

Huilai Chen, Yuanbo Wen, Liangfeng Li, Shaohui Peng, Jingzhe Zhu, Jun Bi, Xuzhi Zhang, Qi Guo, Ling Li, Yunji Chen

Expand abstract

Optimizing OS kernels for specific applications is vital for peak performance, yet existing LLM-based methods struggle with a semantic mismatch between generalized reasoning and low-level system behaviors. As a result, these static, open-loop approaches suffer from runtime blindness, configuration fragmentation, and search drift, ultimately failing to unlock the system’s performance potential. To address this, we propose QiMeng-PerceptOS, an autonomous framework that shifts the paradigm to perception-driven tuning. QiMeng-PerceptOS integrates: (1) a Perception Module that aligns raw telemetry into high-fidelity semantic fingerprints; (2) a Global Search Module utilizing a Bi-level Hierarchical Induction Tree (BHIT) for global navigation and efficient pruning; and (3) a Posterior Enhancement Module to suppress hallucinations via trajectory synthesis. Experiments across diverse workloads show that it achieves significant performance breakthroughs by optimizing kernel configurations, reaching 296.6% of default Redis throughput and surpassing SOTA baselines by 32.6% within only 15 iterations. By establishing a perception-driven closed-loop, QiMeng-PerceptOS provides new insights for fully automated, large-scale system optimization.

Poster
AVGen-Bench: A Task-Driven Benchmark for Multi-Granular Evaluation of Text-to-Audio-Video Generation

Ziwei Zhou, Zeyuan Lai, Rui Wang, Yifan Yang, Yuqing Yang, Qi Dai, Lili Qiu, Chong Luo

Expand abstract

Text-to-Audio-Video (T2AV) generation is rapidly becoming a core interface for media creation, yet its evaluation remains fragmented. Existing benchmarks largely assess audio and video in isolation or rely on coarse embedding similarity, failing to capture fine-grained joint correctness required by realistic prompts. We introduce AVGen-Bench, a task-driven benchmark for T2AV generation, featuring high-quality prompts across 11 real-world categories. To support comprehensive assessment, we propose a multi-granular evaluation framework that combines lightweight specialist models with Multimodal Large Language Models (MLLMs), enabling evaluation from perceptual quality to fine-grained semantic controllability. Our evaluation reveals a pronounced gap between strong audio-visual aesthetics and weak semantic reliability, including persistent failures in text rendering, speech coherence, physical reasoning, and universal breakdown in musical pitch control.

Poster
Learning Adaptive Perturbation-Conditioned Contexts for Robust Transcriptional Response Prediction

Yinhua Piao, Hyomin Kim, Seonghwan Kim, Yunhak Oh, Junhyeok Jeon, Sang-Yeon Hwang, Jaechang Lim, Woo Youn Kim, Chanyoung Park, Sungsoo Ahn

Expand abstract

Predicting high-dimensional transcriptional responses to genetic perturbations is challenging because signals are sparse and experimental noise is severe. Existing methods often suffer from mean collapse, achieving high correlation by predicting the global average expression rather than perturbation-specific responses, which yields false positives and poor interpretability. Methods that add biological knowledge graphs typically treat them as dense, static priors shared across perturbations, propagating noise. We propose AdaPert, which counters mean collapse by extracting a sparse, perturbation-specific subgraph via differentiable node selection, then suppressing spurious variation in non-responsive genes while emphasizing differentially expressed ones. Across multiple benchmarks, AdaPert outperforms existing baselines, with the largest gains on DEG-aware metrics.

Poster
Affine-Scaled Attention: Towards Flexible and Stable Transformer Attention

Jeongin Bae, Baeseong park, Gunho Park, Minsub Kim, Joonhyung Lee, Junhee Yoo, Sunghyeon Woo, Jiwon Ryu, Se Jung Kwon, Dongsoo Lee

Expand abstract

Transformer attention is typically implemented using softmax normalization, which enforces attention weights with unit sum normalization. While effective in many settings, this constraint can limit flexibility in controlling attention magnitudes and may contribute to overly concentrated or unstable attention patterns during training. Prior work has explored modifications such as attention sinks or gating mechanisms, but these approaches provide only limited or indirect control over attention reweighting. We propose Affine-Scaled Attention, a simple extension to standard attention that introduces input-dependent scaling and a corresponding bias term applied to softmax-normalized attention weights. This design relaxes the strict normalization constraint while maintaining aggregation of value representations, allowing the model to adjust both the relative distribution and the scale of attention in a controlled manner. We empirically evaluate Affine-Scaled Attention in large-scale language model pretraining across multiple model sizes. Experimental results show consistent improvements in training stability, optimization behavior, and downstream task performance compared to standard softmax attention and attention sink baselines. These findings suggest that modest reweighting of attention outputs provides a practical and effective way to improve attention behavior in Transformer models.

Poster