Self-Transformer (ST)

• Self-Transformer (ST) is an encoder adaptation that refines its self-attention via fixed-point iterations to achieve adaptive computational depth.
• It replaces standard multi-head self-attention with a Fixed-Point Self-Attention module that iteratively updates hidden states until convergence.
• Empirical results show ST outperforms traditional Transformers in language and vision tasks with improved accuracy and efficient adaptive compute.

The SELF-Transformer (Self-Enhancing Latent Feedback Transformer, or ST) is an architectural adaptation of the encoder Transformer that introduces a fixed-point iterative loop to refine its self-attention alignment matrix. Unlike the standard encoder where attention is computed in a single feed-forward pass, the ST updates its alignment weights via internal iteration, scaling test-time computation with input complexity. This approach enables input-adaptive computational depth, boosting expressive power without increasing the parameter count or relying on explicit token-level autoregression (Mathur et al., 17 Jul 2025).

1. Architectural Modifications and Computational Process

The ST replaces each vanilla multi-head self-attention (MHA) block with a Fixed-Point Self-Attention (FPSA) module. Each attention head maintains a hidden state $Z_k^{(i)}$ (initialized by the standard query projection), which is iteratively updated according to the following procedure:

1. Alignment computation:

$$T_k^{(i)} = \operatorname{softmax}\big( (Z_k^{(i)} W_Q^{(i)})(Z_k^{(i)} W_K^{(i)})^\top/\sqrt{d_h} \big) \in \mathbb{R}^{n \times n}$$

2. Hidden state update:

$$Z_{k+1}^{(i)} = T_k^{(i)} (X W_V^{(i)}) \in \mathbb{R}^{n \times (d/h)}$$

This loop continues until the fixed-point criterion is met (relative change in Frobenius norm falls below threshold $\epsilon$) or a maximum number of iterations $K_{max}$ is reached. At convergence, the concatenated results from all heads are projected as in a standard MHA block. All residual connections, normalization, feed-forward network (FFN), and embedding structures are retained as in vanilla Transformers. No additional parameters or hypernetworks are introduced; only the forward-backward pass in self-attention is modified.

2. Mathematical Formulation and Training Dynamics

Let $X \in \mathbb{R}^{n \times d}$ denote input embeddings, $h$ the number of heads, and $d_h = d/h$ the head dimension. For head $i$, projection matrices $$W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} \in \mathbb{R}^{d \times d_h}$$ and output projection $W_O \in \mathbb{R}^{d \times d}$ are shared across iterations. The alignment computation and hidden state update can alternatively be written directly in terms of the alignment matrix $A_k^{(i)}$:

$$A_k^{(i)} = \operatorname{softmax}\left( (A_{k-1}^{(i)} V W_Q)(A_{k-1}^{(i)} V W_K)^\top/\sqrt{d_h} \right)$$

with $Z_k^{(i)} = A_k^{(i)} V$, where $V = X W_V^{(i)}$.

Convergence criterion:

The FPI stops when:

$$\frac{\| Z_{k+1}^{(i)} - Z_k^{(i)} \|_F}{\| Z_k^{(i)} \|_F} < \epsilon$$

or $k = K_{max}$. At this point, $Z^* = f_\theta(Z^*, X)$, where $f_\theta$ defines the local FPI operator.

For stability, spectral normalization (coefficient 1.0) is applied to $W_Q$/$W_K$ projections, enforcing contractivity as per the theoretical result that softmax attention is contractive under the Wasserstein-1 metric (provided the product of their spectral norms is less than 1). Gradients are clipped by global norm across all parameters to ensure stability under many inner-loop steps.

The outer objective (e.g., cross-entropy) and overall optimization procedures are unchanged relative to standard Transformer training.

3. Relationship to Expressive Power and Circuit Complexity

Standard encoder Transformers, regardless of width and head count, computed in a single feed-forward pass, are provably limited to the constant-depth circuit class $TC^0$ [Hahn et al., 2022]. In contrast, autoregressive Transformers (e.g., LLMs with token-wise generation) and reasoning via externally decoded chains of thought can simulate $O(n)$-depth circuits, breaching the $TC^0$ barrier.

The ST raises the expressive ceiling for encoder architectures by effectively achieving unbounded “latent depth” through its iterative attention refinement. Allowing $T(k)$ FPI iterations per input, an ST layer can simulate a circuit of depth $O(T(k))$, with $T(k)$ being adaptively chosen — thus computational cost grows only on difficult inputs. This formulation aligns the ST’s expressivity with that of deep equilibrium models or autoregressive computation, while retaining a feed-forward encoder interface (Mathur et al., 17 Jul 2025).

4. Empirical Results and Benchmarking

Extensive benchmarks demonstrate the empirical advantage of the ST over baselines, both in language and vision domains.

Model	GLUE Avg (%)	SQuAD v1.1 F1	SQuAD v2.0 F1	WikiText-2 PPL
BERT-Base (110M)	78.3	88.6	73.6	61.1
RoBERTa-Base (125M)	82.1	90.2	80.5	46.2
ELECTRA-Base (110M)	85.0	90.7	81.7	–
SELF-Transformer	88.4	95.2	88.7	28.5

Across 20+ benchmarks, the ST achieved:

• Up to 20% higher accuracy in language (GLUE, SQuAD, WikiText-2)
• +1.7% top-1 and +0.7% top-5 ImageNet accuracy (SELF-ViT vs ViT-B/16, with fewer parameters)
• +4.6 pp mAP@50 in object detection (COCO), with faster inference (87 ms vs 150 ms)
• +0.2–0.4 dB PSNR in image restoration tasks and +0.005 SSIM
• Up to +5.3 pp improvement on image–text retrieval and +3.9 pp on VQA v2.0 tasks

In ablations, the ST closed much of the reasoning gap to chain-of-thought and induction head tasks: 91.1% vs 63.1% accuracy (induction task), with only modest extra compute.

5. Computational Efficiency and Adaptivity

The average extra compute per ST layer is approximately 5× that of vanilla self-attention (reflecting a typical $\sim$5 FPI steps), but because lower layers and “easy” inputs converge quickly, end-to-end model overhead is only $\sim$1.5× baseline, with batch throughput decreasing by approximately 30%. In vision, sublinear scaling of test-time FLOPs was observed, with speedups for longer sequences and larger images (e.g., up to 1.5× acceleration for long inputs in SELF-ViT).

Most inputs converge in fewer than $K_{max}$ iterations (median: 5, max: 22 in toy ViT MNIST, with $K_{max}=100$). Setting $K_{max}=5$ preserves 90% of ST’s gain with <1% performance degradation, while tighter convergence tolerances ($\epsilon \leq 10^{-4}$) yield best performance.

No explicit ponder cost is included, and computational budgets are currently imposed via the $K_{max}$ parameter or, optionally, a separate routing network to adapt iteration counts per sample.

6. Limitations and Theoretical Considerations

The worst-case computational cost remains unbounded up to $K_{max}$ iterations, which could pose challenges for applications with strict real-time constraints. The ST’s convergence guarantees rely on contractive properties of the softmax attention operator; these may not universally hold under atypical weight initializations, though failures are rare with proper spectral normalization.

There is no learnable ponder cost or explicit return of computational effort to the end user; practical deployments must manually enforce computational ceilings. The underlying theoretical characterization of which function classes are computable via FPI-refined attention versus deep equilibrium or autoregressive methods remains open for further investigation (Mathur et al., 17 Jul 2025).

7. Prospects and Potential Extensions

ST architectures suggest several future research directions:

• Scaling to LLMs by optimization of FPI solvers and hardware acceleration for iterative attention
• Integration of learnable ponder cost or controller modules for dynamic accuracy/compute tradeoff
• Formal characterization of the computational expressivity classes enabled by latent attention refinement, both practically and in circuit-complexity terms
• Development of hybrid reasoning architectures that combine internal latent iteration with externalized chain-of-thought processes
• Extension of FPI-based refinement to cross-attention blocks and multimodal architectures

The ST demonstrates that internal latent state refinement—without explicit externalization of intermediate reasoning steps—can bridge much of the gap in expressive power between encoder Transformers and autoregressive/DEQ models, while preserving architectural simplicity and efficiency (Mathur et al., 17 Jul 2025).