More accurate streaming `mean` method for arbitrary iterators

[CameronBieganek]

Is there a reason that we can't have a more numerically stable streaming mean method for arbitrary iterators? Here's a sketch of an implementation:

function streaming_mean(itr)
    a, rest = Iterators.peel(itr)
    for (i, x) in enumerate(rest)
        a += (x - a) / (i + 1)
    end
    return a
end

And here's a demonstration of the numerical performance:

julia> using Statistics

julia> A = rand(Float32, 100_000_000);

julia> itr = (rand(Float32) for _ in 1:100_000_000);

julia> mean(A), mean(itr), streaming_mean(itr)
(0.4999884f0, 0.16777216f0, 0.49996638f0)

(I would imagine that there are some tweaks to be made to optimize the performance of streaming_mean, but I hope by now we've learned that correctness (i.e. numerical accuracy) is more important than performance.)

[mikmoore]

This problem is not unique to Statistics.mean. Here's a simpler and more widespread issue:

julia> sum(1f0 for _ in 1:1e8)
1.6777216f7

That said, there are pathological inputs to most (all?) uncompensated sequential accumulations. It's worth ensuring an algorithm in this flavor doesn't simply introduce different pathologies of equal concern.

One unfortunate, although (arguably) not devastating, consequence of an algorithm like this one is the propensity for results like this:

julia> streaming_mean([-1,-1,-1,1,1,1])
-2.7755575615628914e-17

[PallHaraldsson]

A. It's intriguing, and confusing, why your version is not deterministic:

julia> @time streaming_mean(itr)
  0.769110 seconds (1 allocation: 16 bytes)
0.5000117f0

julia> @time streaming_mean(itr)
  0.699333 seconds (1 allocation: 16 bytes)
0.4999693f0

my versions also:
julia> @time streaming_mean3(itr)
  0.369009 seconds (1 allocation: 16 bytes)
0.50004f0


julia> @time streaming_mean2(itr)
  0.666329 seconds (1 allocation: 16 bytes)
0.5000293f0

julia> @time streaming_mean2(itr)
  0.707315 seconds (1 allocation: 16 bytes)
0.49995688f0

julia> function streaming_mean2(itr)
           a, rest = Iterators.peel(itr)
           for (i, x) in enumerate(rest)
               a += (x - a) / Float32(i + 1)
           end
           return a
       end

The 3rd version that's much faster (but ok?): with a += (x - a) * inv(Float32(i + 1))

I can't add @simd, is there a workaround:
ERROR: LoadError: Base.SimdLoop.SimdError("simd loop index must be a symbol")

julia> @time mean(A)  # apparently never in-deterministic:
  0.049539 seconds (1 allocation: 16 bytes)
0.5000132f0

B. FYI:

julia> sum([1f0 for _ in 1:1e8])  # correct, likely your point
1.0f8

I'm not sure (do you know exactly?) why your example gives 1.6777216f7, except likely since it uses an iterator, which was your point, them broken for this. At first I thought you meant to write what I wrote, and since it's legal for me in that context, likely should be here too:

julia> 1f0 for _ in 1:10
┌ Error: JuliaSyntax parser failed — falling back to flisp!
│   exception =
│    BoundsError: attempt to access 0-element Vector{Any} at index [1]

That's non-ideal (on 1.10rc1) though on 1.9.4:

julia> 1f0 for _ in 1:10
ERROR: syntax: extra token "for" after end of expression

[if this should actually ERROR, then why not when wrapped in a sum?]

[mikmoore]

Parentheses are necessary to construct generators in some contexts.

julia> 1f0 for _ in 1:100 # invalid syntax
ERROR: syntax: extra token "for" after end of expression
Stacktrace:
 [1] top-level scope
   @ none:1

julia> (1f0 for _ in 1:100) # generator
Base.Generator{UnitRange{Int64}, var"#104#105"}(var"#104#105"(), 1:100)

julia> [1f0 for _ in 1:100] |> typeof # comprehension
Vector{Float32} (alias for Array{Float32, 1})

The original example is non-deterministic because it is a generator that calls rand to produce values. Adding brackets makes it a comprehension, which is eagerly evaluated and so will produce a consistent value because the random values are only sampled once.

The discrepancy in the results of sum comes from the type of iterable. sum uses pairwise summation over AbstractArray arguments, but sequential summation over general iterators. In this case, the sequential sum eventually saturates due to limited precision.

julia> maxintfloat(Float32) # adding 1 to this value has no effect
1.6777216f7

julia> sum(1f0 for _ in 1:1e8) # sequential sum
1.6777216f7

julia> sum([1f0 for _ in 1:1e8]) # pairwise sum
1.0f8

There isn't anything intrinsically precluding a pairwise summation algorithm over general iterators, although it's a bit more complicated to implement. That could be used to improve the sum algorithm, although it'd be an outstanding question whether we wanted to use the same for mean.

For what it's worth, the streaming_mean algorithm suffers more-or-less the same issue as sequential summation. Eventually, further accumulation will be no-ops because a will be much greater than abs(x-a) / (i + 1) (since x-a is bounded for bounded inputs and i + 1 grows indefinitely. For example:

julia> streaming_mean(Iterators.flatten((Iterators.repeated(0f0, 10^8),Iterators.repeated(1f0, 10^8))))
0.13370937f0 # not 0.5f0

[CameronBieganek]

Interesting. I don't know much about floating point arithmetic. I would just hope that there would be some way that we could improve the numerical accuracy situation for mean(itr) and sum(itr). 🙂

[PallHaraldsson]

sum uses pairwise summation over AbstractArray arguments, but sequential summation over general iterators.

I know the difference, sort of, i.e. downsides of the latter. Is the former always better, and can't it just also be used for generators?

The (absolute error and) relative seems large, i.e. 100000000/16777216 = 5.96 x, and I know you chose the example to highlight that, but it seems a consequence of floating point, and both can be defined "right". It might not be that common to sum up that many numbers (that diverge up or down, not converging to smaller value), I at least not the other isn't immune:

julia> sum([100f0 for _ in 1:1e9])
1.0000001f11

julia> sum([1000f0 for _ in 1:1e9])
9.999999f11

julia> sum([1f0 for _ in 1:1e10])  #  I had to change the example from summing more vlaues, since I think OOM:
Killed

People should be in the know, use Float64 if worried (or should Float32 sum use Float64 then convert to Float32 in the end? It's not bad, except on GPUs where it may not be available or slower). You end up with problems same problem with Float64, even Float128 and BigFloat... at least if you sum up sufficiently large and/or many values.

is non-deterministic because it is a generator that calls rand to produce values.

That's good to know, and plausible, though I do not see it documented at:

help?> Base.Generator

nor do I see a need, good reason, for rand, but if I'm looking at the right place shouldn't it be documented (if/when user-visible consequences). Also since just adding up same values, 1s, it seems the order shouldn't matter. Why is streaming_mean, what is uses, nondeterministic? Because of enumerate presumably, neither documenting (nor that it uses a generator, I suppose it's not because of Iterators.peel, it's docs do not indicate):

An iterator that yields (i, x) where i is a counter starting at 1, and x is the ith value from the given iterator.

That seems to imply a for loop and no rand needed...

[CameronBieganek]

is non-deterministic because it is a generator that calls rand to produce values.

That's good to know, and plausible, though I do not see it documented at:

help?> Base.Generator

nor do I see a need, good reason, for rand, but if I'm looking at the right place shouldn't it be documented (if/when user-visible consequences).

It's nothing in particular about generators, it's just that I put rand() inside the generator, so every time you iterate itr, it generates a new sequence of random numbers. It's not that much different from repeatedly calling rand(100): you get a different sequence of random numbers every time.

[stevengj]

I've changed the title to "more accurate" rather than "numerically stable". According to the mathematical definition, the current algorithm (naive in-order summation) is numerically stable (the wikipedia article on this subject is truly terrible by the way — it never gives a precise definition) (naive summation is backwards stable). The fact that the error grows with n faster than you would like is not the same thing as instability.

[PetrKryslUCSD]

I've changed the title to "more accurate" rather than "numerically stable". According to the mathematical definition, the current algorithm (naive in-order summation) is numerically stable (the wikipedia article on this subject is truly terrible by the way — it never gives a precise definition) (naive summation is backwards stable). The fact that the error grows with n faster than you would like is not the same thing as instability.

It may be true that the current algorithm in mean may be numerically stable, but the proposed algorithm would appear certainly more robust. Is there any reason for not using that algorithm?

[nalimilan]

I'm not sure of the relative advantages of each approach in terms of accuracy, but at least in terms of performance the advantage of the current one is clear. Taking the example from the OP:

julia> using Statistics, BenchmarkTools

julia> A = rand(Float32, 100_000_000);

julia> itr = (x for x in A);

julia> @btime mean(A);
  23.134 ms (1 allocation: 16 bytes)

julia> @btime streaming_mean(A);
  463.581 ms (1 allocation: 16 bytes)

julia> @btime mean(itr);
  106.404 ms (1 allocation: 16 bytes)

julia> @btime streaming_mean(itr);
  474.276 ms (1 allocation: 16 bytes)

Of course, returning a wrong solution quickly isn't useful, but we have to make sure the proposed algorithm is strictly and significantly better than the existing one before switching.

Also, as noted at JuliaStats/Statistics.jl#140 (comment), for small float types, a good solution would be to accumulate in Float64, which is both accurate and fast:

julia> itr2 = (Float64(x) for x in A);

julia> @btime mean(itr2);
  106.298 ms (1 allocation: 16 bytes)

(The difficulty is how to handle heterogeneous arrays, as discussed in that issue.)

[CameronBieganek]

we have to make sure the proposed algorithm is strictly and significantly better than the existing one before switching

I'm no expert on floating point arithmetic, but based on the tradeoffs described by @mikmoore above, it seems like "strictly better" (in terms of numerical accuracy) is too high a bar. Perhaps "more accurate in the large majority of likely use cases" would be a better decision rule. Of course I imagine it's harder to assess whether that rule has been satisfied.

[stevengj]

The proposed algorithm (discussed in this stackoverflow post) does one division per iteration. If you're willing to pay that price, you might as well do Kahan summation, which does extra additions but no division.

For example:

function streaming_mean_kahan(itr)
    s, rest = Iterators.peel(itr)
    c = zero(s)
    count = 1
    for x in rest
        y = x - c
        t = s + y
        c = (t - s) - y
        s = t
        count += 1
    end
    return s / count
end

(which doesn't handle the empty-iterator case, like the original post).
On the test above, this gives

julia> streaming_mean_kahan(itr)
0.4999869f0

on the test above. Even without optimizing it, it's 20% faster than the division-based algorithm on my machine.

julia> @btime mean($A); @btime mean($itr); @btime streaming_mean($itr); @btime streaming_mean_kahan($itr);
  20.456 ms (0 allocations: 0 bytes)
  150.730 ms (0 allocations: 0 bytes)
  479.155 ms (0 allocations: 0 bytes)
  402.276 ms (0 allocations: 0 bytes)

(Both algorithms can easily be further optimized by summing multiple sub-streams at once with SIMD.)

[stevengj]

A simple thing that can increase performance without losing much accuracy is to do Kahan summation "chunked" to every 4 numbers at a time. That is, we sum 4 consecutive numbers naively, then do a Kahan step, then another 4, and so on:

function streaming_mean_kahan4(itr)
    it = iterate(itr)
    it === nothing && return zero(eltype(itr)) / 1
    s_, state = it
    T = typeof(s_ / 1) # accumulate in the precision of the result
    s::T = s_
    c = zero(s)
    count = 1  
    while true
        it = iterate(itr, state)
        it === nothing && return s / count
        x1::T, state = it
        it = iterate(itr, state)
        it === nothing && return (s+x1) / (count+1)
        x2::T, state = it
        it = iterate(itr, state)
        it === nothing && return (s+x1+x2) / (count+2)
        x3::T, state = it
        it = iterate(itr, state)
        it === nothing && return (s+x1+x2+x3) / (count+3)
        x4::T, state = it
        
        # apply Kahan summation chunked 4 elements at a time
        @fastmath x = x1+x2+x3+x4 # (re-associating this line is allowed)
        y = x - c
        t = s + y
        c = (t - s) - y
        s = t
        count += 4
    end
end

On my machine, this is actually slightly faster than our current mean(itr):

julia> @btime streaming_mean_kahan4($itr);
126.799 ms (0 allocations: 0 bytes)

(vs. 150.730 ms for mean(itr), above), while still attaining excellent accuracy:

julia> streaming_mean_kahan4(itr)
0.50002027f0

Here, compared to the exact result, the most accurate method on this test data (on my machine) is the pairwise mean(A), followed by streaming_kahan(itr), followed by streaming_kahan4(itr), followed by the proposed streaming_mean(itr): Correction: see below