🤖 machine learning

Robust Neural Tucker Factorization with Bias Correction and Adaptive Initialization

This paper proposes KaBiN, a robust neural Tucker factorization model for high-dimensional incomplete tensor completion that leverages Kaiming initialization and output bias correction to stabilize the optimization landscape and outperform existing methods with minimal computational overhead.

Original authors: Yuchao Su, Yixin Ran

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yuchao Su, Yixin Ran

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to reconstruct a giant, 3D puzzle of the world's traffic patterns or weather data. But here's the catch: 90% of the puzzle pieces are missing. You only have a few scattered pieces, and the picture you need to complete is incredibly complex, with traffic jams and storms changing every second.

This is the problem of High-Dimensional Incomplete (HDI) Tensors. It's like trying to guess the temperature in a city at 3 PM on a Tuesday, but you only have data from 10% of the sensors, and the data is messy.

The Old Way: A Rigid Blueprint

For a long time, scientists tried to solve this with "Linear" methods. Think of this like trying to fit a square peg into a round hole, or using a stiff, pre-drawn blueprint to guess a shape that keeps changing. These old methods assume the world is simple and straight. But real-world data (like traffic) is messy, curved, and unpredictable.

Then came Neural Tucker Factorization (NeuTucF). This was a big upgrade. Instead of a rigid blueprint, it used a flexible, "neural" approach. It learned to build a custom, 3D interaction map for every single piece of data it saw. It was much better at handling the complexity.

But, it had two major glitches:

The "Bad Start" Problem: When you start training these neural models, you have to guess the starting numbers (initialization). The old way used random guesses, like rolling dice. Sometimes, the dice roll was so wild that the numbers exploded, making the model confused and unstable right from the start.
The "Missing Offset" Problem: The model tried to learn the shape of the data (the curves and bumps) but forgot to account for the average level (the baseline). It was like trying to draw a mountain range but forgetting to draw the sea level first. The model forced itself to do too much work, trying to use the mountain shapes to explain the sea level, which made the whole picture slightly off.

The New Solution: KaBiN

The authors of this paper introduced a new model called KaBiN. They didn't build a giant, complicated new machine. Instead, they fixed the two glitches in the existing machine with two simple, clever tweaks.

1. The "Kaiming" Start (Better Dice)

Instead of rolling random dice to start the model, they used a specific, mathematically perfect way to pick the starting numbers called Kaiming Uniform Initialization.

The Analogy: Imagine you are teaching a child to balance on a tightrope. The old method threw the child onto the rope with a random push. The new method places them gently in the exact center with the perfect amount of tension. This prevents the "variance explosion" (the child falling off immediately) and keeps the training smooth and stable.

2. The "Bias" Button (The Sea Level)

They added a tiny, single number (a bias term) to the final step of the model.

The Analogy: Remember the mountain range? The old model tried to draw the mountains and the sea level at the same time, getting confused. The new model says, "Hey, let's just draw the mountains first, and then we'll add a little 'offset' button to lift or lower the whole picture to match the sea level." This separates the shape of the data from the average value, making the model much more accurate.

The Results

The researchers tested this on real-world traffic data (New York City taxis) and climate data (temperature sensors).

The Outcome: KaBiN completed the puzzles better than the old Neural Tucker model and better than many other complex methods.
The Cost: It was incredibly cheap to run. They didn't add any heavy machinery or deep layers. They just fixed the starting numbers and added one tiny button. It's like tuning a race car engine rather than building a new car.

In a Nutshell

The paper argues that you don't always need a bigger, more complex AI to get better results. Sometimes, you just need to start with a better plan (Kaiming initialization) and make sure you account for the basics (adding a bias term). By doing these two simple things, the model becomes more stable, learns faster, and predicts missing traffic and weather data with higher accuracy.

Technical Summary: Robust Neural Tucker Factorization with Bias Correction and Adaptive Initialization (KaBiN)

Problem Statement
High-dimensional incomplete (HDI) tensors are prevalent in applications such as traffic monitoring, climate observation, and sensor network analysis. These datasets are characterized by complex, non-linear interactions across multiple modes (e.g., spatial, temporal, and environmental metrics) but suffer from extreme sparsity due to missing or noisy observations. While classical tensor factorization methods (CP and Tucker) offer low-rank approximations, their linear assumptions often fail to capture localized anomalies or saturated patterns. Conversely, Neural Tucker Factorization (NeuTucF) successfully integrates neural output mapping with Tucker-style high-order interactions to model non-linear dynamics. However, the original NeuTucF framework exhibits sensitivity to implementation-level choices, specifically parameter initialization and the absence of a dedicated output bias. Suboptimal random uniform initialization can trigger variance explosion in the cubically expanded interaction spaces, leading to gradient saturation. Furthermore, the lack of a translation parameter forces interaction weights to implicitly absorb global statistical deviations, destabilizing the optimization landscape.

Methodology
The paper proposes KaBiN (Kaiming initialization and Bias-enhanced Neural Tucker Factorization), a lightweight modification of the NeuTucF architecture designed to stabilize training and improve reconstruction accuracy without increasing model complexity. The methodology retains the core NeuTucF pipeline:

Embedding: Discrete mode indices $(i, j, k)$ are mapped to low-dimensional embeddings ( $a_i, b_j, c_k$ ).
Interaction: These embeddings are combined via an outer product to form a sample-wise Tucker-style interaction tensor, which is then flattened.
Prediction: The flattened vector is passed through a linear layer followed by a sigmoid activation to predict the normalized entry value.

KaBiN introduces two specific enhancements to this pipeline:

Kaiming Uniform Initialization: Instead of random uniform initialization, the embedding matrices and the final Tucker linear weight vector are initialized using Kaiming uniform initialization. This strategy scales the initial weights based on the input dimension ( $fan\_in$ ), systematically avoiding variance explosion or vanishing gradients during the early phases of non-linear training.
Bias Correction: An explicit scalar bias term ( $b$ ) is added to the final linear layer ( $z_{ijk} = w^T t_{ijk} + b$ ). This allows the model to decouple global statistical shifts (mean offsets) from local structural representations. Consequently, the interaction weights focus on modeling high-order mode interactions, while the bias term adjusts the global level of the prediction scores.

The model is trained by minimizing the mean squared error (MSE) on observed entries, with optional L2 regularization. The computational cost remains asymptotically identical to the baseline NeuTucF, as the method only introduces a single scalar parameter.

Key Contributions

Improved Initialization and Bias: The paper introduces KaBiN, which replaces random uniform initialization with Kaiming uniform initialization and enables a final-layer bias term in the Tucker linear mapping.
Preservation of Efficiency: The proposed method maintains the original Tucker interaction structure, ensuring the model remains lightweight, parameter-efficient, and easy to implement without adding deep non-linear layers or increasing architectural complexity.
Empirical Validation: Experiments on three real-world HDI tensor datasets (NYCTaxi and PCTemp) demonstrate that KaBiN effectively reduces reconstruction errors (MAE and RMSE) compared to both classical baseline methods and the original NeuTucF framework.

Experimental Results
The authors evaluated KaBiN on three datasets: NYCTaxi p10 (10% observation), NYCTaxi p20 (20% observation), and PCTemp p10 (10% observation). The results indicate consistent improvements over the original NeuTucF (M8) and other baselines (including CP-based, non-negative, and Bayesian methods):

NYCTaxi p10: KaBiN reduced MAE from 3.9759 to 3.9289 and RMSE from 7.1218 to 6.9358.
NYCTaxi p20: KaBiN reduced MAE from 3.8535 to 3.8402 and RMSE from 6.9373 to 6.7992.
PCTemp p10: KaBiN reduced MAE from 0.8575 to 0.8483 and RMSE from 1.1429 to 1.1303.

An ablation study on the PCTemp dataset confirmed that the combination of Kaiming initialization and the bias term yields the lowest error rates, suggesting that the two components are complementary. The bias term specifically helps the model accommodate global non-zero target translations that optimal initialization scales alone cannot balance.

Significance and Claims
The paper claims that KaBiN provides a "simple yet effective" solution to the representational bottlenecks of standard neural tensor mapping. By elegantly decoupling global mean shifts from local structural representations, the framework creates a highly stable and well-conditioned optimization landscape. The authors emphasize that these improvements are achieved with minimal computational overhead, making the approach suitable for scenarios where model complexity must be kept low.

Looking forward, the authors suggest that future research could extend this calibration mechanism to dynamic, time-varying bias translations for streaming tensor environments. Additionally, they propose that the variance-scaled optimization framework could be adapted for federated latent factor learning protocols to support privacy-preserving, cross-organizational spatio-temporal signal recovery.