TF-UNet: Advanced Encoder–Decoder Networks

• TF-UNet architecture is a family of encoder–decoder networks that integrate innovations like dense connectivity, grouped-MLP skip fusion, and transformer modules to enhance feature extraction.
• TF-UNet variants have been applied in traffic forecasting, optical imaging, and medical image segmentation, achieving improved performance metrics and robust state-of-the-art results.
• Methodological advancements in TF-UNet, such as parallel pooling and transformer-based bottleneck integration, enable effective multi-scale feature fusion and capture of non-local dependencies.

TF-UNet defines a group of encoder–decoder neural architectures derived from the U-Net backbone, characterized by innovations that range from enhanced dense connectivity and ensemble learning to advanced skip-connection fusion and transformer integration. The term "TF-UNet" has distinct instantiations in recent literature, most notably for traffic map prediction in the Traffic4cast challenge (Choi, 2020), for resolving optical speckle patterns in single-shot imaging through tapered fibers (Xu et al., 2 Feb 2026), and (as TUnet) for combining transformer modules with U-Net for medical image segmentation (Sha et al., 2021). These models share a common encoder–decoder topology yet diverge substantially in how they fuse multi-scale spatial context, manage non-local dependencies, and integrate domain priors.

1. Topological Overview and Architectural Motifs

Across its variants, TF-UNet retains the U-shaped encoder–decoder backbone: an input tensor is progressively contracted along the encoder path to a bottleneck, with symmetric expansion in the decoder path, and skip connections bridging encoder and decoder at corresponding spatial resolutions. In the Traffic4cast variant (Choi, 2020), the model receives $X_0 \in \mathbb{R}^{H \times W \times C_{in}}$, with $H=495, W=436, C_{in}=115$, propagates inputs through eight stages of dense blocks with average or max pooling, and decodes via transposed convolutions, culminating with a 1×1 convolution to project to future traffic states.

In the tapered-fiber imaging application (Xu et al., 2 Feb 2026), the input is $(512,512)$ speckle images, and the architecture comprises four encoding and four decoding stages. Each encoder block applies "double-conv" (two Conv2D-BN-ReLU layers), progressing through channel widths of 64, 128, 256, and 512. MaxPooling2D ($2\times2$) halves spatial dimensions at each stage, mirrored by $2\times2$ transposed convolutions in the decoder. Skip connections are augmented via grouped-MLP fusion rather than raw concatenation, introducing non-local spatial and channel mixing.

In the TF-UNet "TUnet" (Transformer-UNet) hybrid (Sha et al., 2021), the architecture fuses a four-level CNN encoder–decoder with a parallel pathway: the non-overlapping, linearly embedded raw image patches are processed by a stack of transformer layers, the output of which seeds the decoder bottleneck, while skip connections propagate only CNN features.

2. Encoder and Decoder Design Variants

TF-UNet encompasses several key encoder–decoder implementations:

Traffic4cast Models (Choi, 2020):

• Model 1: Employs eight dense blocks (each with four $3\times3$ convolutions) and average pooling for spatial reduction. Output channels increase per stage (64, 96, 128), fixing at 128 beyond the third block. The bottleneck is a $4\times4\times128$ tensor. Decoder uses deconvolutions and single standard convolutions at each expansion step.
• Model 2: Integrates parallel max-pooling and dense convolutions within each block; their outputs are concatenated and spatially downsampled by a $3\times3$, stride-2 convolution. This hybrid pooling enables the model to ingest heterogeneous compression artifacts.
• Model 3: Applies max-pooling before the dense block in the encoder and combines both deconvolution and bilinear interpolation for upsampling in the decoder.

Grouped-MLP Skip Fusion (Optical Imaging) (Xu et al., 2 Feb 2026):

Encoder blocks remain standard, but the skip connection at each level preprocesses features via grouped-MLP blocks. Each feature tensor $X\in\mathbb{R}^{C\times H\times W}$ is partitioned into $G$ groups along the channel axis, spatially flattened, and processed by a two-layer MLP per group:

$$\hat X^{(g)} = W_2^{(g)} \sigma( W_1^{(g)}\, \mathrm{LN}(X^{(g)}) )$$

with ReLU activation, layer norm, and groupwise spatial mixing ($W_1^{(g)},W_2^{(g)}\in\mathbb{R}^{N\times N}$, $N=H\times W$). Decoding reverts the contraction pathway, concatenating upsampled decoder outputs with MLP-fused skip features.

Transformer Integration (TUnet) (Sha et al., 2021):

A non-overlapping patch embedding is performed on raw images, yielding $N$ tokens that are linearly projected and positionally encoded. These are propagated through $m$ transformer blocks using multi-head self-attention and layer norm, then reshaped as a low-resolution feature map injected at the decoder bottleneck. The decoder path mirrors U-Net, with upsampling and concatenation at each level.

3. Skip Connections and Feature Fusion

TF-UNet models depart from classical skip connection designs by incorporating advanced fusion mechanisms:

• DenseNet-style Concatenation (Choi, 2020): Skip paths propagate the output of each encoder dense block directly to its decoder counterpart via channel-wise concatenation. This preserves both fine and hierarchical spatial context, particularly when dense connectivity is employed.
• Grouped-MLP Skip Bridges (Xu et al., 2 Feb 2026): Rather than plain concatenation, each skip feature is transformed to facilitate non-local spatial relationships typical of tapered-fiber speckle distortions. This enables the network to correct for non-stationary, physically induced artifacts beyond the reach of conventional spatially local convolutions.
• Transformer Feature Injection (Sha et al., 2021): Global context is not fused per se through skip connections; instead, transformer-encoded features replace the bottleneck decoder input, while skip links transport only CNN-based hierarchical features.

4. Mathematical Formulations and Key Operations

The convolutional and pooling operations used throughout TF-UNet adhere to standard forms, e.g.,

$$F^\ell_{i,j,k} = \sum_{p=1}^{K_h} \sum_{q=1}^{K_w} \sum_{c=1}^{C_{in}} W^\ell_{p,q,c,k} F^{\ell - 1}_{i + p - \lfloor K_h/2 \rfloor,\,j + q - \lfloor K_w/2 \rfloor,\,c} + b^\ell_k$$

with $K_h=K_w=3$, 'same' padding, stride 1, and ReLU (or ELU) activations as specified.

Grouped-MLP skip blocks in (Xu et al., 2 Feb 2026):

$$X^{(g)} \to \mathrm{LN} \to \mathrm{ReLU} \to W_1^{(g)} \to \mathrm{ReLU} \to W_2^{(g)}$$

with group outputs concatenated. An explicit orthogonality regularizer $L_{ortho} = \lVert G - I_C \rVert_F^2$ is applied where $G = \frac{1}{N} FF^T$ ($F$ is the mean-centered bottleneck feature), encouraging channel decorrelation analogous to mode disentanglement in the tapered fiber.

In the TUnet (Transformer-UNet) architecture (Sha et al., 2021), multi-head self-attention and MLP operations follow the ViT-style Pre-LN format:

$$z'_l = z_{l-1} + \mathrm{MHA}(\mathrm{LN}(z_{l-1}))$$

$z_l = z'_l + \mathrm{MLP}(\mathrm{LN}(z'_l))$

where $\mathrm{MLP}$ employs ELU activations, and the transformer operates on raw image patches projected into $d$-dimensional embeddings.

5. Training Strategies and Regularization

Traffic4cast TF-UNet models are trained with mean squared error loss,

$$L(\hat y,x) = \|\hat y(x) - y_{\mathrm{true}}(x)\|_2^2$$

using Adam (lr $3 \times 10^{-4}$, hand-tuned decay on plateau), with input normalization and channel-wise static/dynamic feature concatenation. Ensemble learning is employed: predictions from independently trained instances of Models 1, 2, and 3 are aggregated via averaging or median selection, achieving best results with averaged ensembles of six models.

In the optical imaging variant (Xu et al., 2 Feb 2026), the total loss is

$L_{total} = L_{MSE} + \lambda L_{ortho}$

with $\lambda=2$. SGD with learning rate scheduling is used, and data split into 80%/10%/10% for training/validation/test.

TUnet employs AdamW (lr $10^{-3}$, weight decay $10^{-6}$, decay at epochs 60 and 100), with a batch size and training details adapted to the CT82 dataset for pancreas segmentation.

Notably, batch-norm and explicit dropout or weight decay (beyond optimizer defaults) are absent in the Traffic4cast variant (Choi, 2020), while batch norm is present in the double-conv layers for (Xu et al., 2 Feb 2026).

6. Application Domains and Performance Metrics

TF-UNet architectures are applied across diverse scientific and industrial contexts:

• Traffic Map Prediction (Traffic4cast) (Choi, 2020): TF-UNet achieves state-of-the-art performance in future urban traffic prediction, with ensemble MSE as low as $1.1628615 \times 10^{-3}$, surpassing single-model results.
• Micron-Sized Fiber Single-Shot Reconstruction (Xu et al., 2 Feb 2026): On $512\times512$ speckle images, TF-UNet improves on standard U-Net in key metrics: PSNR (9.17 dB vs 8.98 dB), SSIM (0.24 vs 0.18), MS-SSIM (0.32 vs 0.22), LPIPS (0.64 vs 0.69, lower better), and Pearson correlation (0.50 vs 0.39). It achieves recovery of fine neuronal and vascular features with FWHM near 15 $\mu$m and capillaries at 10–20 $\mu$m, supporting downstream tasks such as ROI extraction and functional signal quantification.
• Medical Image Segmentation (TUnet) (Sha et al., 2021): The integration of raw-patch-based transformer modules with the U-Net bottleneck leads to segmentation improvements on the CT82 pancreas dataset compared to prior U-Net based algorithms.

7. Computational Complexity and Resource Considerations

TF-UNet's architectural innovations have direct implications for complexity:

• Grouped-MLP Fusion: For grouped-MLP with $N=512^2 \approx 2.62\times10^5$ pixels, the quadratic cost per channel ($O(C N^2)$) of unfactored MLPs is reduced to $O(G N^2)$ via grouping. Actual inference cost is $\sim2.76\times10^3$ GFLOPs with 172.7M parameters; inference time is $\sim1.5$ s on NVIDIA A100 (Xu et al., 2 Feb 2026).
• Dense/Hybrid Pooling: The use of dense blocks and parallel pooling in Traffic4cast TF-UNet increases model capacity, mitigated by ensembling and careful channel growth.
• Transformer Augmentation: TUnet's transformer bottleneck operates over patchified embeddings, ensuring tractable attention even at $512\times512$ resolution (with patch size $n=16$, $N=1024$ tokens), avoiding full image-resolution attention which would otherwise be prohibitive.

8. Comparative Summary Table

Variant/Domain	Key Innovation	Performance Metric
Traffic4cast (Choi, 2020)	Dense blocks, hybrid pooling, model ensembling	MSE = $1.162 \times 10^{-3}$
Optical Imaging (Xu et al., 2 Feb 2026)	Grouped-MLP skip fusion, physics-inspired loss	PSNR 9.17 dB; SSIM 0.24
TUnet (Sha et al., 2021)	Transformer raw patch bottleneck	Outperforms U-Net in CT82

These TF-UNet architectures demonstrate the adaptability of the U-Net encoder–decoder principle when augmented by advanced feature fusion, domain-driven priors, parallel pooling, and transformer-based global context, supporting state-of-the-art results across diverse scientific and engineering tasks.