Avoid length-n allocations when evaluating the ProxGrad objective

[matthew-liu-1]

Removes most length-n allocations from the ProxGrad objective. The fitted coefficients do not change.

• Lasso, X of size (100000, 200), FISTA(): 470 ms → 377 ms, and 184 MiB → 2.6 MiB.

Note: I've generated this PR with Claude assisting me. I've thoroughly & manually validated the effort of this change.

Why

• We run many LassoRegression fits with ProxGrad inside a larger pipeline, often with a high max_iter.
• In a profile, these allocations were more than half of all the allocations in the pipeline.

Where the allocations come from

• The gradient (smooth_fg!) already reuses scratch.n.
• The backtracking check calls smooth_objective(glr, X, y). For L2 loss, each call allocates:
  • apply_X: X * θ[1:p], then .+ θ[end]
  • LPLoss: y .- ŷ
  • lp: abs2.(v)
• The loss value that smooth_fg! returns allocates one more abs2.(r).
• For a lasso, this is about seven length-n vectors per iteration.

What changes

• New smooth_objective!(glr, X, y, scratch). It writes X*θ into scratch.n2.
• For L2Loss, a specific method forms the residuals in place with get_residuals!.
• ProxGrad uses smooth_objective! when c == 0. The multinomial path (c > 0) does not change.
• The L2 loss value is computed in place as sum(r .= abs2.(r)) / 2:
  • sum(abs2, r) is neater, but it rounds differently from sum(abs2.(r)) in the last bit.
  • That difference can change a backtracking step, so I kept the same summation. The results do not change.
• For the logistic and robust losses, only the X*θ allocations go. The losses themselves still allocate.

Length-n vectors that one objective call allocates (L1 penalty, with intercept):

Loss	Before	After
L2 (lasso)	4	0
Logistic	3	1
Huber	5	3
Quantile	4	2

Benchmark

X of size (100000, 200), FISTA().

LassoRegression(0.1):

before:  Time (median): 469.875 ms   GC (median): 8.25%   Memory estimate: 184.30 MiB, allocs estimate: 2093.
after:   Time (median): 377.069 ms   GC (median): 0.00%   Memory estimate: 2.61 MiB, allocs estimate: 1148.

ElasticNetRegression(0.1, 0.1; scale_penalty_with_samples=false):

before:  Time (median): 445.273 ms   GC (median): 6.40%   Memory estimate: 188.25 MiB, allocs estimate: 2761.
after:   Time (median): 348.912 ms   GC (median): 0.00%   Memory estimate: 2.72 MiB, allocs estimate: 1392.

I ran this script once on dev (1071e68) and once on this branch (f044a36):

using MLJLinearModels, BenchmarkTools, StableRNGs

n, p = 100_000, 200
rng = StableRNG(0)
X = randn(rng, n, p)
θ = randn(rng, p) .* (rand(rng, p) .< 0.3)
y = X * θ .+ 0.5 .+ 0.1 .* randn(rng, n)

for glr in (LassoRegression(0.1),
            ElasticNetRegression(0.1, 0.1; scale_penalty_with_samples=false))
    display(@benchmark fit($glr, $X, $y; solver=FISTA()))
end

versioninfo(), BLAS and package versions

Julia Version 1.13.1
Commit 96ca370cf0 (2026-09-25 19:34 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 20 × 13th Gen Intel(R) Core(TM) i7-13700H
  WORD_SIZE: 64
  LLVM: libLLVM-20.1.8 (ORCJIT, raptorlake)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 20 virtual cores)

BLAS: LBTConfig([ILP64] libopenblas64_.dll), 10 threads
BenchmarkTools v1.8.0, StableRNGs v1.0.4

Checks

• The coefficients are the same before and after for:
  • lasso with FISTA and ISTA, with and without intercept
  • elastic net with scaled and unscaled penalty
  • L1 logistic, huber and multinomial
• The new testset in test/fit/ols-ridge-lasso-elnet.jl checks that:
  • smooth_objective! gives the same value as smooth_objective
  • one call of smooth_objective! or smooth_fg! allocates no length-n vector
  • a FISTA() or ISTA() fit allocates less than 4 length-n vectors (the scratch space holds 3)

---

[tlienart]

thanks for that (and for your issue #187), changes seem right to me, can you bump up the patch release please and then merge?

@ablaom could I leave the path release with you?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.55%. Comparing base (3a2d613) to head (a262cc2).
⚠️ Report is 38 commits behind head on dev.

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #186      +/-   ##
==========================================
- Coverage   94.62%   94.55%   -0.08%     
==========================================
  Files          20       21       +1     
  Lines         875      900      +25     
==========================================
+ Hits          828      851      +23     
- Misses         47       49       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[matthew-liu-1]

@tlienart I have bumped the version to 0.10.4 in Project.toml. The PR is ready to merge.

[ablaom]

Thank you @matthew-liu-1 for your contribution and @tlienart for the prompt review. The quality of the MLJ ecosystem depends crucially on voluntary engagement such as yours.

[ablaom]

Okay, I never tagged a release - or rather, when I tried to I got an error that 0.10.4 was already tagged. @tlienart Do you merge onto master and tag a release at General?

[tlienart]

Yes, sorry if that's not how things should be done these days 🥴

[ablaom]

No all good, I was just doing due diligence.