SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ryan Thompson*  
University of New South Wales  
CSIRO's Data61
Amir Dezfouli†  
BIMLOGIQ
Robert Kohn  
University of New South Wales
Sparse linear models are one of several core tools for interpretable machine learning, a field of emerging importance as predictive models permeate decision-making in many domains. Unfortunately, sparse linear models are far less flexible as functions of their input features than black-box models like deep neural networks. With this capability gap in mind, we study a not-uncommon situation where the input features dichotomize into two groups: explanatory features, which are candidates for inclusion as variables in an interpretable model, and contextual features, which select from the candidate variables and determine their effects. This dichotomy leads us to the contextual lasso, a new statistical estimator that fits a sparse linear model to the explanatory features such that the sparsity pattern and coefficients vary as a function of the contextual features. The fitting process learns this function nonparametrically via a deep neural network. To attain sparse coefficients, we train the network with a novel lasso regularizer in the form of a projection layer that maps the network's output onto the space of  $\ell_1$ -constrained linear models. An extensive suite of experiments on real and synthetic data suggests that the learned models, which remain highly transparent, can be sparser than the regular lasso without sacrificing the predictive power of a standard deep neural network.
Sparse linear models—linear predictive functions in a small subset of features—have a long history in statistics, dating back at least to the 1960s (Garside, 1965). Nowadays, against the backdrop of elaborate, black-box models such as deep neural networks, the appeal of sparse linear models is largely their transparency and intelligibility. These qualities are sought in decision-making settings (e.g., consumer finance and criminal justice) and constitute the foundation of interpretable machine learning, a topic that has recently received significant attention (Murdoch et al., 2019; Molnar et al., 2020; Rudin et al., 2022; Marcinkevics and Vogt, 2023). Interpretability, however, comes at a price when the underlying phenomenon cannot be predicted accurately without a more expressive model capable of well-approximating complex functions, such as a neural network. Unfortunately, one must forgo direct interpretation of expressive models and instead resort to post hoc explanations (Ribeiro et al., 2016; Lundberg and Lee, 2017), which have their own flaws (Laugel et al., 2019; Rudin, 2019).
Motivated by a desire for interpretability and expressivity, this paper focuses on a setting where sparse linear models and neural networks can collaborate together. The setting is characterized by a not-uncommon situation where the input features dichotomize into two groups, which we call explanatory features and contextual features. Explanatory features are features whose effects are of primary interest. They should be modeled via a low-complexity function such as a sparse linear model for interpretability. Meanwhile, contextual features describe the broader predictive context,
*Corresponding author. Email: ryan.thompson1@unsw.edu.au  
†Part of this work was carried out while the author was at CSIRO's Data61.
Figure 1: Fitted coefficient functions from the contextual lasso for the house pricing dataset. Colored points indicate coefficient values at different locations. Grey points indicate coefficients equal to zero.
e.g., the location of the prediction in time or space (see the house pricing example below). These inform which explanatory features are relevant and, for those that are, their exact low-complexity effects. Given this role, contextual features are best modeled via an expressive function class.
The explanatory-contextual feature dichotomy described above leads to the seemingly previously unstudied contextually sparse linear model:
$$
g \left(\operatorname {E} [ y \mid \mathbf {x}, \mathbf {z} ]\right) = \sum_ {j \in S (\mathbf {z})} x _ {j} \beta_ {j} (\mathbf {z}). \tag {1}
$$
To parse the notation,  $y \in \mathbb{R}$  is a response variable,  $\mathbf{x} = (x_{1},\ldots ,x_{p})^{\top} \in \mathbb{R}^{p}$  are explanatory features,  $\mathbf{z} = (z_{1},\dots,z_{m})^{\top} \in \mathbb{R}^{m}$  are contextual features, and  $g$  is a link function (e.g., identity for regression or logit for classification). Via the contextual features, the set-valued function  $S(\mathbf{z})$  encodes the indices of the relevant explanatory features (typically, a small set of  $js$ ), while the coefficient functions  $\beta_{j}(\mathbf{z})$  encode the effects of those relevant features. The model (1) draws inspiration from the varying-coefficient model (Hastie and Tibshirani, 1993; Fan and Zhang, 2008; Park et al., 2015), a special case that assumes all explanatory features are always relevant, i.e.,  $S(\mathbf{z}) = \{1,\dots,p\}$  for all  $\mathbf{z} \in \mathbb{R}^{m}$ . We show that this new model is more powerful for various problems, including energy forecasting and disease prediction. For these tasks, sparsity patterns can be strongly context-dependent.
The main contribution of our paper is a new statistical estimator for (1) called the contextual lasso. The new estimator is inspired by the lasso (Tibshirani, 1996), a classic sparse learning tool with excellent properties (Hastie et al., 2015). We focus on tabular datasets as these are the most common use case for the lasso and its cousins. Whereas the lasso fits a sparse linear model that fixes the relevant features and their coefficients once and for all (i.e.,  $S(\mathbf{z})$  and  $\beta_{j}(\mathbf{z})$  are constant), the contextual lasso fits a contextually sparse linear model that allows the relevant explanatory features and coefficients to change according to the prediction context. To learn the map from contextual feature vector to sparse coefficient vector, we use the expressive power of neural networks. Specifically, we train a feedforward neural network to output a vector of linear model coefficients sparsified via a novel lasso regularizer. In contrast to the lasso, which constraints the coefficients'  $\ell_1$ -norm, our regularizer constrains the expectation of the coefficients'  $\ell_1$ -norm with respect to  $\mathbf{z}$ . To implement this new regularizer, we include a novel projection layer at the bottom of the network that maps the network's output onto the space of  $\ell_1$ -constrained linear models by solving a convex optimization problem.
To briefly illustrate our proposal, we consider data on property sales in Beijing, China, studied in Zhou and Hooker (2022). We use the contextual lasso to learn a pricing model with longitude and latitude as contextual features. The response is price per square meter. Figure 1 plots the fitted coefficient functions of three property attributes (explanatory features) and an intercept. The relevance and effect of these attributes can vary greatly with location. The elevator indicator, e.g., is irrelevant throughout inner Beijing, where buildings tend to be older and typically do not have elevators. The absence of elevators also makes it difficult to access higher floors, hence the negative effect of floor on price. Beyond the inner city, the floor is mostly irrelevant. Naturally, renovations are valuable everywhere, but more so for older buildings in the inner city than elsewhere. The flexibility of the contextual lasso to add or remove attributes by location, and simultaneously determine their coefficients, equips sellers with personalized interpretable models containing only the attributes most relevant to them. At the same time, these models outpredict both the lasso and a deep neural network; see Appendix A.
3The intercept is omitted throughout this paper to ease notation.
The rest of paper is organized as follows. Section 2 introduces the contextual lasso and describes techniques for its computation. Section 3 discusses connections with earlier related work. Section 4 reports experiments on synthetic and real data. Section 5 closes the paper with a discussion.
2 Contextual lasso
This section describes our estimator. To facilitate exposition, we first rewrite the contextually sparse linear model (1) more concisely:
$$
g \left(\operatorname {E} [ y | \mathbf {x}, \mathbf {z} ]\right) = \mathbf {x} ^ {\top} \boldsymbol {\beta} (\mathbf {z}).
$$
The notation  $\beta (\mathbf{z})\coloneqq \left(\beta_{1}(\mathbf{z}),\dots ,\beta_{p}(\mathbf{z})\right)^{\top}$  represents a vector coefficient function which is sparse over its codomain. That is, for different values of  $\mathbf{z}$ , the output of  $\beta (\mathbf{z})$  contains zeros at different positions. The function  $S(\mathbf{z})$ , which encodes the set of active explanatory features in (1), is recoverable as  $S(\mathbf{z})\coloneqq \{j:\beta_j(\mathbf{z})\neq 0\}$ .
At the population level, the contextual lasso comprises a minimization of the expectation of a loss function subject to an inequality on the expectation of a constraint function:
$$
\min  _ {\boldsymbol {\beta} \in \mathcal {F}} \quad \mathrm {E} [ l (\mathbf {x} ^ {\top} \boldsymbol {\beta} (\mathbf {z}), y) ] \quad \text {s . t .} \quad \mathrm {E} [ \| \boldsymbol {\beta} (\mathbf {z}) \| _ {1} ] \leq \lambda , \tag {2}
$$
where the set  $\mathcal{F}$  is a class of functions that constitute feasible solutions and  $l: \mathbb{R}^2 \to \mathbb{R}$  is the loss function, e.g., square loss  $l(\hat{y}, y) = (y - \hat{y})^2$  for regression or logistic loss  $l(\hat{y}, y) = -y\log(\hat{y}) - (1 - y)\log(1 - \hat{y})$  for classification. Here, the expectations are taken with respect to the random variables  $y$ ,  $\mathbf{x}$ , and  $\mathbf{z}$ . The parameter  $\lambda \geq 0$  controls the level of regularization. Smaller values of  $\lambda$  encourage  $\beta(\mathbf{z})$  towards zero over more of its codomain. Larger values have the opposite effect. The contextual lasso thus generalizes the lasso, which learns  $\beta(\mathbf{z})$  as a constant:
$$
\min  _ {\boldsymbol {\beta}} \quad \mathrm {E} \left[ l \left(\mathbf {x} ^ {\top} \boldsymbol {\beta}, y\right) \right] \qquad \text {s . t .} \quad \| \boldsymbol {\beta} \| _ {1} \leq \lambda .
$$
To reiterate the difference: the lasso coaxes the fixed coefficients  $\beta$  towards zero, while the contextual lasso coaxes the expectation of the function  $\beta(\mathbf{z})$  to zero. The result for the latter is coefficients that can change in value and sparsity with  $\mathbf{z}$ , provided the function class  $\mathcal{F}$  is suitably chosen.
Given a sample  $(y_{i},\mathbf{x}_{i},\mathbf{z}_{i})_{i = 1}^{n}$ , the data version of the population problem (2) replaces the unknown expectations with their sample counterparts:
$$
\min  _ {\boldsymbol {\beta} \in \mathcal {F}} \frac {1}{n} \sum_ {i = 1} ^ {n} l \left(\mathbf {x} _ {i} ^ {\top} \boldsymbol {\beta} (\mathbf {z} _ {i}), y _ {i}\right) \quad \text {s . t .} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {\beta} (\mathbf {z} _ {i}) \| _ {1} \leq \lambda . \tag {3}
$$
The set of feasible solutions to optimization problem (3) are coefficient functions that lie in the  $\ell_1$ -ball of radius  $\lambda$  when averaged over the observed data. To operationalize this estimator, we take  $\mathcal{F} = \{\beta_{\mathbf{w}}(\mathbf{z}) : \mathbf{w}\}$ , where  $\beta_{\mathbf{w}}(\mathbf{z})$  is a certain architecture of feedforward neural network (described shortly) parameterized by weights  $\mathbf{w}$ . This choice leads to our core proposal:
$$
\min  _ {\mathbf {w}} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} l \left(\mathbf {x} _ {i} ^ {\top} \boldsymbol {\beta} _ {\mathbf {w}} (\mathbf {z} _ {i}), y _ {i}\right) \quad \text {s . t .} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {\beta} _ {\mathbf {w}} (\mathbf {z} _ {i}) \| _ {1} \leq \lambda . \tag {4}
$$
Configuring a feedforward neural network such that its outputs are sparse and satisfy the  $\ell_1$ -constraint is not trivial. We introduce a novel network architecture to address this challenge.
2.2 Network architecture
The neural network architecture—depicted in Figure 2—involves two key components. The first and most straightforward component is a vanilla feedforward network  $\eta(\mathbf{z}) : \mathbb{R}^m \to \mathbb{R}^p$ . The purpose of the network is to capture the nonlinear effects of the contextual features on the explanatory features.
4The  $\ell_1$  -ball is the convex compact set  $\{\mathbf{x}\in \mathbb{R}^p:\| \mathbf{x}\| _1\leq \lambda \}$
Figure 2: Network architecture. The contextual features  $\mathbf{z}$  pass through a series of hidden layers. The resulting dense coefficients  $\eta_{1},\ldots ,\eta_{p}$  then enter a projection layer to produce sparse coefficients  $\beta_{1},\ldots ,\beta_{p}$ . Here, the last coefficient is gray to illustrate that it is zeroed-out by the projection layer.
Since the network involves only hidden layers with standard affine transformations and nonlinear maps (e.g., rectified linear activation functions), the coefficients they produce generally do not satisfy the contextual lasso constraint and are not sparse. To enforce the constraint and attain sparsity, we employ a novel projection layer as the second main component of our network architecture.
The projection layer takes the dense coefficients  $\pmb{\eta}(\mathbf{z})$  from the network and maps them to sparse coefficients  $\beta(\mathbf{z})$  by performing a projection onto the  $\ell_1$ -ball. Because the contextual lasso does not constrain each coefficient vector to the  $\ell_1$ -ball, but rather constrains the average coefficient vector, we project all  $n$  coefficient vectors  $\pmb{\eta}(\mathbf{z}_1), \dots, \pmb{\eta}(\mathbf{z}_n)$  together. That is, we take the final sparse coefficients  $\beta(\mathbf{z}_1), \dots, \beta(\mathbf{z}_n)$  as the minimizing arguments of a convex optimization problem:
$$
\beta \left(\mathbf {z} _ {1}\right), \dots , \beta \left(\mathbf {z} _ {n}\right) := \underset {\boldsymbol {\beta} _ {1}, \dots , \boldsymbol {\beta} _ {n}: \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {\beta} _ {i} \| _ {1} \leq \lambda} {\arg \min } \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {\eta} \left(\mathbf {z} _ {i}\right) - \boldsymbol {\beta} _ {i} \| _ {2} ^ {2}. \tag {5}
$$
The minimizers of this optimization problem are typically sparse thanks to the geometry of the  $\ell_1$ -ball. The idea of including optimization as a layer in a neural network is explored in previous works (see, e.g., Amos and Kolter, 2017; Agrawal et al., 2019). Yet, to our knowledge, no previous work has studied optimization layers (also known as implicit layers) for inducing sparsity in a neural network.
The optimization problem (5) does not admit an analytical solution, though it is solvable by general purpose convex solvers (see, e.g., Boyd and Vandenberghe, 2004). However, because (5) is a highly structured problem, it is also amenable to more specialized algorithms. Such algorithms facilitate the type of scalable computation necessary for deep learning. Duchi et al. (2008) provide a low-complexity algorithm for solving (5) when  $n = 1$ . Algorithm 1 below is an extension to  $n \geq 1$ . The algorithm consists of two main steps: (1) computing a thresholding parameter  $\theta$  and (2) soft-
Algorithm 1 Projection onto  $\ell_1$ -ball
input Dense coefficients  $\pmb{\eta}_1,\dots ,\pmb{\eta}_n$  and radius  $\lambda$    
Set  $\pmb {\mu} = (|\pmb{\eta}_1^\top |,\dots ,|\pmb{\eta}_n^\top |)^\top$    
Sort  $\pmb{\mu}$  in decreasing order:  $\mu_{i}\geq \mu_{j}$  for all  $i <   j$    
Set  $k_{\mathrm{max}} = \max \left\{k:\mu_k > \left(\sum_{l = 1}^k\mu_l - n\lambda\right) / k\right\}$    
Set  $\theta = \left(\sum_{k = 1}^{k_{\mathrm{max}}}\mu_k - n\lambda\right) / k_{\mathrm{max}}$    
Compute  $\beta_{1},\ldots ,\beta_{n}$  as  $\beta_{ij} = \mathrm{sign}(\eta_{ij})\max (|\eta_{ij}| - \theta ,0)$  for  $i = 1,\dots ,n$  and  $j = 1,\dots ,p$  output Sparse coefficients  $\beta_{1},\ldots ,\beta_{n}$
thresholding the inputs using the computed  $\theta$ . Critically, the operations comprising Algorithm 1 are suitable for computation on a GPU, meaning the model can be trained end-to-end at scale.
Although the projection algorithm involves some nondifferentiable operations, most deep learning libraries provide gradients for these operations, e.g., sort is differentiated by permuting the gradients and abs is differentiated by taking the subgradient zero at zero. The gradient obtained by automatically differentiating through Algorithm 1 is the same as that from implicitly differentiating through a convex solver (e.g., using cvxpylayers of Agrawal et al., 2019), though the latter is slower. In subsequent work, Thompson et al. (2023) derive analytical gradients for a projection layer that maps matrices onto an  $\ell_1$ -ball. Their result readily adapts to the vector case dealt with here.
Computation of the thresholding parameter is performed only during training. For inference, the estimate  $\hat{\theta}$  from the training set is used for soft-thresholding. That is, rather than using Algorithm 1
as an activation function when performing inference, we use  $T(x) \coloneqq \mathrm{sign}(x) \max(|x| - \hat{\theta}, 0)$ . The purpose of using the estimate  $\hat{\theta}$  rather than recomputing  $\theta$  via the algorithm is because the  $\ell_1$ -constraint applies to the expected coefficient vector. It need not be the case that every coefficient vector produced at inference time lies in the  $\ell_1$ -ball, which would occur if the algorithm is rerun.
Algorithm 1 computes the thresholding parameter using the full batch of  $n$  observations. The algorithm can also be applied to mini-batches during training. Once training is complete, the estimate  $\hat{\theta}$  for inference can be obtained via a single forward pass of the full batch through the network.
2.3 Grouped explanatory features
In certain settings, the explanatory features may be organized into groups such that all the features in a group should be selected together. These groups may emerge naturally (e.g., genes in the same biological path) or be constructed for a statistical task (e.g., basis expansions for nonparametric regression). The prevalence of such problems has led to the development of sparse estimators capable of handling groups, one of the most well-known being the group lasso (Yuan and Lin, 2006; Meier et al., 2008). Perhaps unsurprisingly, the contextual lasso extends gracefully to grouped selection.
Let  $\mathcal{G}_1, \ldots, \mathcal{G}_g \subseteq \{1, \ldots, p\}$  be a set of  $g$  nonoverlapping groups, and let  $\beta_k(\mathbf{z})$  and  $\mathbf{x}_k$  be the coefficient function and explanatory features restricted to group  $\mathcal{G}_k$ . In the noncontextual setting, the group lasso replaces the  $\ell_1$ -norm  $\| \beta \|_1$  of the lasso with a sum of group-wise  $\ell_2$ -norms  $\sum_{k=1}^{g} \| \beta_k \|_2$ . To define the contextual group lasso, we make the analogous modification to (4):
$$
\min  _ {\mathbf {w}} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} l \left(\sum_ {k = 1} ^ {g} \mathbf {x} _ {i k} ^ {\top} \boldsymbol {\beta} _ {k, \mathbf {w}} (\mathbf {z} _ {i}), y _ {i}\right) \quad \text {s . t .} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \sum_ {k = 1} ^ {g} \| \boldsymbol {\beta} _ {k, \mathbf {w}} (\mathbf {z} _ {i}) \| _ {2} \leq \lambda .
$$
Similar to how the absolute values of the  $\ell_1$ -norm are nondifferentiable at zero, which causes individual explanatory features to be selected, the  $\ell_2$ -norm is nondifferentiable at the zero vector, causing grouped explanatory features to be selected together. To realize the grouped estimator, we adopt the same architecture as before but replace the previous (ungrouped) projection layer with its grouped counterpart. This change demands a different projection algorithm, presented in Appendix B.
2.4 Side constraints
Besides the contextual (group) lasso constraint, our architecture readily accommodates side constraints on  $\beta (\mathbf{z})$  via modifications to the projection. For instance, we follow Zhou and Hooker (2022) in the house pricing example (Figure 1) and constrain the coefficients on the elevator and renovation features to be nonnegative. Such sign constraints reflect domain knowledge that these features should not impact price negatively. Appendix C presents the details and proofs of this extension.
2.5 Pathwise optimization
The lasso regularization parameter  $\lambda$ , controlling the size of the  $\ell_1$ -ball and thus the sparsity level, is typically treated as a tuning parameter. For this reason, algorithms for the lasso usually provide multiple models over a grid of varying  $\lambda$ , which can then be compared (Friedman et al., 2010). Towards this end, it can be computationally efficient to compute the models pathwise by sequentially warm-starting the optimizer. As Friedman et al. (2007) point out, pathwise computation for many  $\lambda$  can be as fast as for a single  $\lambda$ . For the contextual lasso, warm starts also reduce run time compared with initializing at random weights. More importantly, however, pathwise optimization improves the training quality. This last advantage is a consequence of the network's nonconvex optimization surface. Building up a sophisticated network from a simple one helps the optimizer navigate this surface. Appendix D presents our pathwise algorithm and an approach for setting the  $\lambda$  grid.
2.6 Relaxed fit
A possible drawback to the contextual lasso, and indeed all lasso estimators, is bias of the linear model coefficients towards zero. This bias, which is a consequence of shrinkage from the  $\ell_1$ -norm, can help or hinder depending on the data. Typically, bias is beneficial when the number of observations is low or the level of noise is high, while the opposite is true in the converse situation (see, e.g., Hastie et al., 2020). This consideration motivates a relaxation of the contextual lasso that unwinds some, or all, of
the bias imparted by the  $\ell_1$ -norm. We describe an approach in Appendix E that extends the proposal of Hastie et al. (2020) for relaxing the lasso. Their relaxation, which simplifies an earlier proposal by Meinshausen (2007), involves a convex combination of the lasso's coefficients and "polished" coefficients from an unregularized least squares fit on the lasso's selected features. We extend this idea from the lasso's fixed coefficients to the contextual lasso's varying coefficients. The benefits of the relaxation are demonstrated empirically in Appendix E, where we present an ablation study.
2.7 Computational complexity
A forward or backward pass through the vanilla feedforward component of the network takes  $O(md + hd^2 + pd)$  time, where  $h$  is the number of hidden layers,  $d$  is the number of nodes per layer,  $m$  is the number of contextual features, and  $p$  is the number of explanatory features. A forward or backward pass through the projection algorithm takes  $O(p)$  time (Duchi et al., 2008). The time complexity for a pass through the full network over  $n$  observations is thus  $O(nd(m + hd + p))$ . This result suggests that the training time is linear in the sample size  $n$  and number of features  $m$  and  $p$ . Actual training times are reported in Appendix F that demonstrate linear complexity empirically.
2.8 Package
We implement the contextual lasso as described in this section in the Julia (Bezanson et al., 2017) package ContextualLasso. For training the neural network, we use the deep learning library Flux (Innes et al., 2018). Though the experiments throughout this paper involve square or logistic loss functions, our package supports any differentiable loss function, e.g., those in the family of generalized linear models (Nelder and Wedderburn, 1972). ContextualLasso is available at
https://github.com/ryan-thompson/ContextualLasso.jl.