TT-NF: TENSOR TRAIN NEURAL FIELDS

Abstract

Learning neural fields has been an active topic in deep learning research, focusing, among other issues, on finding more compact and easy-to-fit representations. In this paper, we introduce a novel low-rank representation termed Tensor Train Neural Fields (TT-NF) for learning neural fields on dense regular grids and efficient methods for sampling from them. Our representation is a TT parameterization of the neural field, trained with backpropagation to minimize a non-convex objective. We analyze the effect of low-rank compression on the downstream task quality metrics in two settings. First, we demonstrate the efficiency of our method in a sandbox task of tensor denoising, which admits comparison with SVD-based schemes designed to minimize reconstruction error. Furthermore, we apply the proposed approach to Neural Radiance Fields, where the low-rank structure of the field corresponding to the best quality can be discovered only through learning.

1. INTRODUCTION

Following the growing interest in deep neural networks, learning neural fields has become a promising research direction in areas concerned with structured representations. However, precision is usually at odds with the computational complexity of these representations, which makes training them and sampling from them a challenge. In this paper, we investigate interpretable low-rank neural fields defined on dense regular grids and efficient methods for learning them. Since, in extreme cases, the dimensionality of such fields can exceed the memory size of a typical computer by several orders of magnitude, we look at the problem of learning such fields from the angle of stochastic methods. Tensor decompositions have become a ubiquitous tool for dealing with structured sparsity of intractable volumes of data. Within the large family of tensor decompositions, we focus on the Tensor Train (TT) (Oseledets, 2011) , also known as the Matrix Product State in physics. TT is notable for its high-capacity representation, efficient algebraic operations in the low-rank space, and support of SVD-based algorithms for data approximation. As such, we consider TT-SVD (Oseledets, 2011) and TT-cross (Oseledets & Tyrtyshnikov, 2010) methods for obtaining a low-rank representation of the full tensor. While TT-SVD requires access to the full tensor at once (which might already be problematic in specific scenarios), TT-cross requires access to data through a black-box function, computing (or looking up) elements by their coordinates on demand. Both methods operate under the assumption of noise-free data and are not guaranteed to output sufficiently good approximations in the presence of noise. While noise in observations is challenging for SVD-based schemes and requires devising tailored approaches to different noise types and magnitude (Zhou et al., 2022) , exploiting the low-rank structure of the field driven by data is even more challenging (Novikov et al., 2014; Boyko et al., 2020) and typically resorts to the paradigm of data updates through algebraic operations on TT. In this work, we take a step back and leverage the modern deep learning paradigm to parameterize neural fields as TT, coined TT-NF. Through deep learning tooling with support for automatic differentiation and our novel sampling methods, we obtain mini-batches of samples from the parameterized neural field and perform optimization of a non-convex objective defined by a downstream task. The optimization comprises the computation of parameter gradients with backpropagation and parameter updates with any suitable technique, such as SGD. We analyze TT-NF and several sampling techniques on a range of problem sizes and provide reference charts for choosing a sampling method based on memory and computational constraints. Next, we define a synthetic task of low-rank tensor denoising and demonstrate the superiority of the proposed 1

