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:

$H=495, W=436, C_{in}=115$0

with ReLU activation, layer norm, and groupwise spatial mixing ($H=495, W=436, C_{in}=115$1, $H=495, W=436, C_{in}=115$2). 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 $H=495, W=436, C_{in}=115$3 tokens that are linearly projected and positionally encoded. These are propagated through $H=495, W=436, C_{in}=115$4 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.,

$H=495, W=436, C_{in}=115$5

with $H=495, W=436, C_{in}=115$6, 'same' padding, stride 1, and ReLU (or ELU) activations as specified.

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

$H=495, W=436, C_{in}=115$7

with group outputs concatenated. An explicit orthogonality regularizer $H=495, W=436, C_{in}=115$8 is applied where $H=495, W=436, C_{in}=115$9 ($(512,512)$0 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:

$(512,512)$1

$(512,512)$2

where $(512,512)$3 employs ELU activations, and the transformer operates on raw image patches projected into $(512,512)$4-dimensional embeddings.

5. Training Strategies and Regularization

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

$(512,512)$5

using Adam (lr $(512,512)$6, 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

$(512,512)$7

with $(512,512)$8. SGD with learning rate scheduling is used, and data split into 80%/10%/10% for training/validation/test.

TUnet employs AdamW (lr $(512,512)$9, weight decay $2\times2$0, 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 $2\times2$1, surpassing single-model results.
• Micron-Sized Fiber Single-Shot Reconstruction (Xu et al., 2 Feb 2026): On $2\times2$2 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 $2\times2$3m and capillaries at 10–20 $2\times2$4m, 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 $2\times2$5 pixels, the quadratic cost per channel ($2\times2$6) of unfactored MLPs is reduced to $2\times2$7 via grouping. Actual inference cost is $2\times2$8 GFLOPs with 172.7M parameters; inference time is $2\times2$9 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 $2\times2$0 resolution (with patch size $2\times2$1, $2\times2$2 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 = $2\times2$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.