HUTFormer: Hierarchical U-Net Transformer for Long-Term Traffic Forecasting

In plain words

Traffic forecasts are most useful when they look a full day ahead, so a driver can plan a route. Most published traffic models only look a short time ahead, because the graph-based machinery they use becomes too expensive on long histories. This paper asks what it takes to forecast a whole day.

The authors' answer is HUTFormer, a model shaped like a U-Net, the layered design used in image segmentation. The encoder looks at the history at several zoom levels: fine detail in the first layers, broader patterns in the later ones. The decoder then goes back down the stack and uses those zoom levels to sharpen a first, rough forecast. To keep the cost down, the history is first chopped into short chunks of time steps, and each sensor is given learned labels for its position and for the time of day and week.

TL;DR

HUTFormer is a Transformer for long-term traffic forecasting: one day of 5-minute steps for all sensors, with a U-Net-style encoder that builds multi-scale features and a decoder that refines a coarse forecast. Training is two-stage: the encoder is trained first, then frozen while the decoder is trained. The authors report the lowest MAE in every horizon of their speed and flow tables, against six traffic models and seven long-sequence models (PatchTST included).

Problem

Graph neural networks for traffic forecasting model correlations between sensors through graph convolution, which is costly when the input is long. Most such models therefore predict about 1 hour (12 steps) ahead §1. The authors study predicting 1 day ahead, which they define as 288 future steps from 288 past steps of 5-minute data §3 §5.1. They argue two problems are specific to this horizon: the signal has regular daily cycles plus sharp, short congestion peaks, so the model needs features at several scales, and the input and output windows are not aligned in time, so the U-Net skip connections of vision models cannot be added directly §1 §4.1.

Method

Overview. HUTFormer has a hierarchical encoder, an intermediate forecast made from the top encoder layer, and a hierarchical decoder that refines it. Training runs in two stages §4.1.

Input embedding §4.2:

1. Divide the history of sensor i, X^i in R^(T x C), into P non-overlapping segments of length L, so T = P * L. Embed each segment with one linear map shared by all segments and sensors §4.2 eq. 1.
2. Add spatial-temporal positional encoding. Concatenate each segment embedding with a learned per-sensor vector E (size N x d1), a learned time-of-day vector (N_D slots) and a learned day-of-week vector (N_W = 7 slots). Then apply one linear layer §4.2 eq. 2. The paper presents this as a replacement for standard positional encoding and as a cheaper substitute for graph convolution §4.2 §5.5.

Hierarchical encoder §4.3:

3. Window Transformer layer. Standard pre-norm Transformer layer, except that self-attention runs only inside non-overlapping windows of size 3 tokens, so the receptive field stays small §4.3 eq. 3 §5.1.
4. Segment merging. Concatenate each pair of neighbouring tokens. This halves the token count and doubles the width §4.3 eq. 4 fig. 4.
5. Repeat window layer and merging to get a stack of scales. The model depth is set to 4 §5.1; the text does not say whether this counts encoder blocks, decoder blocks or both.
6. Intermediate forecast. Flatten all top-layer tokens and apply one linear layer to get the full T_f x C forecast for each sensor §4.3 eq. 5. Stage 1 minimises mean absolute error (MAE) on this forecast §4.3 eq. 6.

Hierarchical decoder §4.4:

7. Cross-scale attention. Queries come from the decoder tokens; keys and values come from the encoder's multi-scale features. The encoder features pass through a linear layer to match widths. The scores are scaled by the square root of the decoder width §4.4 eq. 7.
8. Cross-scale Transformer layer. Same residual and layer-norm structure as the window layer, with multi-head cross-scale attention in place of self-attention, followed by an MLP §4.4 eq. 8.
9. Final forecast. One linear layer, shared across all decoder segments, maps each decoder token to its L x C slice of the forecast §4.4 eq. 9. Stage 2 minimises MAE on the full forecast §4.4 eq. 10. The encoder's parameters are fixed during this stage §4.4 §4.1.

# HUTFormer forward pass, rebuilt from the text (the paper has no algorithm box)
# Input:  history (T, N, C) with T = 288; time-of-day and day-of-week indices
# Output: forecast (T_f, N, C) with T_f = 288

def hutformer(history, time_of_day, day_of_week):
    # Cut each sensor's series into P non-overlapping segments of length L
    # [§4.2] [eq. 1]
    segments = segment(history, length=12)       # (N, P=24, L*C)
    tokens = linear_embed(segments)              # (N, P, d), d = 32

    # Concatenate segment, spatial and temporal embeddings, then project
    # [§4.2] [eq. 2]
    tokens = project(concat(tokens, spatial_emb, tid_emb, diw_emb))

    # Hierarchical encoder: window attention, then merge neighbouring tokens
    # [§4.3]
    encoder_features = []
    for block in encoder_blocks:             # number of blocks not stated
        # Self-attention inside non-overlapping windows of size 3 [§4.3] [eq. 3]
        tokens = window_transformer(tokens, window_size=3)
        # Which blocks feed the decoder is not stated in the paper
        encoder_features.append(tokens)
        # Concatenate pairs of tokens: half as many tokens, double width
        # [§4.3] [eq. 4]
        tokens = segment_merging(tokens)

    # Coarse forecast from the top-layer tokens, one linear map [§4.3] [eq. 5]
    coarse = reshape(linear_head(flatten(tokens)), (N, T_f, C))

    # Decoder tokens built from the coarse forecast: not stated in the paper
    dec = embed_segments(coarse)
    for block in decoder_blocks:
        # Queries from decoder tokens, keys and values from encoder features
        # [§4.4] [eq. 7] [eq. 8]
        dec = cross_scale_transformer(dec, encoder_features)

    # One linear map per segment, shared across segments [§4.4] [eq. 9]
    forecast = reshape(linear_out(dec), (N, T_f, C))
    return forecast, coarse

# Two-stage training  [§4.1] [§4.3] [§4.4]
# Stage 1 trains the encoder and the coarse head; stage 2 trains the decoder

def train_stage_one(history, target):
    # Mean absolute error on the intermediate forecast [§4.3] [eq. 6]
    _, coarse = hutformer(history, time_of_day, day_of_week)
    return mean(abs(coarse - target))

def train_stage_two(history, target):
    # Encoder parameters are frozen here [§4.4]
    forecast, _ = hutformer(history, time_of_day, day_of_week)
    # Mean absolute error on the final forecast [§4.4] [eq. 10]
    return mean(abs(forecast - target))

Spatial adjacency is not used by HUTFormer itself. Graph-based baselines use two adjacency definitions, a thresholded Gaussian kernel on road distance for METR-LA and PEMS-BAY §5.1 eq. 11, and a binary 3.5-mile threshold for PEMS04 and PEMS08 §5.1 eq. 12.

Hyperparameters

Results

Author-reported. Setup: four traffic datasets tab. 1; MAE and MAPE at horizons 12, 48, 96, 144, 192 and 288 steps (1 to 24 hours) §5.2; six traffic baselines (DCRNN, Graph WaveNet, MTGNN, STID, STEP, D²STGNN) and six long-sequence baselines (HI, DLinear, Informer, Autoformer, FEDformer, Pyraformer), plus PatchTST in the tables §5.1 §5.2 tab. 3 tab. 4. Crossformer ran out of memory and is omitted §5.2. Each baseline is a single number per cell; the paper reports no run-to-run variance.

Headline cells, taken from the tables. "Best baseline" is the lowest MAE among the baselines in that table, chosen by us; the asterisk marks baselines run with segment embedding §5.1.

Ablation on METR-LA, horizon 288 only. Each row changes one part of the full model §5.5 tab. 6:

Other results: on ETTh1, ETTm1 and Weather (non-traffic), HUTFormer reports the best MAE and MSE in most cells against seven Transformer baselines tab. 5 §5.4. An efficiency study on METR-LA compares GPU memory and running time of the ablation variants §5.3 fig. 6. The text gives the qualitative ranking only: removing segment embedding raises cost most, and spatial-temporal encoding, unlike graph convolution, does not add much complexity §5.3.

Paper vs code

No official code.

Search, in order: the LaTeX source and arXiv metadata contain no code link or project page [raw/src/2307.14596]. A GitHub repository search for "HUTFormer" returned one repository, owned by an account that is not an author. Its README calls it a rough reimplementation that falls short of the paper's accuracy, so it is not used here. A second search restricted to one author's GitHub account returned nothing; that account was not verified as the authors' own. The paper trains all models in the BasicTS pipeline §5.1, but that framework is not HUTFormer's code. Without the official code, the open points below cannot be checked against an implementation.

Pitfalls

• Encoder depth is ambiguous. The text gives "depth 4" §5.1 but writes the encoder as "S blocks" without a value §4.3. It is not stated whether the four layers are encoder blocks, decoder blocks, or both. The authors also say the depth may change with the segment length so the receptive field covers the whole sequence §5.6.
• Token counts under segment merging. Starting from P = 24 tokens, repeated halving gives 12, 6 and 3. The window size of 3 divides each of these evenly, but a merge of 3 tokens is not defined in the paper. Any fourth merge needs padding or a rule the paper does not state.
• Equation typos. The width update reads d^(l+1) = 2d^(l+1) §4.3. Read as d^(l+1) = 2d^l, which matches the prose ("doubles the hidden dimension"). The token update reads P^(l+1) = P^l / 2 §4.3.
• Which encoder features the decoder uses. Keys and values are "the multi-scale features from the encoder" §4.4, but the text does not say which blocks are used, how many scales, or how they are combined.
• Decoder input is not defined. The paper says the decoder fine-tunes the intermediate forecast §4.4 but does not say how that forecast becomes the decoder's input tokens (segmentation, embedding, or raw values).
• Attention details. Equation 8 applies layer norm to a pair of inputs without saying which one is normalised §4.4 eq. 8. The head count and MLP width are not stated.
• Normalisation and missing data. The paper does not say whether inputs are scaled, how zero or missing sensor readings are handled, or how MAPE is computed when a true value is zero. MAPE on flow data is sensitive to zeros §5.2 tab. 4.
• Time indices. The mapping from a timestamp to the time-of-day index (N_D = 288 at 5-minute sampling §4.2) and day-of-week index is not stated, and neither is the alignment of the forecast's first step.
• Two stages. The epochs for each stage and whether stage 2 uses the same schedule are not stated §5.1. The "end2end" and "w/o fix" ablations differ from the full model by 0.07 MAE or less at horizon 288 tab. 6, and no variance is reported, so these gaps may be within run-to-run noise.
• Baselines were modified. The STGNN baselines were changed to use segment embedding, with the same segment and stride size as HUTFormer §5.1. Their results therefore reflect the authors' modified versions, not the unmodified published models. Long-sequence baselines use default settings §5.1.
• Hardware differs. Training ran on RTX 3090 §5.1; the cost comparison ran on a V100 32 GB §5.3. Timing numbers are not comparable across the two.
• New sensors. Spatial embeddings must be trained for each sensor. The authors say HUTFormer cannot naturally handle new nodes without retraining §5.8.
• Long horizons. The authors note that uncertainty grows with the forecast period and do not quantify it §5.8.

Lineage

• Builds on: Swin Transformer for window self-attention and staged token merging in vision §4.3; U-Net and Swin-UNet for the U-shaped design §4.1; STID for the spatial-temporal identity idea behind the positional encoding §4.2; STEP for segment embedding §4.2.
• Compared against: STGNN traffic models (DCRNN, Graph WaveNet, MTGNN, STID, STEP, D²STGNN) §5.1 tab. 3 tab. 4; long-sequence Transformers and linear models (Informer, Autoformer, FEDformer, Pyraformer, PatchTST, DLinear, HI) §5.1 tab. 3 tab. 4; Crossformer and Triformer on the non-traffic datasets §5.4 tab. 5.
• Superseded by: none recorded in this wiki.