Smoothing Slot Attention Iterations and Recurrences

Conference: ICML 2026
arXiv: 2508.05417
Code: https://github.com/Genera1Z/SmoothSA (available)
Area: Multimodal VLM / Object-Centric Learning / Slot Attention
Keywords: Object-Centric Learning, Slot Attention, Query Warm-up, Video OCL, Self-Distillation

TL;DR

Addressing two long-standing but overlooked issues in Slot Attention for image and video object-centric learning—namely, "insufficient information in cold-start queries" and "forced unification of aggregation transformations for first/non-first frames"—the authors propose SmoothSA: a self-distillation-based lightweight warm-up module that injects sample-specific information into queries, and a scheduling scheme where the first frame undergoes three full iterations while non-first frames only run one. This approach achieves new SOTA on both image and video OCL benchmarks.

Background & Motivation

Background: Object-centric learning (OCL) is a paradigm that represents visual scenes as a set of independent object/background vectors (slots). Such structured, compact representations often outperform dense feature maps in downstream reasoning, video prediction, and generative tasks. Most mainstream implementations are built on Slot Attention (SA): image features are treated as key/value, \(n\) query slots act as "competitors," and several rounds of iterative cross-attention assign patches to different slots, learning object-level representations in a self-supervised manner via reconstruction loss. For video OCL, the standard approach (STEVE series) recursively applies image-based SA across frames: the first frame's query is initialized as in the image case, while non-first frame queries are predicted from the previous frame's slots via a Transformer encoder.

Limitations of Prior Work: The authors identify two issues that have been implicitly accepted but never directly addressed. First, "query cold start": whether initialized from learnable Gaussians or positional priors, the initial query slots contain only dataset-level priors, lacking any sample-specific cues; this undermines aggregation quality on the first image/video frame, forcing the model to "guess" through more iterations. Second, "transformation homogenization": in videos, the first frame's query is cold-started and information-poor, while non-first frame queries inherit rich sample information from previous slots, yet most methods treat both identically, applying the same three SA iterations and ignoring the information gap.

Key Challenge: Aggregation accuracy depends on the sample information carried by the query, but the SA framework lacks differentiated processing paths for queries with different prior information. As a result, all queries either share a crude cold-start or all frames use a fixed iteration count, regardless of information state.

Goal: Without altering the core OCL model, address: (i) how to inject sample-level information into cold-start queries for images/video first frames; (ii) how to enable SA to use different aggregation strengths for first and non-first video frames.

Key Insight: The authors observe that a trained OCL model can already output "good" slots, so the "ideal query" can be supervised by existing slots—naturally supporting self-distillation. Moreover, three iterations are designed for cold-start queries; for non-first frame queries already close to the true distribution, further iterations may cause over-adjustment.

Core Idea: Insert a small module before SA to warm up cold-start queries into informative queries approximating slots, and decouple the "three iterations for first frame" from "single iteration for non-first frames"—thus symmetrizing information content and transformation strength to address both issues.

Method

Overall Architecture

SmoothSA fully retains the classic OCL encode-aggregate-decode structure: the encoder maps images/video frames to features \(F \in \mathbb{R}^{h\times w\times c}\), the aggregator \(\phi_a\) (i.e., SA module) aggregates features into \(n\) slots \(S \in \mathbb{R}^{n\times c}\), and the decoder reconstructs the input from slots for self-supervision. SmoothSA introduces two minor modifications at the aggregator input and inter-frame scheduling: (1) before the cold-start query \(Q_1\) enters SA, it passes through a "warmer" \(\phi_p\) to produce an informative query \(\tilde Q_1\); (2) for videos, only the first frame runs the standard 3 SA iterations to obtain slot \(S_1\), while non-first frames run only 1 SA iteration, with their queries directly predicted from the previous frame's slots via the standard transition module. The changes are minimal and can be directly applied to any SA-based image/video OCL model.

Key Designs

1. Query Warmer \(\phi_p\) (Self-distillation to Move Cold-Start Query Closer to Slot Distribution):

   • Function: Jointly maps the information-poor cold-start query \(Q_1\) and input feature \(F_1\) into an informative query \(\tilde Q_1\) approximating the current sample's slot, providing a better starting point for subsequent SA iterations.
   • Mechanism: \(\phi_p\) is a lightweight module, structurally akin to a single query-feature cross-attention/MLP; it takes \(Q_1\) and \(F_1\) as input and outputs \(\tilde Q_1 \in \mathbb{R}^{n\times c}\). The supervision signal comes from the OCL model's own slot output \(S_1\) for the current batch—i.e., enforce \(\tilde Q_1 \approx S_1\) via stop-gradient self-distillation, with loss \(\mathcal{L}_p = \|\tilde Q_1 - \text{sg}(S_1)\|^2\). During training, \(\phi_p\) is co-optimized with the OCL backbone; at inference, \(\tilde Q_1\) replaces \(Q_1\) as SA input.
   • Design Motivation: Previous works (e.g., BO-QSA, MetaSlot) attempted to enrich query priors via multiple Gaussians or object prototype codebooks, but still failed to inject "current sample" information into the query. \(\phi_p\) is the first to explicitly inject sample features into the query via a differentiable channel, effectively shifting the SA iteration curve forward—starting from a point already close to the optimal slot, thus reducing iteration error.

2. Heterogeneous Iteration Scheduling Across Video Frames (Multiple Iterations for First Frame, Single for Non-First Frames):

   • Function: Eliminates the waste and disturbance of applying 3 SA iterations to all frames in STEVE-like frameworks, matching "iteration strength" to "query information content."
   • Mechanism: For the first frame (cold-start query, large information gap), \(\Phi_a\) unfolds as \(S_1^{(i)}, M_1^{(i)} = \phi_a(S_1^{(i-1)}, F_1)\), \(i=1,2,3\), with \(S_1 := S_1^{(3)}\); for non-first frames \(t\ge 2\), query \(Q_t\) is directly obtained from the previous slot \(S_{t-1}\) via the standard transition network (e.g., STEVE's Transformer encoder), then only one SA iteration is run: \(S_t, M_t = \phi_a(Q_t, F_t)\). Visualization shows that multiple iterations for non-first frames can cause unnecessary alignment oscillations for queries already close to the true distribution.
   • Design Motivation: The 3-iteration setting was originally intended to allow multiple alignment steps for information-poor queries; if the query is already close to the true slot, further iterations introduce redundant updates and dilute the temporal information. Differentiating "number of transformations" essentially equips queries in different information states with appropriately strong alignment paths, aligning with the intuition of "harder inputs require more computation."

3. Tightly Coupled Training with Backbone OCL Self-Supervised Loss:

   • Function: Enables the warmer and backbone OCL to interact end-to-end under the same self-supervised objective, requiring no external labels or extra training stages.
   • Mechanism: The total loss is \(\mathcal{L} = \mathcal{L}_{rec} + \lambda \mathcal{L}_p\), where \(\mathcal{L}_{rec}\) is the OCL's reconstruction loss (e.g., cross-entropy for dVAE tokens or pixel MSE), and \(\mathcal{L}_p\) is the stop-gradient distillation term between \(\phi_p\)'s output and the current slot, with \(\lambda\) as the balancing weight. This ensures \(\phi_p\) tracks the backbone OCL's slot distribution at every step without feeding noise back into the SA learning.
   • Design Motivation: Treating \(\phi_p\) as a regular trainable pre-layer and training it end-to-end with SA can cause unstable coupling, where \(\phi_p\) pulls queries toward noisy slots and SA is disturbed by noisy queries. Stop-gradient self-distillation structurally breaks this feedback loop, ensuring \(\phi_p\) always tracks the OCL's current optimum, akin to BYOL and related frameworks.

Loss & Training

All modules are jointly optimized during training, with objective \(\mathcal{L} = \mathcal{L}_{rec} + \lambda \mathcal{L}_p\). The slot side of \(\mathcal{L}_p\) uses stop-gradient, so the warmer module unidirectionally tracks the backbone OCL's optimal slot distribution. During video training, iteration count differentiation is implemented via an if-branch in the forward pass, incurring no extra memory or training stages.

Key Experimental Results

Main Results

The authors attach SmoothSA to various backbones (DINOSAUR, STEVE, SAVi, RandSF.Q, etc.) on standard image OCL (COCO, Movi-E, etc.) and video OCL (Movi-D, Movi-E, YT-VIS, Physion) benchmarks, comparing with existing SA variants on object discovery (mIoU/FG-ARI), object recognition, and visual reasoning tasks.

Task	Dataset	Metric	Backbone Baseline	+ SmoothSA	Gain
Image Object Discovery	COCO	mIoU / FG-ARI	DINOSAUR baseline	Improved	Consistent increase
Video Object Discovery	Movi-E	FG-ARI	STEVE / RandSF.Q	Improved	Further SOTA gains
Visual Reasoning	Physion	Accuracy	SA baseline	Improved	Significant

(See paper tables for specific numbers; the core finding is that SmoothSA consistently improves all three metrics across backbones and datasets.)

Ablation Study

Configuration	FG-ARI / mIoU Trend	Notes
Full SmoothSA	Best	Warmer + heterogeneous iteration
w/o Warmer \(\phi_p\)	Significant drop	Validates query cold start as main bottleneck for first-frame aggregation
w/o Heterogeneous Iteration (non-first frames still 3x)	Drop	Shows multiple iterations for information-rich queries are detrimental
\(\phi_p\) without stop-gradient	Unstable training/performance drop	Validates necessity of stop-gradient in self-distillation
Larger \(\phi_p\)	Saturated gains	Smaller warmer is better; a few thousand parameters suffice

Key Findings

• The extremely lightweight warmer module yields stable gains, indicating that "initialization bias" accounts for much of SA's iteration error; many traditional iterations merely compensate for poor starting points rather than true aggregation.
• Single iteration for non-first frames not only avoids performance drop but improves results, suggesting a general rule for SA: "the more informative the query, the fewer iterations needed"—providing direct experience for future adaptive iteration-count SA designs.
• Improvements appear simultaneously in FG-ARI (proxy for object discovery quality) and downstream reasoning, indicating that enhancing query informativeness benefits both segmentation and downstream representation learning.

Highlights & Insights

• The warmer leverages OCL's own outputs for self-distillation, incurring almost no extra annotation or computation cost, yet addresses the long-ignored "query cold start" problem in the simplest way, with a minimal and theoretically clear structure.
• Treating "iteration count" as a hyperparameter adjustable according to query information content, rather than a rigid 3, aligns "information content → computation" and can be transferred to any iterative refinement framework (e.g., iterative query decoders, recurrent mask refinement).
• The paper unifies two seemingly independent issues (image first-frame cold start, video inter-frame homogenization) under the abstraction of "smooth SA iterations and recurrences," providing a clear conceptual framework—this approach of subsuming multiple engineering issues under a unified perspective is valuable for inspiring new methods.

Limitations & Future Work

• The warmer module relies on the backbone OCL's ability to output "sufficiently good slots" as distillation targets; during early training when the OCL backbone collapses, the supervision signal for \(\phi_p\) is noisy and may require warm-up scheduling.
• The inter-frame iteration count is hardcoded as "3 for first frame, 1 for others," without considering scene changes or shot transitions that may require re-cold-starting; a reasonable extension is to introduce a lightweight signal to determine whether the current frame's query should be treated as cold-start.
• Experiments are mainly conducted on synthetic videos and medium-scale real data, not covering large-scale long videos or real open-world scenarios; whether query information degrades under long-term drift remains to be explored.

Related Work & Insights

• vs Slot Attention / BO-QSA: BO-QSA enriches query distribution via multiple Gaussians, but queries remain "dataset-level" priors, lacking "sample-level" information; SmoothSA directly injects sample features into queries via self-distillation, representing a fundamental path difference.
• vs MetaSlot: MetaSlot performs a draft aggregation, then reinitializes queries with object prototype codebooks; essentially still "augmenting queries with priors" rather than "with current sample," and requires extra codebooks and discretization, increasing complexity.
• vs STEVE / SAVi / RandSF.Q: These methods focus on "better query propagation across frames," but all assume identical SA iterations for all frames; SmoothSA is the first to make "inter-frame transformation strength" adjustable, orthogonal to query transition mechanisms, and can be stacked atop SOTA methods like RandSF.Q for further gains.

Rating

• Novelty: ⭐⭐⭐⭐ First to explicitly formalize the long-overlooked issues of "query cold start" and "transformation homogenization," though the solution (self-distillation + scheduling differentiation) is relatively simple.
• Experimental Thoroughness: ⭐⭐⭐⭐ Consistency validated across multiple image/video OCL backbones and downstream tasks, with complete ablations.
• Writing Quality: ⭐⭐⭐⭐ Clearly unifies both modifications under the "smooth iterations / smooth recurrences" concept.
• Value: ⭐⭐⭐⭐ Near-zero cost, can be attached to any SA-based OCL backbone for stable gains, high engineering value.

Related Papers

• [ICML 2026] Certified Robustness under Heterogeneous Perturbations via Hybrid Randomized Smoothing
• [ICML 2026] Large Vision-Language Models Get Lost in Attention
• [ICLR 2026] Directional Embedding Smoothing for Robust Vision Language Models
• [CVPR 2026] MASQuant: Modality-Aware Smoothing Quantization for Multimodal Large Language Models
• [ICLR 2026] Constructive Distortion: Improving MLLMs with Attention-Guided Image Warping