TADrop: Tensor-Wise Adaptive Drop in Deep Networks

• TADrop is a method that applies stochastic dropout on low-rank tensor decompositions to enhance robustness in neural network training.
• It adapts sparsification per tensor using quantile-based statistics to prune redundant weights while preserving informative components.
• Empirical evaluations demonstrate improved accuracy and adversarial resilience across benchmarks like ImageNet and multimodal tasks.

Tensor-Wise Adaptive Drop (TADrop) refers to two influential but distinct methodologies in neural network learning and model merging: (1) randomized low-rank tensor regularization for robust learning in tensor regression layers, and (2) adaptive, distribution-aware sparsification of parameter tensors in model merging pipelines. Both approaches exploit the statistical diversity inherent in neural network parameters at the tensor level, enhancing robustness or mitigating parameter interference across a range of domains (Kolbeinsson et al., 2019, Luo et al., 8 Aug 2025).

1. Formal Definitions and Core Principles

TADrop in its original formulation for robust learning operates on low-rank multilinear parameterizations of weight tensors, such as those arising in Tucker or CP (CANDECOMP/PARAFAC) decompositions. For a weight tensor $$W\in\mathbb{R}^{I_0\times I_1\times\cdots\times I_N}$$, TADrop stochastically drops entire rank-1 components by introducing binary masks on tensor factors or the core, thereby randomizing the effective rank at each iteration. In the context of adaptive model merging, TADrop refers to a per-tensor pruning strategy, where the sparsity level is individually set based on each tensor's empirical magnitude distribution, preserving heavy-tailed (informative) tensors and pruning redundant (near-zero) tensors more aggressively. This contrasts with global, uniform sparsification rates that ignore parameter heterogeneity (Kolbeinsson et al., 2019, Luo et al., 8 Aug 2025).

2. Mathematical Foundation and Algorithmic Formulation

Rank-Dropout on Low-Rank Factorizations

• Tucker Decomposition:

Given ranks $(R_0,\ldots,R_N)$, represent $W$ as $G\times_0 U^{(0)} \times_1 \cdots \times_N U^{(N)}$, with $G$ the core and $U^{(n)}$ are factor matrices. For each $n$ and $r_n$, TADrop samples mask variables $\lambda_{r_n}^{(n)}\sim\mathrm{Bernoulli}(\theta)$, builds diagonal masks $M^{(n)}$, and applies these to $G$ and $U^{(n)}$.

• CP Decomposition:

$$W=\sum_{r=1}^R u_r^{(0)}\otimes \cdots \otimes u_r^{(N)}$$, with mask $\lambda_r\sim\mathrm{Bernoulli}(\theta)$ applied per component.

• Expected Loss Regularization:

Training with TADrop in the Tucker case results in the following expectation over mask variables:

$$\mathbb{E}_{\lambda}\sum_k(y_k-\langle X_k, \widetilde{W}\rangle)^2 = \sum_k (y_k - \theta^N\langle X_k, W\rangle)^2 + \theta^N(1-\theta^N)\sum_k\langle G^{\odot 2} \times_0 (U^{(0)})^{\odot 2}\cdots, (X_k)^{\odot 2}\rangle$$

where $\odot 2$ denotes elementwise square. This introduces a data-dependent $\ell_2$-type penalty. Similar results hold for CP.

Tensor-Wise Adaptive Drop for Model Merging

• Quantile-Ratio Statistic:

For task vector tensors $\tau_k^i = \theta_k^i - \theta_{pre}^i$, drop ratio $s_i$ for tensor $i$ is determined by

$$d_k^i = \frac{Q_a(|\tau_k^i|)}{Q_b(|\tau_k^i|)+\epsilon}$$

with $Q_p(\cdot)$ the $p$-th quantile and recommended $a=0.50$, $b=0.95$.

• Sparse-then-Scale:

Entries with $|\tau_k^i|\leq T_i$ (the $s_i$-quantile) are set to zero, resulting in $\hat{\tau}_k^i$. The tensor is then rescaled to preserve $\ell_2$ norm:

$$\tau_k^{\prime\,i} = \hat{\tau}_k^i \cdot \frac{\|\tau_k^i\|_2}{\|\hat{\tau}_k^i\|_2 + \epsilon}$$

• Integration:

TADrop replaces global sparsification in merging methods such as Task Arithmetic, TIES-Merging, and EMR-Merging, and is plug-and-play within their preprocessing steps.

3. Algorithmic Implementation and Pseudocode

Rank-Dropout in Tensor Regression Layers

1. Initialize: Factor/core tensors, maximal ranks $R_k$ or $R$, drop-probability $\theta$.
2. Training Epoch (Per Mini-batch):
   • Sample Masks: For each factor/core mode, draw independent Bernoulli masks.
   • Form Dropped Factors: Apply masks to factors and core, reconstruct dropped-rank weight.
   • Forward: Compute predictions via contraction.
   • Loss and Backprop: Compute loss and update all parameters using SGD/Adam.
3. Inference: Use full parameterization scaled by $\theta^N$ (Tucker) or $\theta$ (CP); no masks are sampled.

Per-Tensor Adaptive Drop for Model Merging

for each task k in 1..K: for each tensor i in 1..n: A = abs(tau_k^i) q_a = quantile(A, a) q_b = quantile(A, b) s_i = q_a/(q_b + epsilon) T_i = quantile(A, s_i) M_i = (A > T_i) hat_tau_k^i = tau_k^i * M_i tau_k'^i = hat_tau_k^i * (norm(tau_k^i) / (norm(hat_tau_k^i) + epsilon)) tau_k^i = tau_k'^i

Norm-preserving scaling is essential; omission leads to substantial accuracy collapse (Luo et al., 8 Aug 2025).

4. Hyperparameters, Complexity, and Best Practices

Rank-Dropout for Robust Learning

• Drop-probability $\theta$: Typical range $[0.6, 0.9]$.
• Rank hyperparameters $R_k$ (Tucker) or $R$ (CP): Control expressive capacity vs. parameter efficiency.
• Sampling with Bernoulli masks or with-replacement is effective.
• Overhead compared to standard dropout is negligible for low ranks, as only the factor/core tensors are masked and reconstructed (Kolbeinsson et al., 2019).

TADrop for Model Merging

• Quantiles: $a=0.50$, $b=0.95$ recommended; values are robust across datasets.
• No per-task hyperparameter tuning required.
• Compatible with all merging methods employing sparsification.
• Critical to scale each tensor's $\ell_2$ norm after pruning.

5. Empirical Performance and Robustness

Setting	Baseline	TADrop Variant	Metric/Delta
ImageNet, ResNet-101	77.1%	78.0%	Top-1 Acc, +0.9
CIFAR-100, ResNet-18	75.9% (FC)	76.05% (Tucker), 76.19% (CP)	Top-1 Acc, +0.3
UK Biobank Brain MRI	MAE=2.96	TADrop-Tucker: 2.65, TADrop-CP: 2.58	Lower better
Model Merging, ViT-B/32	88.7% (EMR)	90.7% (TADrop)	Acc, +2.0
BEiT3, COCO	79.5%	80.9%	Retrieval Acc, +1.4

On ImageNet, TADrop in tensor regression layers yields up to 0.9 percentage point gains in top-1 accuracy, while maintaining or increasing robustness under adversarial perturbations (up to 10–40 percentage point improvement against FGSM, BIM, PGD) (Kolbeinsson et al., 2019). In model merging scenarios, TADrop increases accuracy by up to 2.0 points in ViT-B/32 pipelines and 1.4–2.7 points on multimodal benchmarks (Luo et al., 8 Aug 2025). Ablations demonstrate the necessity of norm-preserving scaling and the superiority of adaptive sparsification over global or random approaches.

6. Insights, Limitations, and Extensions

TADrop effectively addresses parameter and statistical heterogeneity both in robust learning and model merging. For rank-dropout, the variance term in the expected loss connects naturally to data-dependent $\ell_2$ regularization, strengthening the link with classical nonlinear dropout. In model merging, TADrop's quantile-ratio based adaptation is robust to outliers and does not require data access.

Limitations include assumptions of unimodal, heavy-tailed parameter distributions; extensions could leverage additional statistics (e.g., kurtosis), or parametric mappings from per-tensor statistics to drop rates. While TADrop is data-free in design, hybridization with calibration-based or activation-aware adaptation remains largely unexplored. Extension to convolutional or graph architectures is an open direction (Luo et al., 8 Aug 2025).

7. Applications and Recommendations

TADrop is recommended for insertion at tensor regression layers in deep models, effectively replacing dense FC heads to yield robustness improvements. In model merging, it replaces global sparsification and can be used as a direct drop-in for leading approaches including Task Arithmetic, TIES-Merging, and EMR-Merging. Parameters should be initialized according to memory or computation budgets, and default hyperparameters suffice for most settings. Standard optimizers and learning rate schedules are compatible with both TADrop formulations. Benefits include improved clean and noisy-data accuracy, increased adversarial robustness, and negligible overhead even at high compression rates (Kolbeinsson et al., 2019, Luo et al., 8 Aug 2025).