Attention Broadcast Decoder for Latent Dynamics

• ABCD is a plug-and-play module for autoencoder latent dynamics that generates pixel-accurate attention maps to localize each latent dimension’s influence.
• It replaces the standard convolutional decoder with a composite architecture (encoder, attention processor, broadcast decoder) to reconstruct images and filter static backgrounds.
• ABCD improves multi-step prediction accuracy and provides physically interpretable latent representations, benefiting control tasks in soft robotics.

The Attention Broadcast Decoder (ABCD) is a plug-and-play module designed for autoencoder-based latent dynamics learning. Its core contribution is generating pixel-accurate attention maps that spatially localize each latent dimension’s influence while filtering static backgrounds. Introduced in the context of learning physically interpretable models from video for soft continuum robots, ABCD enables direct on-image visualization of learned dynamics—such as masses, stiffness, and forces—without requiring prior knowledge or manual annotation. ABCD achieves significant multi-step prediction accuracy improvements for both Koopman operator and oscillator network models and enables smooth latent space extrapolation beyond training data (Krauss et al., 23 Nov 2025).

1. Architecture and Components

ABCD replaces the standard convolutional decoder of a β-VAE autoencoder with a composite architecture comprised of three principal parts:

Encoder (φ):

The encoder takes an input image $o \in \mathbb{R}^{C \times H \times W}$ and employs a series of three convolutional layers (Conv2D 4×4, stride 2, LeakyReLU activations) with output channels [C → 32 → 64 → 128]. A subsequent linear layer computes the latent mean $\mu \in \mathbb{R}^k$ and log-variance $\log \sigma^2 \in \mathbb{R}^k$. Latents $z \sim \mathcal{N}(\mu,\, \sigma^2)$ are then sampled for use by downstream modules.

Attention Map Generator ("attention processor"):

Given latent vector $z \in \mathbb{R}^k$ and spatial coordinate channels $c_{xy} \in [-1,1]^2$ (per pixel), a small fully connected network or a stack of $1 \times 1$ convolutions computes per-latent attention logits $\ell_j(x, y) = f_{\text{att}}([z, c_{xy}])_j$ for $j = 1, \ldots, k$. A fixed background logit $b$ is appended, and a softmax (optionally Gumbel-annealed) over $j=1,\ldots,k+1$ yields attention maps $a_j(x, y)$ and $a_b(x, y)$ for latents and background, respectively.

Broadcast Decoder:

Each latent scalar $z_j$ is mapped to a per-pixel representation $$\bar{z}_j(x, y) = W_j z_j + b_j \in \mathbb{R}^{n_f}$$ (broadcast in spatial dimensions). A learnable static-background feature map $\bar{z}_b \in \mathbb{R}^{n_f \times H \times W}$ handles immobile scene content. Attended features are summed pixel-wise:

$$\tilde{Z}(x, y) = \sum_{j=1}^k a_j(x, y)\, \bar{z}_j(x, y) + a_b(x, y)\, \bar{z}_b(x, y)$$

Concatenating coordinates, the result passes through four $1\times1$ convolutional layers to the final output channels.

The architecture’s essential intuition assigns “ownership” of image regions to specific latents, with background handled purely by the background map.

2. Mathematical Formulation of Attention

For time index $i$ and spatial location $(x, y)$ (with normalized $c_{xy}$):

• Latent-Conditioned Attention Logits:

$$\ell_j^{(i)}(x, y) = f_{\text{att}}([z^{(i)}, c_{xy}])_j \quad (j=1,\ldots,k)$$

The background logit $\ell_{k+1}^{(i)}(x, y) \equiv b$ (fixed).

• Softmax Attention:

$$a_j^{(i)}(x, y) = \frac{\exp(\ell_j^{(i)}(x, y))}{\sum_{m=1}^{k+1} \exp(\ell_m^{(i)}(x, y))}$$

for $j=1,\ldots,k$; $a_b$ for background.

• Latent Broadcast Expansion:

$$\bar{z}_j^{(i)}(x, y) = W_j z_j^{(i)} + b_j \in \mathbb{R}^{n_f}$$

• Feature Map Synthesis:

$$\tilde{Z}^{(i)}(x,y) = \sum_{j=1}^k a_j^{(i)}(x, y) \, \bar{z}_j^{(i)}(x, y) + a_b^{(i)}(x, y) \, \bar{z}_b(x, y)$$

• Image Generation:

$\tilde{Z}^{(i)}(x, y)$

concatenated with $c_{xy}$ is mapped by four $1\times1$ convolutions to reconstruct the image $\hat{o}^{(i)}$.

• Attention-Consistency Regularization:

The loss penalizes changes in attention where image content is static:

$$\mathcal{L}_{\text{attn-cons}} = \frac{1}{k h w} \sum_{j, x, y} |\partial_t a_j(x, y)| \cdot [1 - |\partial_t \hat{o}(x, y)|]$$

3. Integration with Autoencoder and Loss Functions

The ABCD is integrated into autoencoder-based models for video-based dynamics learning, employing the following training process at each time step $i$:

1. Encode observation $o^{(i)}$ to obtain $\mu^{(i)}$, $\sigma^{(i)}$ and sample $z^{(i)}$.
2. Calculate latent velocity $\dot{z}^{(i)}$ via forward-mode automatic differentiation on the encoder, using finite-differenced image velocities.
3. Image is reconstructed using ABCD, producing $\hat{o}^{(i)}$.
4. For prediction:
   • Koopman operator: Stack $\xi = [z;\ \dot{z}] \in \mathbb{R}^{2k}$, update via $\xi^{(i+1)} = A \xi^{(i)} + B(u^{(i)})$.
   • Oscillator network: Step symplectic Euler on the system $M\ddot{z} + D\dot{z} + K(z - z_0) = B(u)$ to obtain $z^{(i+1)}, \dot{z}^{(i+1)}$.
5. Predicted $z^{(i+1)}$ is decoded by ABCD.
6. The batch loss is:

   $$\mathcal{L}_{\text{basic}} = \frac{1}{N} \sum_{i=1}^N [ \textrm{MSE}(\hat{o}^{(i)}, o^{(i)}) + \lambda_d \textrm{MSE}(\hat{o}^{(i+1)}, o^{(i+1)}) + \beta\,\textrm{KL}[\mathcal{N}(\mu, \sigma)\Vert \mathcal{N}(0,I)] + \lambda_z(\textrm{MSE}(z^{(i+1)}, \hat{z}^{(i+1)}) + \textrm{MSE}(\Delta t\dot{z}^{(i+1)}, \Delta t\hat{\dot{z}}^{(i+1)}) ] + \lambda_{\text{attn}} \mathcal{L}_{\text{attn-cons}}$$

   with optional additional terms for oscillator cases: $$\lambda_{\text{attn-coup}}\mathcal{L}_{\text{attn-coupling}}$$ and $\lambda_s \mathcal{L}_{\text{steady-state}}$.

This setup supports both single- and multi-step prediction, and the compact latent oscillator representation discovered autonomously is interpretable and suitable for control tasks.

4. Coupling to 2D Oscillator Networks

A key feature of ABCD is its ability to link each latent subspace to a physically interpretable 2D oscillator. The latent dimensions are grouped (for $k$ even) into $n = k/2$ paired oscillators, where for oscillator $l = 1, ..., n$

$$q_l = [z_{2l-1}, z_{2l}]^\top,\quad \dot{q}_l = [\dot{z}_{2l-1}, \dot{z}_{2l}]^\top$$

The mass matrix $M$ is constrained so that $M_{2l-1, 2l-1} = M_{2l, 2l}$ for each oscillator, ensuring equal mass per 2D oscillator.

The ABCD attention processor emits one map $a_l(x, y)$ per oscillator. The center-of-mass (COM) position in image space for oscillator $l$ is:

$$p_l = \frac{\sum_{x, y} [a_l(x, y)]^2 c_{xy}}{\sum_{x, y} [a_l(x, y)]^2}$$

with velocities $\dot{p}_l$ computed via auto-diff through the softmax.

An attention-coupling loss enforces that relative motions between latent and image-space oscillator COMs are consistent:

$$\mathcal{L}_{\text{attn-coupling}} = \mathbb{E}_{l \neq m} \left[ \left( \frac{\dot{d}_{lm}^{\text{lat}}}{\bar{v}_{lm}^{\text{lat}}} - \frac{\dot{d}_{lm}^{\text{img}}}{\bar{v}_{lm}^{\text{img}}} \right)^2 \right]$$

where $d_{lm}$ are signed rates of change of distances, and $\bar{v}_{lm}$ the mean speeds.

This coupling provides direct interpretability of dynamically learned parameters (mass, stiffness, damping, input) projected back onto observable image coordinates.

5. Training Protocols and Hyperparameters

ABCD is typically trained on subsampled video at $3 \times 32 \times 32$ resolution with $\Delta t = 1/60$ s, using batch size 32 and latent dimension $k = 8$ (single-segment) or $k = 10$ (two-segment). The feature width per latent $n_f$ is in the range 32–64. The background logit $b=1.0$ is fixed. Typical loss weights are $\lambda_d=1.0$, $\beta=4.0$, $\lambda_z=0.1$, $\lambda_{\text{attn-cons}}=0.5$, $\lambda_{\text{attn-coupling}}=0.1$, $\lambda_s=1.0$.

Optimization is performed with AdamW using separate learning rates (e.g., $1 \times 10^{-4}$ for encoder/decoder, $1 \times 10^{-3}$ for dynamics). Training involves up to 300 epochs, with a 5-epoch “warmup” at reduced learning rates to allow background feature absorption, and early stopping at ~100–150 epochs when validation loss plateaus. The attention softmax can employ Gumbel noise with annealed temperature to produce sharper maps.

Reproduction of results requires the exact implementation of spatial broadcasting, attention, coupling losses, and symplectic integration for the oscillator networks, as detailed in the published code repository.

6. Empirical Performance and Applications

Evaluations demonstrate ABCD-based models achieve significant improvements in multi-step prediction for soft continuum robots. On the two-segment setting, error reductions were observed of 5.7× for Koopman operators and 3.5× for oscillator networks relative to standard approaches. The learned oscillator networks recovered chain structures corresponding to physical robot segments without supervision. Standard methods lacked the ABCD capability for smooth latent space extrapolation beyond the training data and did not offer physical interpretability of latent variables.

A plausible implication is that ABCD provides a data-driven path to compact, physically grounded control-oriented representations for robots and other dynamical systems with complex, high-dimensional observations, whereas classical pipelines typically require task-specific priors or significant manual engineering (Krauss et al., 23 Nov 2025).