Accelerated Learning of High Dimensional Functions with a Tensor-Featured Training Network

This work presents a method to accelerate the optimization of learning high dimensional functions using deep neural networks (DNNs) by introducing contextual features into the first layer of a DNN. The optimization procedure keeps the input-feature basis fixed while DNN parameters are optimized via standard gradient descent; after optimization, the feature layer is updated before DNN optimization resumes. The feature layer contains two types of functions: rank-1 features that can be evaluated quickly in a matrix-free way on the domain, and tensor features that must first be decomposed using tensor network (TN) decomposition strategies. In particular, the study examines the effect of adding features that distill pretrained DNNs into TNs using a discretize and decompose strategy, leveraging a randomized tensor decomposition to efficiently decompose high-dimensional functions constructed from discretized DNNs. Using randomization, the storage cost of decomposing high dimensional functions is reduced by at least 8 orders of magnitude, enabling efficient training of models between 5 and 40 dimensions.

The method is motivated by the inefficiency of DNN optimization, which is non-convex and can require hundreds of compute hours, especially for physics-informed neural networks. Tensor network methods are introduced as a way to break the curse of dimensionality by factorizing high-dimensional functions into sums of outer products of low-dimensional component functions. The idea is similar to a warm-start initialization technique: a collection of functions is generated and applied to input data to a DNN, with the output added to the input layer. A straightforward feature would be a previously optimized DNN, but this fails; instead, a TN decomposition strategy is used. The authors identify two subclasses of basis functions: those that can be applied quickly in a matrix-free way and those that can be approximated via a TN decomposition. A 2-step optimization process is introduced where, given input features, DNN parameters are optimized recursively by gradient descent for a fixed number of iterations, after which the basis functions are evaluated and updated. This draws parallels with random feature training, as the DNN can be considered a combination of its input layer and neuron layers, with the feature layer having an outer-product structure similar to a rank-1 TN decomposition, hence the name Tensor-Featured training.

The theoretical background covers the canonical polyadic decomposition (CPD), which represents a function of d variables as a sum of outer products of single variables, reducing storage complexity from N d to NDR CP. The CPD is optimized via alternating least squares (ALS), which is non-convex but effective in practice. To address the curse of dimensionality in CPD-ALS, the authors use a sampled least squares method (CPD-RALS) where the number of samples is proportional to the rank of the CPD and independent of dimensionality. The CPD-RALS reduces computational complexity from exponential to polynomial in function dimension. The authors also establish a connection between CPD-ALS and Gaussian processes, showing that the CPD-ALS iteratively constructs 1-dimensional basis functions using kernel matrices, with the squared Khatri-Rao product viewed as a product kernel of many 1-dimensional kernel matrices. The randomized solver samples the leverage score distribution of each factor matrix to determine a compact set of basis points, effectively modifying basis function weights.

In the computational experiments, the authors restrict to functions of dimension d ∈ [5, 40], using a DNN with 3 layers and 100 features per layer, with activation functions [relu, relu, identity]. The ADAM gradient descent algorithm is used, stopping after 3000 epochs. The spatial domain is (-1,1) and time domain (0,1), with 10,000 points per mode on a regular grid. Training sets contain 1000 random points and validation sets contain 5000 random points. For the CPD, functions are decomposed on a grid of 500 points per mode to a rank of 30, with 3000 samples per least squares subproblem in the CPD-RALS, reducing the target tensor size from 500 × 500(d-1) to 500 × 3000, a 20 million times reduction in storage cost in the smallest case. The CPD is optimized using 100 ALS iterations, with nearest neighbor interpolation for off-grid points.

The first set of results investigates using a pretrained DNN as a tensor feature for learning the solution to a nonlinear elliptic equation, u(x) = sin(π/2 (1 - x2) squared.5), in dimensions 5, 10, and 15. Three training strategies are compared: conventional training, NN featured training (using the DNN output as a feature), and Tensor-Featured training (using the CPD interpolation of the DNN as a feature). The results show that decomposing the conventionally optimized DNN into a TN significantly improves accuracy, as the tensor decomposition smooths the DNN function over the domain. Using the DNN directly as a feature provides no practical improvement over conventional training, but using the CPD interpolation as a feature yields improvement. However, the CPD approximation alone outperforms the Tensor-Featured DNN when the conventionally trained model is not accurate relative to the ground truth. The authors hypothesize that if the conventional training is closer to the ground truth, Tensor-Featured training will be significantly accelerated.

To test this hypothesis, the training sample batch size is increased from 1,000 to 10,000, improving the accuracy of the conventionally trained DNN. With a more accurate reference, the CPD approximation captures more character of the true function, and the Tensor-Featured training is greatly improved, with mean squared error in the validation set falling dramatically. The authors demonstrate that this 2-step optimization process can be repeated iteratively. In a 3-step procedure for the 15-dimensional example, a DNN is trained conventionally, decomposed into a CPD, passed as a feature into a new DNN, trained again, decomposed again, and passed to a third DNN. The results show that this multi-step training process iteratively reduces the error in the trained DNNs, with relative L2 error decreasing from 0.272 (conventional) to 0.119 (first Tensor-Featured iteration) to 0.0619 (second Tensor-Featured iteration), and absolute mean squared error decreasing from 6.32 × 10−2 to 1.20 × 10−2 to 3.26 × 10−3.

The second set of results addresses the challenge that Tensor-Featured training requires features to provide meaningful context. If the target function is unknown and conventional DNN training converges slowly, the Tensor-Featured training as presented may not improve optimization. Therefore, the authors introduce two types of features: rank-1 features (elementary functions that can be evaluated quickly, such as sum of inputs, norm, sine/cosine of norm, sum of sines/cosines, exponential of norm) and TN features (more complex functions that must be decomposed, such as partially trained DNNs). For learning the solution to the hyperbolic wave equation u(t, x) = (exp(t2) - 1) sin(π/2 (1 - x) squared.5) in 15 dimensions, the authors test the effect of different features. Adding a single feature that aligns with the problem improves training efficiency, while a feature that does not align does not significantly improve performance. Adding multiple rank-1 features in parallel systematically improves accuracy over conventional training and the best single feature, though the process can fail if a feature puts the DNN close to a local extrema (e.g., adding the conventionally trained DNN as a feature causes convergence back to the conventional solution, but replacing it with its CPD approximation resolves this issue).

Using Tensor-Featured training with only rank-1 features, the authors can improve DNN training efficiency without prior knowledge of the function. Comparing true and predicted function values on the validation set, the Tensor-Featured training with rank-1 features more accurately predicts nearly all points compared to conventional training, and converges more quickly when the training batch size is increased. After training with rank-1 features, the DNN is discretized and decomposed into a TN, and the CPD approximation has little impact on accuracy. The decomposition effectively cannibalizes all features from the preliminary DNN optimization into a single function. In a second round of Tensor-Featured training, using the CPD interpolation as a feature (TN feature) is more accurate than using the original DNN as a feature. Combining the CPD with the original rank-1 features yields a new Tensor-Featured training that is more accurate than the original, and modifying the feature list (removing some rank-1 features) can outperform the CPD approximation alone. This two-step process can be repeated iteratively until convergence.

Finally, the authors test the method in 40 dimensions for the wave equation. Tensor-Featured training with only rank-1 features more accurately learns the objective function in just 3000 epochs compared to conventional training, though increasing dimensionality slows convergence of both procedures. The cost comparison shows that Tensor-Featured training is only marginally more expensive than conventional training, with training time of 77.47 seconds versus 76.93 seconds, parameters of 15,001 versus 14,301, and storage of 117.500 KiB versus 112.031 KiB.

In conclusion, the authors introduce a strategy to add context-oriented features into the input layer of a DNN and leverage tensor decomposition to develop a 2-step optimization strategy. They find that DNNs optimized with a pretrained DNN feature quickly converge to the same local minima as their pretrained feature, so they discretize the DNN and decompose it using efficient randomized tensor decomposition (CPD-RALS), reducing memory requirements by between 8 and 34 orders of magnitude. The tensor decomposition has a smoothing effect on functions encoded in trained DNNs, and introducing the TN approximation as a feature helps kick optimizations out of local minima. Initial testing shows mixed results when only the TN decomposition of a DNN is used as a feature: if the trained DNN is close to a global extrema, Tensor-Featured training significantly improves subsequent DNN optimization, but if not, the tensor feature does not improve training. To address this, the authors introduce a more robust strategy where any number of functions acting on any portion of the domain can be added as features, split into rank-1 features (easy to apply) and TN features (must be decomposed). Using a combination of rank-1 and TN features, they significantly accelerate convergence of model error in the validation set compared to conventional DNN training. Future work will consider methods to efficiently choose rank-1 features, determine quadrature grid points for TN decomposition, pick appropriate TN decomposition strategies, study how the product structure of features influences DNN parameter optimization, and improve CPD optimization for extremely high-dimensional functions by leveraging symmetry and sparsity of PDEs.

Based on this paper, here are the specific improvements I can make to AI systems, along with what the improved systems can do:

1. Accelerated Training of Physics-Informed Neural Networks (PINNs)

• Improvement: Implement a two-step optimization loop where a DNN is trained conventionally, then decomposed into a tensor network (CPD) via randomized decomposition, and the CPD approximation is added as a contextual feature to the input layer of a new DNN. Repeat iteratively.

• What the improved system can do: Solve high-dimensional PDEs (5–40 dimensions) with significantly lower validation error and faster convergence than standard PINN training. For example, in the 15D wave equation case, relative L2 error dropped from 0.272 to 0.0619 over three iterations, and absolute MSE decreased by 20x.

2. Robust Feature Engineering Without Prior Knowledge

• Improvement: Add a set of rank-1 features (e.g., sum of inputs, norm, sine/cosine of norm, exponential of norm) to the input layer of any DNN, alongside TN-decomposed features from partially trained models. This removes the need for problem-specific feature design.

• What the improved system can do: Automatically improve DNN training efficiency for unknown target functions, even when conventional training is slow. In the 15D hyperbolic wave equation, rank-1 features alone improved accuracy over conventional training, and combining them with a CPD feature further reduced error.

3. Memory-Efficient High-Dimensional Function Representation

• Improvement: Use CPD-RALS (randomized tensor decomposition) to compress a trained DNN into a sum of outer products of 1D functions. This reduces storage cost by 8–34 orders of magnitude compared to full grid discretization (e.g., from 500 × 500(d-1) to 500 × 3000 samples).

• What the improved system can do: Store and reuse complex high-dimensional functions (e.g., learned PDE solutions) in a compact form, enabling deployment on memory-constrained devices or transfer to other tasks without retraining.

4. Escaping Local Minima in DNN Optimization

• Improvement: Replace a pretrained DNN feature with its tensor-decomposed (CPD) approximation. The decomposition smooths the function, preventing the optimizer from collapsing back to the same local minimum.

• What the improved system can do: Avoid the failure mode where a DNN with a pretrained feature quickly converges to the pretrained model’s suboptimal solution. The CPD feature kicks optimization into better basins, as shown when CPD-based features outperformed direct DNN features in both 15D and 40D tests.

5. Iterative Self-Improvement Loop for DNNs

• Improvement: Implement a multi-step training procedure: train → decompose → add as feature → retrain → decompose again → repeat. Each step uses the previous model’s TN approximation as a new feature.

• What the improved system can do: Continuously reduce error over multiple iterations without manual intervention. The 3-step procedure in 15D reduced relative L2 error by over 4x, demonstrating a self-correcting optimization path.

6. Cost-Effective High-Dimensional Learning

• Improvement: Use rank-1 features (which require no decomposition) for initial training, then decompose the trained model only once for a second round. This keeps overhead minimal—training time increased by only 0.7% (77.47s vs 76.93s) and storage by 5% in the 40D case.

• What the improved system can do: Achieve superior accuracy in high-dimensional regression or PDE solving with negligible additional computational cost, making it practical for real-time or resource-limited applications.

7. Generalizable Feature Selection for Unknown Targets

• Improvement: Dynamically add or remove rank-1 features based on validation performance, and combine them with TN features from partially trained models. The paper shows that removing some rank-1 features after CPD addition can outperform the CPD alone.

• What the improved system can do: Adapt its feature set during training to avoid local extrema and improve generalization, without requiring prior knowledge of the target function’s structure.

8. Scalable Decomposition for Very High Dimensions

• Improvement: Use CPD-RALS with leverage-score sampling to reduce the least-squares subproblem size from exponential to polynomial in dimension, independent of d.

• What the improved system can do: Handle functions in 40+ dimensions efficiently, as demonstrated by successful training and decomposition in the 40D wave equation case, where conventional methods would be computationally infeasible.

9. Improved Transfer Learning via Tensor Features

• Improvement: Use the CPD decomposition of a trained DNN as a feature for a new DNN on a related task, rather than using the raw DNN weights or outputs.

• What the improved system can do: Transfer knowledge more effectively between tasks, as the smoothed TN representation provides a more robust prior than the original DNN, avoiding negative transfer due to overfitting or local minima.

10. Faster Convergence with Larger Batches

• Improvement: Increase training batch size (e.g., from 1,000 to 10,000) when using Tensor-Featured training, as the paper shows this dramatically improves the quality of the CPD feature and subsequent training.

• What the improved system can do: Leverage larger datasets more effectively in high-dimensional settings, achieving lower error in fewer epochs than conventional training with the same batch size.