From 3830ebe7829ce69641d6ed56615384065fd6dafc Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Thu, 23 Jul 2026 15:15:55 +0200 Subject: [PATCH] feat(nn): optional bias in Linear MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Linear's initBias is now nullable with a null default, creating a bias-less projection (y = x W^T) — the equivalent of PyTorch's nn.Linear(bias=False), needed for GPT-style qkv_bias=False projections and weight-tied output heads. A bias-less layer registers only its weight parameter, so parameter counts and checkpoints match architectures defined without bias. Adds a biasOrNull() helper next to the throwing bias() accessor. Note for Java callers: the @JvmOverloads overload set changes — the 4-arg (in, out, weights, bias) convenience overload is replaced by (in, out, weights); Kotlin callers bind the full constructor and are unaffected. --- .../sk/ainet/exec/nn/LinearNoBiasTest.kt | 61 +++++++++++++++++++ .../api/android/skainet-lang-core.api | 4 +- .../api/jvm/skainet-lang-core.api | 4 +- .../kotlin/sk/ainet/lang/nn/Linear.kt | 46 ++++++++------ .../lang/nn/topology/ModuleParameters.kt | 4 ++ .../ainet/lang/nn/LinearOptionalBiasTest.kt | 57 +++++++++++++++++ 6 files changed, 156 insertions(+), 20 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/LinearNoBiasTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/LinearOptionalBiasTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/LinearNoBiasTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/LinearNoBiasTest.kt new file mode 100644 index 00000000..ff5c48be --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/LinearNoBiasTest.kt @@ -0,0 +1,61 @@ +package sk.ainet.exec.nn + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.nn.Linear +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 + +/** Value-level behavior of bias-less Linear on the real CPU backend. */ +class LinearNoBiasTest { + + private val ctx = DirectCpuExecutionContext() + + @Test + fun forward_without_bias_is_a_pure_projection() { + // W = [[1, 2], [3, 4]] (rows = out features), x = [[1, 1], [2, 0]] + val layer = Linear( + inFeatures = 2, outFeatures = 2, name = "lin", + initWeights = ctx.fromFloatArray(Shape(2, 2), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)), + initBias = null, + ) + val x = ctx.fromFloatArray(Shape(2, 2), FP32::class, floatArrayOf(1f, 1f, 2f, 0f)) + + val y = layer.forward(x, ctx) + + // y = x @ W^T: [[1*1+1*2, 1*3+1*4], [2*1+0*2, 2*3+0*4]] = [[3, 7], [2, 6]] + assertContentEquals(floatArrayOf(3f, 7f, 2f, 6f), y.data.copyToFloatArray()) + } + + @Test + fun no_bias_equals_zero_bias() { + val w = ctx.fromFloatArray(Shape(3, 2), FP32::class, floatArrayOf(0.5f, -1f, 2f, 0.25f, -0.75f, 1.5f)) + val noBias = Linear(2, 3, "a", initWeights = w, initBias = null) + val zeroBias = Linear( + 2, 3, "b", + initWeights = w, + initBias = ctx.fromFloatArray(Shape(3), FP32::class, FloatArray(3)), + ) + val x = ctx.fromFloatArray(Shape(4, 2), FP32::class, FloatArray(8) { (it - 3).toFloat() }) + + assertContentEquals( + zeroBias.forward(x, ctx).data.copyToFloatArray(), + noBias.forward(x, ctx).data.copyToFloatArray(), + ) + } + + @Test + fun vector_input_works_without_bias() { + val layer = Linear( + inFeatures = 3, outFeatures = 2, name = "lin", + initWeights = ctx.fromFloatArray(Shape(2, 3), FP32::class, floatArrayOf(1f, 0f, 1f, 0f, 1f, 0f)), + initBias = null, + ) + val x = ctx.fromFloatArray(Shape(3), FP32::class, floatArrayOf(1f, 2f, 3f)) + + val y = layer.forward(x, ctx) + + assertContentEquals(floatArrayOf(4f, 2f), y.data.copyToFloatArray()) + } +} diff --git a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api index c9c152e6..1fad544e 100644 --- a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api @@ -755,10 +755,11 @@ public abstract class sk/ainet/lang/nn/InternalMixedPrecisionModule : sk/ainet/l } public final class sk/ainet/lang/nn/Linear : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;)V public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Z)V public synthetic fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V - public fun (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V + public fun (IILsk/ainet/lang/tensor/Tensor;)V public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; public fun getParams ()Ljava/util/List; @@ -2034,6 +2035,7 @@ public abstract interface class sk/ainet/lang/nn/topology/ModuleParameters { public final class sk/ainet/lang/nn/topology/ModuleParametersKt { public static final fun bias (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$BiasParameter; + public static final fun biasOrNull (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$BiasParameter; public static final fun by (Ljava/util/List;Ljava/lang/String;)Lsk/ainet/lang/nn/topology/ModuleParameter; public static final fun weights (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$WeightParameter; public static final fun zeroGrad (Ljava/lang/Iterable;)V diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 4c7cbad8..6b22f0f4 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -990,10 +990,11 @@ public final class sk/ainet/lang/nn/LayerScale : sk/ainet/lang/nn/Module, sk/ain } public final class sk/ainet/lang/nn/Linear : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters { + public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;)V public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Z)V public synthetic fun (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V - public fun (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V + public fun (IILsk/ainet/lang/tensor/Tensor;)V public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; public fun getParams ()Ljava/util/List; @@ -2262,6 +2263,7 @@ public abstract interface class sk/ainet/lang/nn/topology/ModuleParameters { public final class sk/ainet/lang/nn/topology/ModuleParametersKt { public static final fun bias (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$BiasParameter; + public static final fun biasOrNull (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$BiasParameter; public static final fun by (Ljava/util/List;Ljava/lang/String;)Lsk/ainet/lang/nn/topology/ModuleParameter; public static final fun weights (Ljava/util/List;)Lsk/ainet/lang/nn/topology/ModuleParameter$WeightParameter; public static final fun zeroGrad (Ljava/lang/Iterable;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt index 8bf3019b..e55cbfa1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Linear.kt @@ -6,18 +6,24 @@ import sk.ainet.lang.tensor.* import sk.ainet.lang.types.DType import sk.ainet.lang.nn.topology.ModuleParameter import sk.ainet.lang.nn.topology.ModuleParameters -import sk.ainet.lang.nn.topology.bias +import sk.ainet.lang.nn.topology.biasOrNull import sk.ainet.lang.nn.topology.weights /** * Linear layer (a.k.a. fully connected dense layer). This layer applies a linear transformation to the input data. * The weights and biases are learned during training. * + * Passing `null` for [initBias] creates a bias-less projection (`y = x W^T`), + * the equivalent of PyTorch's `nn.Linear(..., bias=False)` — used for example + * by GPT-style Q/K/V projections (`qkv_bias=False`) and weight-tied output + * heads. A bias-less layer registers only the weight parameter, so parameter + * counts and checkpoints match architectures defined without bias. + * * @param inFeatures Number of input features * @param outFeatures Number of output features * @param name Name of the module * @param initWeights Initial weights - * @param initBias Initial bias + * @param initBias Initial bias, or `null` for a layer without bias */ public class Linear @kotlin.jvm.JvmOverloads constructor( @@ -25,7 +31,7 @@ public class Linear @kotlin.jvm.JvmOverloads constructor( outFeatures: Int, override val name: String = "Linear", initWeights: Tensor, - initBias: Tensor, + initBias: Tensor? = null, public val trainable: Boolean = true ) : Module(), ModuleParameters { @@ -35,23 +41,26 @@ public class Linear @kotlin.jvm.JvmOverloads constructor( require(initWeights.rank == 2 && wShape[0] == outFeatures && wShape[1] == inFeatures) { "Linear($name): initWeights shape must be [outFeatures, inFeatures]=[${outFeatures}, ${inFeatures}], but was ${initWeights.shape}" } - // Validate bias shape: allow [outFeatures] or [1, outFeatures] - val bShape = initBias.shape.dimensions - val biasOk = when (initBias.rank) { - 1 -> bShape[0] == outFeatures - 2 -> bShape[0] == 1 && bShape[1] == outFeatures - else -> false - } - require(biasOk) { - "Linear($name): initBias shape must be [outFeatures] or [1, outFeatures] with outFeatures=${outFeatures}, but was ${initBias.shape}" + // Validate bias shape (when present): allow [outFeatures] or [1, outFeatures] + if (initBias != null) { + val bShape = initBias.shape.dimensions + val biasOk = when (initBias.rank) { + 1 -> bShape[0] == outFeatures + 2 -> bShape[0] == 1 && bShape[1] == outFeatures + else -> false + } + require(biasOk) { + "Linear($name): initBias shape must be [outFeatures] or [1, outFeatures] with outFeatures=${outFeatures}, but was ${initBias.shape}" + } } } - override val params: List> = listOf( - ModuleParameter.WeightParameter("$name.weight", initWeights, trainable), - ModuleParameter.BiasParameter("$name.bias", initBias, trainable) - ) - + override val params: List> = buildList { + add(ModuleParameter.WeightParameter("$name.weight", initWeights, trainable)) + if (initBias != null) { + add(ModuleParameter.BiasParameter("$name.bias", initBias, trainable)) + } + } override val modules: List> get() = emptyList() @@ -59,10 +68,11 @@ public class Linear @kotlin.jvm.JvmOverloads constructor( @Suppress("UNCHECKED_CAST") override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { val weight = params.weights().value - val bias = params.bias().value + val bias = params.biasOrNull()?.value val weightTransposed = weight.t() val matmulResult = input.matmul(weightTransposed) + if (bias == null) return matmulResult // If input is a 1D vector, ensure bias is also 1D to avoid broadcasting to [1, out] val result = if (input.rank == 1 && bias.rank == 2 && bias.shape.dimensions[0] == 1) { diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/ModuleParameters.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/ModuleParameters.kt index f0397b6a..7916be37 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/ModuleParameters.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/topology/ModuleParameters.kt @@ -74,6 +74,10 @@ public fun List>.bias(): ModuleParameter.Bi this.filterIsInstance>() .firstOrNull() ?: throw NoSuchElementException("No bias parameter found!") +// Returns the first BiasParameter or null for modules without a bias (e.g. Linear(bias = null)). +public fun List>.biasOrNull(): ModuleParameter.BiasParameter? = + this.filterIsInstance>().firstOrNull() + // Returns the first WeightParameter or throws a NoSuchElementException if none is found. public fun List>.weights(): ModuleParameter.WeightParameter = this.filterIsInstance>() diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/LinearOptionalBiasTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/LinearOptionalBiasTest.kt new file mode 100644 index 00000000..4e17201a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/LinearOptionalBiasTest.kt @@ -0,0 +1,57 @@ +package sk.ainet.lang.nn + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertNull +import sk.ainet.lang.nn.topology.bias +import sk.ainet.lang.nn.topology.biasOrNull +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 + +/** Parameter registration for bias-less Linear layers (value behavior is covered in backend-cpu). */ +class LinearOptionalBiasTest { + + private val ctx = DefaultNeuralNetworkExecutionContext() + + private fun weights(out: Int, inF: Int) = + ctx.fromFloatArray(Shape(out, inF), FP32::class, FloatArray(out * inF) { 1f }) + + @Test + fun without_bias_only_the_weight_parameter_is_registered() { + val layer = Linear( + inFeatures = 3, outFeatures = 2, name = "lin", + initWeights = weights(2, 3), initBias = null, + ) + assertEquals(1, layer.params.size) + assertEquals("lin.weight", layer.params.single().name) + assertNull(layer.params.biasOrNull()) + assertFailsWith { layer.params.bias() } + } + + @Test + fun with_bias_both_parameters_are_registered() { + val bias = ctx.fromFloatArray(Shape(2), FP32::class, floatArrayOf(0f, 0f)) + val layer = Linear( + inFeatures = 3, outFeatures = 2, name = "lin", + initWeights = weights(2, 3), initBias = bias, + ) + assertEquals(2, layer.params.size) + assertEquals("lin.bias", layer.params.bias().name) + } + + @Test + fun trainable_parameter_count_reflects_the_missing_bias() { + val withBias = Linear( + inFeatures = 4, outFeatures = 3, name = "a", + initWeights = weights(3, 4), + initBias = ctx.fromFloatArray(Shape(3), FP32::class, FloatArray(3)), + ) + val withoutBias = Linear( + inFeatures = 4, outFeatures = 3, name = "b", + initWeights = weights(3, 4), initBias = null, + ) + assertEquals(15, withBias.trainableParameters().sumOf { it.value.volume }) + assertEquals(12, withoutBias.trainableParameters().sumOf { it.value.volume }) + } +}