SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\tilde {\eta} _ {i j} = \left\{ \begin{array}{l l} 0 & i f \eta_ {i j} <   0 \wedge j \in \mathcal {P} \\ 0 & i f \eta_ {i j} > 0 \wedge j \in \mathcal {N}, \quad i = 1, \ldots , n, j = 1, \ldots , p. \\ \eta_ {i j} & o t h e r w i s e \end{array} \right.
$$
Then optimization problem (6) admits the same optimal solution as
$$
\min  _ {\boldsymbol {\beta} _ {1}, \dots , \boldsymbol {\beta} _ {n}} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \| \tilde {\boldsymbol {\eta}} _ {i} - \boldsymbol {\beta} _ {i} \| _ {2} ^ {2} \tag {7}
$$
$$
\text {s . t .} \quad \frac {1}{n} \sum_ {i = 1} ^ {n} \| \beta_ {i} \| _ {1} \leq \lambda .
$$
The proof of this proposition requires the following lemma.
Lemma C.2. A solution  $\beta_{1},\ldots ,\beta_{n}$  to (6) must satisfy the inequality  $\eta_{ij}\beta_{ij}\geq 0$  for all  $i$  and  $j$
Proof. We proceed using proof by contradiction along the lines of Lemma 3 in Duchi et al. (2008). Suppose there exists a solution  $\beta_{1},\ldots ,\beta_{n}$  such that  $\eta_{ij}\beta_{ij} < 0$  for some  $i$  and  $j$ . Take  $\beta_1^\star ,\dots ,\beta_n^\star$  equal to  $\beta_{1},\ldots ,\beta_{n}$  except at index  $(i,j)$ , where we set  $\beta_{ij}^{\star} = 0$ . Note that  $\beta_1^\star ,\dots ,\beta_n^\star$  continues to satisfy the sign constraints and  $\ell_1$ -constraint, and hence remains feasible for (6). We also have that
$$
\sum_ {i = 1} ^ {n} \left\| \boldsymbol {\eta} _ {i} - \boldsymbol {\beta} _ {i} \right\| _ {2} ^ {2} - \sum_ {i = 1} ^ {n} \left\| \boldsymbol {\eta} _ {i} - \boldsymbol {\beta} _ {i} ^ {\star} \right\| _ {2} ^ {2} = \beta_ {i j} ^ {2} - 2 \eta_ {i j} \beta_ {i j} > \beta_ {i j} ^ {2} > 0.
$$
Thus,  $\beta_1^\star, \ldots, \beta_n^\star$  attains a lower objective value than the solution  $\beta_1, \ldots, \beta_n$ . This contradiction yields the statement of the lemma.
The proof of Proposition C.1 now follows.
Proof. If  $\eta_{ij} < 0$  and  $j \in \mathcal{P}$ , then by Lemma C.2, a solution to (6) must satisfy  $\beta_{ij} = 0$ . Likewise, if  $\eta_{ij} > 0$  and  $j \in \mathcal{N}$ , then it must hold that  $\beta_{ij} = 0$ . Hence, by setting any  $\eta_{ij}$  that violate the sign constraints to zero (i.e.,  $\tilde{\eta}_{ij}$ ), and noting that a solution to (7) must satisfy  $\beta_{ij} = 0$  when  $\tilde{\eta}_{ij} = 0$ , we arrive at the result of the proposition.
D Pathwise optimization
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}\|\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. Linear spacing of the sequence of  $\lambda^{(t)}$  generally yields linearly spaced sparsity levels. The sequence should be decreasing so the optimizer can build on networks that increase in sparsity.
Algorithm 3 summarizes the complete pathwise optimization process, 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 \boldsymbol {\beta}_{\mathbf{w}}(\mathbf{z}_i),y_i)$  is the loss as a function of the network's weights w given  $\lambda$ , and  $\nabla_{\mathbf{w}}L(\mathbf{w};\lambda)$  is its gradient.
E Relaxed fit
Denote by  $\hat{\beta}_{\lambda}(\mathbf{z})$  a contextual lasso network fit with regularization parameter  $\lambda$ . To unwind bias in  $\hat{\beta}_{\lambda}(\mathbf{z})$ , we train a polished network  $\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} \to \{0,1\}^{p}$  that outputs a vector with elements equal to one wherever  $\hat{\beta}_{\lambda}(\mathbf{z})$  is nonzero and elsewhere is zero. We then fit the
Algorithm 3 Pathwise optimization
input Initial weights  $\hat{\mathbf{w}}^{(0)}$  , step size  $\alpha$  , and number of regularization parameters  $T$    
Initialize  $\lambda^{(1)} = \infty$    
for  $t = 1,\dots ,T$  do Initialize  $\mathbf{w}_{(0)} = \hat{\mathbf{w}}^{(t - 1)}$  Initialize  $m = 0$  while Not converged do Update  $\mathbf{w}_{(m + 1)} = \mathbf{w}_{(m)} - \alpha \cdot \nabla_{\mathbf{w}}L(\mathbf{w}_{(m)};\lambda^{(t)})$  Update  $m = m + 1$  end while Set  $\hat{\mathbf{w}}^{(t)} = \mathbf{w}_{(m)}$  if  $t = 1$  then Set  $\lambda^{(1)} = n^{-1}\sum_{i = 1}^{n}\| \beta_{\hat{\mathbf{w}}^{(1)}}(\mathbf{z}_i)\| _1$  and  $\lambda^{(T)} = 0$  Equispace  $\lambda^{(2)},\ldots ,\lambda^{(T - 1)}$  between  $\lambda^{(1)}$  and  $\lambda^{(T)}$  end if   
end for   
output Fitted weights  $\hat{\mathbf{w}}^{(1)},\dots ,\hat{\mathbf{w}}^{(T)}$
polished network as  $\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  $\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 observations 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 combine  $\hat{\beta}_{\lambda}(\mathbf{z})$  and the fitted polished network  $\hat{\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 {8}
$$
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 computational expense over and above that of the two networks. In practice, we choose among the possibilities by tuning  $\gamma$  on a validation set.
To illustrate the benefits of relaxation, we present Figure 6, which compares the relaxed and nonrelaxed variants of the contextual lasso under the synthetic experimental design of Section 4. As far as prediction accuracy is concerned, the benefits of relaxation are marginal. However, for selection and interpretation, the story is quite different. The relaxation yields models that are both sparser and contain more true positives and/or fewer false positives. These gains are most pronounced in larger samples. In smaller samples, relaxation is less beneficial because the bias from shrinkage helps stabilize the models. Yet, the relaxed variant of the contextual lasso typically does no worse than its nonrelaxed counterpart because it adapts to the best level of bias (shrinkage) by tuning  $\gamma$ .
F Run times
The complexity analysis in Section 2.7 suggests that the training time for the contextual lasso should be linear in the sample size  $n$  and number of explanatory features  $p$ . To verify this result, we record the time taken to fit the contextual lasso over a sequence of 50 values of  $\lambda$ . The number of hidden layers is three and the number of neurons per layer is 100. Figure 7 reports the results. The run times are indeed linear as a function of  $n$  and  $p$  and, overall, quite reasonable for real world applications.
G Localized lasso
We compare the contextual lasso with the localized lasso of Yamada et al. (2017). The graph information required by the localized lasso is estimated from the contextual features using the nearest neighbors approach of Yamada et al. (2017). We focus on smaller sample sizes than in the main
Figure 6: Relaxation ablation on synthetic regression data. Metrics are aggregated over 10 synthetic datasets. Solid points are averages and error bars are standard errors. Dashed horizontal lines in the middle row indicate the true sparsity level.
Figure 7: Run time in seconds to fit the contextual lasso over a sequence of 50 values of  $\lambda$ , measured over 10 synthetic datasets. The number of explanatory features  $p = 50$  in the left plot and the sample size  $n = 1,000$  in the right plot. The number of contextual features  $m = 5$ . Solid points are averages and error bars are one standard errors.
experiments since each observation requires a new coefficient vector that each constitute additional optimization variables. Figure 8 reports the results for regression (the localized lasso does not support classification). The localized lasso's prediction loss improves with growing  $n$ , but not nearly as well as the contextual lasso. The graph information may be insufficient to encode the underlying nonlinearity. Furthermore, the localized lasso's regularizer never induces fully sparse coefficient vectors (all zeros), which may be limiting if there are no relevant explanatory features for certain  $\mathbf{z}$ .
H Implementation details
The contextual lasso is fit using our Julia package ContextualLasso. The network is configured with three hidden layers. The number of neurons, which are spread equally across these hidden layers, is set so that the dimensionality of the weights  $\mathbf{w}$  is approximately  $32 \times p \times m$ . This setting ensures the network size scales roughly linearly with the number of features. The contextual linear model uses the same architecture, excluding the projection layer. The deep neural network is set up similarly. These methods all use rectified linear activation functions in the hidden layers and are optimized
Figure 8: Comparisons with the localized lasso on synthetic regression data. Metrics are aggregated over 10 synthetic datasets. Solid points are averages and error bars are standard errors. The dashed horizontal line in the middle indicates the true sparsity level.
using Adam (Kingma and Ba, 2015) with a learning rate of 0.001. Convergence is monitored on a validation set with the optimizer terminated after 30 iterations without improvement.
The lasso is fit using the Julia package GLMNet (Friedman et al., 2010). The group lasso is fit using the R package grpred (Breheny and Huang, 2015). Since contextual features are always relevant, the regularizer for the lasso is applied only to explanatory features and interactions, not the contextual features. The (group) lasso and contextual lasso all allow for relaxed fits, as discussed in Section 2.6. The regularization parameter  $\lambda$  is swept over a grid of 50 values computed automatically from the data using ContextualLasso, GLMNet, or grpred. For each value of  $\lambda$ , the relaxation parameter  $\gamma$  is swept over the grid  $\{0, 0.1, \dots, 1\}$ .
The lassonet is fit using the Python package lassonet (Lemhadri et al., 2021b), which also performs its own relaxation. LLSPIN is fit using the authors' Python implementation (see Yang et al., 2022). Their implementation relies on the hyperparameter optimization framework Optuna (Akiba et al., 2019) to determine the regularization parameter and learning rate. We use the default grid for tuning the learning rate, but increase the maximum regularization parameter to 10, which is roughly the smallest value required to achieve a fully sparse solution in our experiments. LLSPIN does not shrink and so does not admit a relaxation. Lassonet and LLSPIN use the same convergence criterion, number of regularization parameters, and network configuration as the other deep learning methods.
For all methods, the input features are standardized prior to training. Standardization of the explanatory features is necessary for the lasso estimators as it places all coefficients on the same scale, ensuring equitable regularization. ContextualLasso automates standardization and expresses all final coefficients on their original scale.
All experiments are run on a Linux platform with NVIDIA RTX 4090 GPUs.
I Synthetic data generation
The explanatory features  $\mathbf{x}_1, \ldots, \mathbf{x}_n$  are generated iid as  $p$ -dimensional  $N(\mathbf{0}, \mathbf{\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 as
$$
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 as
$$
y _ {i} \sim \operatorname {B e r n o u l l i} (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,\dots ,n$ . Here,  $\kappa >0$  controls the signal strength vis-à-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
9grpreg does not provide support for this functionality.
Figure 9: Illustration of coefficient function (9) 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.
is five. The coefficient function  $\beta (\mathbf{z})\coloneqq \left(\beta_{1}(\mathbf{z}),\dots ,\beta_{p}(\mathbf{z})\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} \left(\mathbf {z} _ {i}\right) = \left\{ \begin{array}{l l} 1 - \frac {1}{2 r _ {j}} \left\| \mathbf {z} _ {i} - \mathbf {c} _ {j} \right\| _ {2} & \text {i f} \| \mathbf {z} _ {i} - \mathbf {c} _ {j} \| _ {2} \leq r _ {j} \\ 0 & \text {o t h e r w i s e} \end{array} . \right. \tag {9}
$$
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,\dots,r_p$  are chosen to achieve sparsity levels that vary between 0.05 and 0.15 (average 0.10). Figure 9 provides a visual illustration. This function is inspired by the house pricing example in Figure 1, where the coefficients are nonzero in some central region of the contextual feature space.
J Classification results
J.1 Synthetic data
Figure 10 reports the experimental results for classification on synthetic data, analogous to those for regression in Section 4. The findings are broadly in line with the regression ones. The contextual lasso performs best overall and is the only method ever able to recover the true nonzeros accurately.
J.2 News popularity data
We turn to a real dataset 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 4 reports the results over 10 random splits of the dataset's  $n = 39,643$  observations into training, validation, and testing sets in 0.6-0.2-0.2 proportions. In contrast to the previous datasets, all
Table 4: Comparisons on the news popularity data. Metrics are aggregated over 10 random splits of the data. Averages and standard errors are reported.
Relative loss	Avg. sparsity
Deep neural network	0.903 ± 0.003	51.0 ± 0.0
Contextual linear model	0.906 ± 0.003	51.0 ± 0.0
Lasso	0.914 ± 0.002	22.4 ± 0.7
Lassonet	0.894 ± 0.002	50.7 ± 0.3
LLSPIN	0.923 ± 0.009	51.0 ± 0.0
Contextual lasso	0.906 ± 0.002	12.6 ± 0.9
methods predict similarly well here. The lassonet performs marginally best overall, while LLSPIN performs marginally worst. Though predicting neither best nor worst, the contextual lasso retains a
Deep neural network Lasso LLSPIN Contextual linear model Lassonet Contextual lasso
Figure 10: Comparisons on synthetic classification data. Metrics are aggregated over 10 synthetic datasets. Solid points are averages and error bars are standard errors. Dashed horizontal lines in the middle row indicate the true sparsity level.
significant lead in terms of sparsity, being twice as sparse as the next sparsest method (lasso). Sparsity is crucial for this task as it allows the author or editor to focus on a small number of changes necessary to improve the article's likelihood of success. Other methods such as the uninterpretable deep neural network, or fully dense contextual linear model, are not nearly as useful for the same purpose.
K High-dimensional and fixed coefficient results
Besides interpretability, a major appeal of the lasso is its good performance in high-dimensional regimes, where the number of features is comparable to the sample size. It is intriguing to consider whether the contextual lasso remains useful in this setting. To this end, we extend the experiments of Section 4 to  $p = 1,000$  explanatory features. Typically, when there are so many features, only a small number are relevant for predicting the response, so we adjust the synthetic data generation process so that only 10 explanatory features are relevant for some values of the contextual features  $\mathbf{z}$ . The remaining 990 explanatory features remain irrelevant for all  $\mathbf{z}$ . Figure 11 reports the results. The contextual lasso performs highly competitively, even when  $n < 1,000$  and there are more explanatory features than observations. As with the lower-dimensional experiments, the contextual lasso can still filter out the irrelevant explanatory features and achieves near-perfect support recovery for large  $n$ .
A second regime that might arise in practice is where contextual features are present but have no effect on the explanatory features. That is, the explanatory features have a fixed sparsity pattern and fixed coefficients for all  $\mathbf{z}$ . Figure 12 reports the results in this setting. Here, the contextual lasso is outperformed by the lasso (which is in its home territory) and the lassonet, both of which assume the sparsity pattern is fixed. Nonetheless, the contextual lasso is able to recover the true nonzeros by learning a constant function for  $\beta (\mathbf{z})$  via the bias term in each neuron. Provided the sample size is sufficiently large, it is reasonable to expect the contextual lasso to remain competitive with the lasso.
Figure 11: Comparisons on high-dimensional synthetic regression data. Metrics are aggregated over 10 synthetic datasets. Solid points are averages and error bars are standard errors. The dashed horizontal line in the middle indicates the true sparsity level.
Figure 12: Comparisons on fixed coefficient synthetic regression data. Metrics are aggregated over 10 synthetic datasets. Solid points are averages and error bars are standard errors. The dashed horizontal line in the middle indicates the true sparsity level.
L Stability analysis
Since the contextual lasso involves a highly-parameterized neural network, it is insightful to consider its selection stability when trained using different random initializations of the network weights. To this end, we report Figure 13, which shows selection instability as measured by the Hamming distance (scaled by  $p$ ) between two independently initialized networks. Unsurprisingly, the statistically difficult
Figure 13: Selection instability of the contextual lasso over 10 synthetic datasets. Solid points represent averages and error bars denote standard errors.
regimes, where the sample size is small or the number of explanatory/contextual features is large, correspond to higher instability. In any case, the contextual lasso grows increasingly stable with the sample size. Even when  $p = 50$  and  $m = 5$ , the contextual lasso is almost entirely stable by the time  $n = 10,000$ , which is typical of our real data examples. The concave shape of the stability curve
when  $p = 50$  is consistent with the proportion of nonzeros reported in Figure 3, where the contextual lasso initially produces highly sparse models that are more restricted and hence more stable.
M Dataset availability
The datasets used throughout this paper are publicly available at the following URLs.
- House pricing data in Section 1: https://www.kaggle.com/datasets/ruiqurm/lianjia.
- Energy consumption data in Section 4: https://archive.ics.uci.edu/ml/datasets/Apliances+energy+prediction.
- Parkinson's telemonitoring data in Section 4: https://archive.ics.uci.edu/ml/datasets/parkinsons+telemonitoring.
- News popularity data in Appendix J: https://archive.ics.uci.edu/ml/datasets/online+news+popularity.
The contextual lasso has strong prediction, interpretation, and selection properties. Naturally, it also has several weaknesses. As with the ordinary lasso, the contextual lasso may select a single feature from a group of highly correlated features. This shortcoming could be remedied by introducing an  $\ell_2^2$ -regularizer in a spirit similar to the elastic net (Zou and Hastie, 2005). Another potential drawback is that the contextual lasso does not guarantee the complete exclusion of an explanatory feature. That is, it cannot ensure the coefficient for a feature is nonzero for every possible  $\mathbf{z}$ , which might limit interpretability in certain settings. Full exclusion could be achieved by adding an  $\ell_1$ -regularizer on the weights of the first layer, similar to the lassonet. Finally, as is typical of neural networks, the objective function of the contextual lasso is nonconvex, which can complicate analysis of its properties.