The Contextual Lasso: Sparse Linear Models via Deep Neural Networks
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.
Ryan Thompson, Amir Dezfouli, Robert Kohn
# The Contextual Lasso: Sparse Linear Models via Deep Neural Networks

###### Abstract

Sparse linear models are a gold standard tool 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 we wish to explain the model's predictions, and contextual features, which we wish to determine the model's explanations. This dichotomy leads us to propose the contextual lasso, a new statistical estimator that fits a sparse linear model whose sparsity pattern and coefficients can vary with the contextual features. The fitting process involves learning a nonparametric map, realized via a deep neural network, from contextual feature vector to sparse coefficient vector. 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. Extensive experiments on real and synthetic data suggest 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.

Machine Learning, ICML

## 1 Introduction

Sparse linear models--linear predictive functions in a small subset of features--have a long and rich 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 highly-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; Marcinkevics and Vogt, 2020; Molnar et al., 2020; Rudin et al., 2022). 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 flaws of their own (Laugel et al., 2019; Rudin, 2019).

Motivated by a desire for interpretability and expressivity, this paper focuses on a statistical learning setting where sparse linear models and neural networks can collaborate together. The setting is characterized by a not-uncommon situation where the input features naturally 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. On the other hand, contextual features describe the broader predictive context, e.g., the location of the prediction in time or space, as in the house pricing example below. These inform which explanatory features are relevant and, for those that are, the sign and magnitude of their linear effects. Given this crucial 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\,|\,\mathbf{x},\mathbf{z}]\right)=\sum_{j\in S( \mathbf{x})}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},\ldots,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).1 Via the contextual features, the set-valued function \(S(\mathbf{z})\) encodes the indices of the relevant explanatory features (typically, a small set of \(j\)'s), 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,\ldots,p\}\) for all \(\mathbf{z}\in\mathbb{R}^{m}\). We show throughout the paper that this new model is powerful in various decision-making settings, including energy forecasting and news optimization. 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). Whereas the lasso fits a sparse linear model that fixes the relevant features and 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 coefficient's \(\ell_{1}\)-norm, our regularizer constraints the _expectation_ of the coefficient's \(\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 constrained quadratic program.

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 model with five 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 yields higher floors undesirable, hence the negative effect of floor on price. Beyond the inner city, the floor is irrelevant. Naturally, renovations are valuable everywhere, but more so for older buildings in the inner city than elsewhere. Meanwhile, the number of living rooms and bathrooms are not predictive of price per square meter anywhere, so they remain inactive. The flexibility of the contextual lasso to add or remove attributes by location equips sellers with personalized, readily interpretable linear models containing only attributes most relevant to them.

The rest of the paper is organized as follows. Section 2 introduces the contextual lasso and describes techniques for its computation. Section 3 discusses connections with earlier work. Section 4 reports extensive experimental analyses on synthetic data. Section 5 presents two applications to real-world data. Section 6 concludes the paper.

## 2 Contextual Lasso

This section describes our estimator. To facilitate exposition, we first rewrite the contextually sparse linear model (1) more concisely as

\[g\left(\mathrm{E}[y\,|\,\mathbf{x},\mathbf{z}]\right)=\mathbf{x}^{\top} \boldsymbol{\beta}(\mathbf{z}).\]

The notation \(\boldsymbol{\beta}(\mathbf{z}):=\left(\beta_{1}(\mathbf{z}),\ldots,\beta_{p}( \mathbf{z})\right)^{\top}\) represents a vector coefficient function which is sparse over its domain. That is, for different values of \(\mathbf{z}\), the output of \(\boldsymbol{\beta}(\mathbf{z})\) con

Figure 1: Coefficients as a function of longitude and latitude for the estimated house pricing model. Colored points indicate values of coefficients at different locations. Grey points indicate locations where coefficients are zero.
tains 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}):=\{j:\beta_{j}(\mathbf{z})\neq 0\}\).

### Problem Formulation

The contextual lasso, at the population level, comprises a minimization of the expectation of a loss function subject to an inequality on the expectation of a constraint function:

\[\begin{split}\min_{\mathbf{\beta}\in\mathcal{F}}& \quad\mathrm{E}\left[l\left(\mathbf{x}^{\top}\mathbf{\beta}(\mathbf{z}),y \right)\right]\\ \mathrm{s.\,t.}&\quad\mathrm{E}\left[\|\mathbf{\beta}( \mathbf{z})\|_{1}\right]\leq\lambda,\end{split} \tag{2}\]

where the set \(\mathcal{F}\) is a class of functions that constitute feasible solutions and \(l:\mathbb{R}^{2}\rightarrow\mathbb{R}\) is the loss function, e.g., square loss \(l(z,y)=(y-z)^{2}\) for regression or logistic loss \(l(z,y)=-y\log(z)-(1-y)\log(1-z)\) for classification. Here, the expectations are taken with respect to the random variables \(y\), \(\mathbf{x}\), and \(\mathbf{z}\). The parameter \(\lambda>0\) controls the level of regularization. Smaller values of \(\lambda\) encourage \(\mathbf{\beta}(\mathbf{z})\) to contain more zeros over its domain. Larger values have the opposite effect. The contextual lasso thus differs from the regular lasso, which learns \(\mathbf{\beta}(\mathbf{z})\) as a constant function:

\[\begin{split}\min_{\mathbf{\beta}}&\quad\mathrm{E} \left[l\left(\mathbf{x}^{\top}\mathbf{\beta},y\right)\right]\\ \mathrm{s.\,t.}&\quad\|\mathbf{\beta}\|_{1}\leq\lambda. \end{split}\]

To reiterate the difference: the lasso coaxes _the parameter \(\mathbf{\beta}\)_ towards zero, while the contextual lasso coaxes the _expectation of the function \(\mathbf{\beta}(\mathbf{z})\)_ to zero. The result for the latter is coefficients that can change in value and sparsity with \(\mathbf{z}\).

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:

\[\begin{split}\min_{\mathbf{\beta}\in\mathcal{F}}& \quad\frac{1}{n}\sum_{i=1}^{n}l\left(\mathbf{x}_{i}^{\top}\mathbf{ \beta}(\mathbf{z}_{i}),y_{i}\right)\\ \mathrm{s.\,t.}&\quad\frac{1}{n}\sum_{i=1}^{n}\| \mathbf{\beta}(\mathbf{z}_{i})\|_{1}\leq\lambda.\end{split} \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.2 To operationalize this estimator, we take the function class \(\mathcal{F}\) to be the family of neural networks parameterized by weights \(\mathbf{w}\), denoted \(\mathbf{\beta}_{\mathbf{w}}\). This choice leads to our core proposal:

Footnote 2: The \(\ell_{1}\)-ball is the convex compact set \(\{\mathbf{x}\in\mathbb{R}^{p}:\|\mathbf{x}\|_{1}\leq\lambda\}\).

\[\begin{split}\min_{\mathbf{w}}&\quad\frac{1}{n}\sum_ {i=1}^{n}l\left(\mathbf{x}_{i}^{\top}\mathbf{\beta}_{\mathbf{w}}(\mathbf{z}_{i}),y _{i}\right)\\ \mathrm{s.\,t.}&\quad\frac{1}{n}\sum_{i=1}^{n}\| \mathbf{\beta}_{\mathbf{w}}(\mathbf{z}_{i})\|_{1}\leq\lambda.\end{split} \tag{4}\]

Training a neural network such that its outputs satisfy the \(\ell_{1}\)-constraint is not trivial. We introduce a novel network architecture that addresses this challenge.

### Network Architecture

The neural network architecture--depicted in Figure 2--involves two key components. The first and most straightforward component is a feedforward network \(\mathbf{\eta}(\mathbf{z}):=\left(\eta_{1}(\mathbf{z}),\ldots,\eta_{p}(\mathbf{z} )\right)^{\top}\) comprised of \(p\) fully-connected subnetworks. These subnetworks each tend to the coefficient function of a single explanatory feature. The purpose of the subnetworks is to capture the nonlinear effects of the contextual features on the explanatory features. Since these subnetworks involve only hidden layers with standard affine transformations and nonlinear maps (e.g., rectified linear activations), the coefficients they produce generally do not satisfy the contextual lasso constraint and hence are not sparse. To enforce the constraint, we employ a novel projection layer as the second component of our network.

The projection layer takes the dense coefficients \(\mathbf{\eta}(\mathbf{z})\) from the subnetworks and maps them to sparse coefficients \(\mathbf{\beta}(\mathbf{z}):=\left(\beta_{1}(\mathbf{z}),\ldots,\beta_{p}(\mathbf{ z})\right)^{\top}\) by performing an orthogonal 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 \(\mathbf{\eta}(\mathbf{z}_{1}),\ldots,\mathbf{\eta}(\mathbf{z}_{n})\) together. That is, we take the final sparse coefficients \(\mathbf{\beta}(\mathbf{z}_{1}),\ldots,\mathbf{\beta}(\mathbf{z}_{n})\) as the minimizing arguments of a constrained quadratic program:

\[\begin{split}\mathbf{\beta}(\mathbf{z}_{1}),\ldots,\mathbf{\beta}( \mathbf{z}_{n}):=\\ \underset{\mathbf{\beta}_{1},\ldots,\mathbf{\beta}_{n}:\frac{1}{n}\sum_ {i=1}^{n}\|\mathbf{\beta}_{i}\|_{1}\leq\lambda}{\arg\min}&\quad\frac{1}{ n}\sum_{i=1}^{n}\|\mathbf{\eta}(\mathbf{z}_{i})-\mathbf{\beta}_{i}\|_{2}^{2}.\end{split} \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 (Amos and Kolter, 2017; Agrawal et al., 2019). Yet, to our knowledge, no previous work has studied optimization layers for inducing sparsity.

The program (5) does not admit an analytical solution, though it is solvable by general purpose convex optimization algorithms (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-thresholding the inputs using the computed \(\theta\). Critically, the operations comprising Algorithm 1 are scalable on a GPU. Moreover, we can differentiate through these
operations, allowing end-to-end training of the model.

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):=\operatorname{sign}(x)\max(|x|-\hat{\theta},0)\). The purpose of using the estimate \(\hat{\theta}\) rather than recomputing \(\theta\) via Algorithm 1 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 Algorithm 1 is rerun.

### Side Constraints

Besides the contextual lasso constraint, our architecture accommodates side constraints on \(\mathbf{\beta}(\mathbf{z})\) via modifications to Algorithm 1. For instance, we follow Zhou and Hooker (2022) in the housing example (Figure 1) and constrain the coefficients on the elevator, renovation, living room, and bathroom features to be nonnegative. Such sign constraints reflect domain knowledge that these features should not impact price negatively. Appendix A describes this extension.

### Pathwise Optimization

The lasso regularization parameter \(\lambda\) controlling the size of the \(\ell_{1}\)-ball and thus the sparsity of the coefficient vectors is typically treated as a tuning parameter. For this reason, algorithms for the lasso typically do not return a model for a single value of \(\lambda\) but instead return multiple models with varying \(\lambda\), which can then be compared, e.g., using cross-validation (Friedman et al., 2010). Towards this end, it can be computationally efficient to compute multiple models pathwise by sequentially warm-starting the optimizer. As Friedman et al. (2007) point out, pathwise computation for many values of \(\lambda\) can be as fast as for a single \(\lambda\).

The contextual lasso benefits from pathwise optimization in even more ways than the lasso. Warm starts reduce overall runtime compared with initializing at random weights (running for a sequence of \(\lambda\) is the same order as for a single \(\lambda\)). More importantly, however, pathwise optimization improves the training quality. This last advantage, which is irrelevant for the lasso, 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. Lemhadri et al. (2021) note a similar benefit from pathwise optimization with their regularized neural networks.

In a spirit similar to Friedman et al. (2007), we take the sequence of regularization parameters \(\{\lambda^{(t)}\}_{t=1}^{T}\) as a grid of values that yields a path between the unregularized model (no sparsity) and the fully regularized model (all coefficients zero). Specifically, we set \(\lambda^{(1)}\) such that the contextual lasso regularizer does not impart any regularization, i.e., \(\lambda^{(1)}=n^{-1}\sum_{i=1}^{n}\|\mathbf{\beta}_{\hat{\mathbf{w}}^{(1)}}( \mathbf{z}_{i})\|_{1}\), where the weights \(\hat{\mathbf{w}}^{(1)}\) are a solution to (4) from setting \(\lambda=\infty\). We then construct the sequence as a grid of linearly spaced values between \(\lambda^{(1)}\) and \(\lambda^{(T)}=0\), the latter forcing all coefficients to zero. It is critical here the sequence of \(\lambda^{(t)}\) is decreasing so the optimizer can build on networks that increase in sparsity.

Algorithm 2 summarizes the complete pathwise optimization process outlined above, with gradient descent employed as the optimizer. To parse the notation used in the algorithm, \(L(\mathbf{w};\lambda)=n^{-1}\sum_{i=1}^{n}l(\mathbf{x}_{i}^{\top}\mathbf{\beta}_{ \mathbf{w}}(\mathbf{z}_{i}),y_{i})\) represents
the loss as a function of the network's weights \(\mathbf{w}\) given \(\lambda\), and \(\nabla_{\mathbf{w}}L(\mathbf{w};\lambda)\) represents the associated gradient.

### 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 samples 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 us to consider a relaxation of the contextual lasso that unwinds some, or all, of the bias imparted by the \(\ell_{1}\)-norm. We present here an approach 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. Our proposal extends this idea from the lasso's fixed coefficients to the contextual lasso's varying coefficients.

Denote by \(\hat{\boldsymbol{\beta}}_{\lambda}(\mathbf{z})\) a contextual lasso network fit with regularization parameter \(\lambda\). To unwind bias in \(\hat{\boldsymbol{\beta}}_{\lambda}(\mathbf{z})\), we train a polished network \(\boldsymbol{\beta}_{\lambda}^{p}(\mathbf{z})\) that selects the same explanatory features but does not impose any shrinkage. For this task, we introduce the function \(\hat{\mathbf{s}}_{\lambda}(\mathbf{z}):\mathbb{R}^{m}\rightarrow\{0,1\}^{p}\) that outputs a vector with elements equal to one wherever \(\hat{\boldsymbol{\beta}}_{\lambda}(\mathbf{z})\) is nonzero and elsewhere is zero. We then fit the polished network as \(\boldsymbol{\beta}_{\lambda}^{p}(\mathbf{z})=\boldsymbol{\eta}(\mathbf{z}) \circ\hat{\mathbf{s}}_{\lambda}(\mathbf{z})\), where \(\circ\) means element-wise multiplication and \(\boldsymbol{\eta}(\mathbf{z})\) is the same architecture as used for the original contextual lasso network before the projection layer. The effect of including \(\hat{\mathbf{s}}_{\lambda}(\mathbf{z})\), which is fixed when training \(\boldsymbol{\beta}_{\lambda}^{p}(\mathbf{z})\), is twofold. First, it guarantees the coefficients from the polished network are nonzero in the same positions as the original network, i.e., the same features are selected. Second, it ensures explanatory features only contribute to gradients for samples in which they are active, i.e., \(x_{ij}\) does not contribute if the \(j\)th component of \(\hat{\mathbf{s}}_{\lambda}(\mathbf{z}_{i})\) is zero. Because the polished network does not project onto an \(\ell_{1}\)-ball, its coefficients are not shrunk.

To arrive at the relaxed contextual lasso fit, we convexly combine \(\hat{\boldsymbol{\beta}}_{\lambda}(\mathbf{z})\) and the fitted polished network \(\hat{\boldsymbol{\beta}}_{\lambda}^{p}(\mathbf{z})\):

\[\hat{\boldsymbol{\beta}}_{\lambda,\gamma}(\mathbf{z}):=(1-\gamma)\hat{ \boldsymbol{\beta}}_{\lambda}(\mathbf{z})+\gamma\hat{\boldsymbol{\beta}}_{ \lambda}^{p}(\mathbf{z}),\quad 0\leq\gamma\leq 1. \tag{6}\]

When \(\gamma=0\), we recover the original biased coefficients, and when \(\gamma=1\), we attain the unbiased polished coefficients. Between these extremes lies a continuum of relaxed coefficients with varying degrees of bias. Since the original and polished networks need only be computed once, we may consider any relaxation on this continuum at virtually no additional computational expense. In practice, we choose among the possibilities by tuning \(\gamma\) on a validation set.

### Package

We implement the contextual lasso and its optimization strategy 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 used throughout the entire family of generalized linear models (Nelder and Wedderburn, 1972). ContextualLasso will be available open source on GitHub shortly.

## 3 Related Work

Contextual explanation networks (Al-Shedivat et al., 2020) can be a considered a cousin of the contextual lasso. These neural networks input contextual features and output an interpretable model in explanatory features. They include non-sparse contextual linear models. Al-Shedivat et al. (2020) implement the contextual linear model as a weighted combination of finitely many individual linear models. Though sparsity is not the focus of their work, they add a small amount of \(\ell_{1}\)-regularization to the individual models to prevent overfitting. This type of regularization is fundamentally different from that studied here since no mechanism encourages the network to combine these sparse models such that the combined model remains sparse. In contrast, the contextual lasso guides the network towards sparse models by directly regularizing the sparsity of the models it produces.

The contextual lasso is also related to several estimators that allow for sparsity patterns that vary by sample. Yamada et al. (2017) devised the first of these estimators--the localized lasso--which fits a linear model with a different coefficient vector for each sample. The coefficients are sparsified using
a lasso regularizer that relies on the availability of graph information to link the samples. Yang et al. (2022) and Yoshikawa & Iwata (2022) followed with neural networks that produce linear models with varying sparsity patterns via gating mechanisms. These approaches are quite distinct from our own, however. First, the sparsity patterns vary with every feature, severely restricting the interpretability of the output linear models. Second, the sparsity level or nonzero coefficients are fixed across samples, making them unsuitable for the contextual setting where both may vary.

More broadly, our work advances the literature at the intersection of feature sparsity and neural networks, an area that has gained momentum over the last few years. See, e.g., the lassonet of Lemhadri et al. (2021a;b) which selects features in a residual neural network using an \(\ell_{1}\)-regularizer on the skip connection. This regularizer is combined with constraints that force a feature's weights on the first hidden layer to zero whenever its skip connection is zero. See also Scardapane et al. (2017) and Feng & Simon (2019) for earlier ideas based on the group lasso, and Chen et al. (2021) for other approaches. Though related, these methods differ from the contextual lasso in that they involve uninterpretable neural networks with fixed sparsity patterns. The underlying optimization problems also differ--whereas these methods regularize the network's weights, ours regularizes its output.

## 4 Experimental Analysis

The properties of the contextual lasso are evaluated here via experimentation on synthetic data. As benchmark methods, we consider a nonsparse contextual linear model (i.e., no projection layer) and a deep neural network as a function of all contextual and explanatory features. We also compare against the lasso and a pairwise lasso with main effects plus pairwise interactions between the explanatory and contextual features. Appendix B contains implementation details.

### Data Generation

The explanatory features \(\mathbf{x}_{1},\ldots,\mathbf{x}_{n}\) are generated iid as \(p\)-dimensional \(N(\mathbf{0},\boldsymbol{\Sigma})\) random variables, where the covariance matrix \(\boldsymbol{\Sigma}\) has elements \(\Sigma_{ij}=0.5^{[i-j]}\). The contextual features \(\mathbf{z}_{1},\ldots,\mathbf{z}_{n}\) are generated iid as \(m\)-dimensional random variables uniform on \([-1,1]^{m}\), independent of the \(\mathbf{x}_{i}\). With the features drawn, we simulate a regression response:

\[y_{i}\sim N(\mu_{i},1),\quad\mu_{i}=\kappa\cdot\mathbf{x}_{i}^{\top} \boldsymbol{\beta}(\mathbf{z}_{i}),\]

or a classification response via a logistic function:

\[y_{i}\sim\mathrm{Bernoulli}(p_{i}),\quad p_{i}=\frac{1}{1+\exp\left(-\kappa \cdot\mathbf{x}_{i}^{\top}\boldsymbol{\beta}(\mathbf{z}_{i})\right)},\]

for \(i=1,\ldots,n\). Here, \(\kappa>0\) is a parameter controlling the signal strength vis-a-vis the variance of \(\kappa\cdot\mathbf{x}_{i}^{\top}\boldsymbol{\beta}(\mathbf{z}_{i})\). We first estimate the variance of \(\mathbf{x}_{i}^{\top}\boldsymbol{\beta}(\mathbf{z}_{i})\) on the training set and then set \(\kappa\) so the variance of the signal is five. The coefficient function \(\boldsymbol{\beta}(\mathbf{z}_{i}):=\left(\beta_{1}(\mathbf{z}_{i}),\ldots, \beta_{p}(\mathbf{z}_{i})\right)^{\top}\) is constructed such that \(\beta_{j}(\mathbf{z}_{i})\) maps to a nonzero value whenever \(\mathbf{z}_{i}\) lies within a hypersphere of radius \(r_{j}\) centered at \(\mathbf{c}_{j}\):

\[\beta_{j}(\mathbf{z}_{i})=\begin{cases}1-\frac{1}{2r_{j}}\|\mathbf{z}_{i}- \mathbf{c}_{j}\|_{2}&\text{if }\|\mathbf{z}_{i}-\mathbf{c}_{j}\|_{2}\leq r_{j}\\ 0&\text{otherwise}\end{cases}. \tag{7}\]

This function attains the maximal value one when \(\mathbf{z}_{i}=\mathbf{c}_{j}\) and the minimal value zero when \(\|\mathbf{z}_{i}-\mathbf{c}_{j}\|_{2}>r_{j}\). The centers \(\mathbf{c}_{1},\ldots,\mathbf{c}_{p}\) are generated with uniform probability on \([-1,1]^{p}\), and the radii \(r_{1},\ldots,r_{p}\) are chosen to achieve sparsity levels that vary between 0.05 and 0.15 (average 0.10 across all features). Figure 3 provides a visual illustration.

This coefficient function is inspired by the house pricing example in Figure 1, where the coefficients are typically nonzero in one central region of the contextual feature space.

### Statistical Performance

As a prediction metric, we report the square or logistic loss relative to the intercept-only model. As an interpretability metric, we report the proportion of nonzero features. As a selection metric, we report the F1-score of the selected features; a value of one indicates perfect feature selection (all true nonzeros selected and no false positives).3 All three metrics are evaluated on a testing set with hyperparameters tuned on a validation set, both constructed by drawing \(n\) samples independently and identically to the training set.

Footnote 3: The \(\mathrm{F1}\)-score \(:=2\,\mathrm{TP}\,/(2\,\mathrm{TP}+\mathrm{FP}+\mathrm{FN})\), where \(\mathrm{TP}\), \(\mathrm{FP}\), and \(\mathrm{FN}\) are the number of true positive, false positive, and false negative selections.

We consider three different settings of increasing complexity: (1) \(p=10\) and \(m=2\), (2) \(p=50\) and \(m=2\), and (3) \(p=50\) and \(m=5\). Within each setting, the sample size ranges from \(n=10^{2}\) to \(n=10^{5}\). Figure 4 reports regression results for these settings over 10 independent replications. Appendix C reports the classification results.

Due to its regularizer, the contextual lasso performs comparably with the lasso, pairwise lasso, and deep neural network

Figure 3: Illustration of coefficient function (7) for \(p=3\) explanatory features and \(m=2\) contextual features. The centers \(\mathbf{c}_{j}\) correspond to the dark red in the middle of each sphere.
when the sample size is small. On the other hand, the contextual linear model (the contextual lasso's unregularized counterpart) can perform poorly here. As \(n\) increases, the contextual lasso begins to outperform other methods in prediction, interpretability, and selection. Eventually, it learns the correct map from contextual features to relevant explanatory features, recovering only the true nonzeros. Though its unregularized counterpart performs nearly as well in terms of prediction for large \(n\), it remains much less interpretable, using all explanatory features. In contrast, the contextual lasso uses just 10% of the explanatory features on average.

Unlike the contextual lasso, the deep neural network performs about as well as the regular lasso for most \(n\). Only for large sample sizes does it begin to approach the prediction performance of the contextual lasso. The two methods should predict equally well for large enough \(n\), though the function learned by the deep neural network remains opaque. The regular lasso makes small gains with increasing sample size. Adding pairwise interactions between the explanatory and contextual features yields a modest improvement to prediction accuracy. Nonetheless, the lasso lacks the expressive power of the contextual lasso needed to adapt to the complex sparsity pattern underlying the true model.

## 5 Data Analyses

The contextual lasso is now applied to model two real datasets. Appendix D has links to the datasets.

### Energy Consumption Regression Dataset

The first dataset contains measurements of energy use over five months for a low-energy home in Mons, Belgium (Candaredo et al., 2017). Besides this continuous response feature, the dataset also contains \(p=27\) explanatory features in the form of temperature and humidity readings in different rooms of the house and local weather data. We define several contextual features from the time stamp to capture seasonality: month of year, day of week, hour of day, and an indicator for the weekend. To reflect their cyclical nature, the first three contextual features are transformed using sine and cosine functions, leading to \(m=7\) contextual features.

The dataset, containing \(n=19,375\) samples, is randomly split into training, validation, and testing sets in 0.6-0.2-0.2 proportions. We repeat this random split of the data 10 times, each time recording performance on the testing set, and report the aggregate results in Table 1. As performance metrics, we consider the relative loss as in Section 4 and the mean number of nonzero features. Among all methods, the contextual lasso leads to the lowest test loss, outperforming even the deep neural network. Importantly, this excellent prediction performance is achieved while maintaining a high level of interpretability. In contrast to the deep neural network and the contextual linear model, which use all available explanatory features, the predictions from the contextual lasso arise from linear models containing just 2.4 explanatory features on average! These linear models are

Figure 4: Comparisons of methods for regression over 10 synthetic datasets. Solid points represent averages and error bars denote standard errors. Dashed horizontal lines in the middle row of plots indicate the true sparsity level. Since relative loss for the contextual linear model can be large for small \(n\), we omit it from the plots in some cases to maintain the aspect ratio.
also much simpler than those from the regular lasso, which typically involve more than four times as many features.

The good predictive performance of the contextual lasso suggests a seasonal pattern of sparsity. To investigate this phenomenon, we apply the fitted model to a randomly sampled testing set and plot the resulting sparsity levels as a function of the hour of day in Figure 5.

The model is typically highly sparse in the late evening and early morning. Between 11 pm and 5 am, the median sparsity level is no more than 5%. There is likely little or no activity inside the house at these times, so sensor readings from within the house--which constitute the majority of the explanatory features--are irrelevant. The number of active explanatory features rises later in the day, reaching a peak at around 6 pm. Overall, a major benefit of the contextual lasso, besides its good predictions, is the ability to identify a parsimonious set of factors driving energy use at any given time of day.

### News Popularity Classification Dataset

The second dataset consists of articles posted to the news platform Mashable (Fernandes et al., 2015). The task is to predict if an article will be popular, defined in Fernandes et al. (2015) as more than 1400 shares. In addition to the zero-one response feature for popularity, the dataset has predictive features that quantify the articles (e.g., number of total words, positive words, and images). The data channel feature, which identifies the category of the article (lifestyle, entertainment, business, social media, technology, world, or viral), is taken as the contextual feature. It is expressed as a sequence of indicator variables yielding \(m=6\) contextual features. There remain \(p=51\) explanatory features.

Table 2 reports the results over 10 random splits of the dataset (\(n=39,643\)) into training, validation, and testing sets in the same proportions as before.

In contrast to the previous dataset, all the methods predict similarly well here. The deep neural network performs marginally best overall, while the lasso performs marginally worst. Though predicting neither best nor worst, the contextual lasso retains a significant lead in terms of sparsity, being twice as sparse as the next best method. Sparsity is crucial for this task as it allows the author to focus on a small number of changes necessary to improve the article's likelihood of success. The uninterpretable deep neural network, or fully dense contextual linear model, are not nearly as useful here.

## 6 Concluding Remarks

Contextual sparsity is an important extension of the classical notion of feature sparsity. Rather than fix the relevant features once and for all, contextual sparsity allows feature relevance to depend on the prediction context. To tackle this intricate statistical learning problem, we devise the contextual lasso. This new estimator utilizes the expressive power of deep neural networks to learn interpretable sparse linear models with sparsity patterns that vary with the contextual features. The optimization problem of the contextual lasso is solvable at scale using modern deep learning frameworks, and we make our implementation open source. An extensive experimental analysis of the new estimator illustrates its good prediction, interpretation, and selection properties. To the best of our knowledge, the contextual lasso is the only available tool for handling the contextually sparse setting.

One direction that continues this line of work is to combine contextual sparsity with classic feature sparsity. This blended sparsity would be suitable for removing explanatory features that are not predictive in any context. Another direction is to extend the notion of contextual sparsity beyond the lasso to other sparsity-inducing regularizers. Other regularizers can, however, give rise to computational difficulties beyond those considered here, e.g., points of discontinuity.

\begin{table}
\begin{tabular}{l l l} \hline \hline  & Relative loss & Mean nonzeros \\ \hline Deep neural network & \(0.406\pm 0.004\) & - \\ Contextual linear model & \(0.353\pm 0.006\) & \(25.0\pm 0.0\) \\ Lasso & \(0.694\pm 0.004\) & \(10.8\pm 0.4\) \\ Pairwise lasso & \(0.589\pm 0.003\) & \(24.5\pm 0.2\) \\ Contextual lasso & \(0.330\pm 0.005\) & \(2.4\pm 0.2\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Comparisons of methods for modeling energy consumption. Metrics are aggregated over 10 random splits of the data. Averages and standard errors are reported.

\begin{table}
\begin{tabular}{l l l} \hline \hline  & Relative loss & Mean nonzeros \\ \hline Deep neural network & \(0.899\pm 0.002\) & - \\ Contextual linear model & \(0.908\pm 0.003\) & \(51.0\pm 0.0\) \\ Lasso & \(0.916\pm 0.002\) & \(21.2\pm 0.9\) \\ Pairwise lasso & \(0.910\pm 0.002\) & \(29.3\pm 0.5\) \\ Contextual lasso & \(0.910\pm 0.002\) & \(9.4\pm 0.6\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Comparisons of methods for modeling news popularity. Metrics are aggregated over 10 random splits of the data. Averages and standard errors are reported.

Figure 5: Explanatory feature sparsity as a function of hour of day for the estimated energy consumption model. The sparsity level varies within each hour because the other contextual features vary.