How to build a tensorflow.js model to predict the square of a number?

Question

I'm trying to create a tensorflow.js model (With javascript) to predict the square of a number, so I have a function to create the xData(x) and the yData(x²).

I've tried to use many hiden layers with different units and to train with thousand epochs but the loss is realy hight and the predictions are not good at all.

here is my code:

const tf=require("@tensorflow/tfjs-node");

let dataX=[
//0, 1, 2, 3, ..., 1000;
];

let dataY=[
//0, 1, 4, 9, ..., 1000000;
]
for (let i = 0; i < 1000; i++) {
    dataX.splice(0,0, i);
    dataY.splice(0,0, i**2);
}
class AI{
    compile(){
        const model=tf.sequential();

        //Input & hidenA layer
        model.add(tf.layers.dense({
            inputShape: [1],
            units: 32,
            activation: 'sigmoid',
            useBias: true
        }));
        //HidenB layer
        model.add(tf.layers.dense({
            units: 32,
            activation: 'sigmoid',
            useBias: true
        }))


        //Output layer

        model.add(tf.layers.dense({
            units: 1,
            activation: 'sigmoid',
            useBias: true
        }))
        let sgdOpt=tf.train.sgd(0.00001)
        model.compile({
            loss: 'meanSquaredError',
            optimizer: sgdOpt
        })

        return model;
    }

    run(){
        const model=this.compile();
        const xs=tf.tensor(dataX)

        const ys=tf.tensor(dataY)
try{
    model.summary();
model.fit(xs, ys, {
            epochs: 1000,
        }).then(()=>{
            const data=tf.tensor([10]);

            const prediction=model.predict(data);

            prediction.print();
        })
}catch(e){

}
        
    }
}

const ai=new AI();

ai.run();

Answer 1

I know this is going to sound strange, but, I believe you're missing a few critical features where the AI could derive an abstract method.

I'm going to use https://playground.tensorflow.org/ to try to illustrate what I'm thinking. They have a classification problem ( I know yours is regression ) that I think really illustrates some key points.

Problem 1: Sigmoid

Firstly, you're using sigmoid activation at each hidden layer. Sigmoid has it's share of issues with training.

After 1,000 epochs with a sigmoid activation, you'll notice the network derives a pretty linear solution to this circular problem:

Link to playground

Simply by switching the activation to ReLU, you can see a much better solution in 1/10th the epochs.

Problem 2: Linear-Features

As you see from the above image, the solution with ReLU was very pointy. It should have ideally created a circle. You're having the same issues in the graph of a line vs a number squared. I actually cover this exact problem in Chapter 8 of my book, and while I can teach the AI to get a perfectly acceptable answer for 7**2 === ~~49, it doesn't do well for larger numbers.

The reason, I believe is quite simple. If I want the model to be able to determine a square of a number, I will need to derive features that are going to help it see that difference.

Let's go back to our playground example. See how it has sharp edges? What if we give some extra features to our model input, like squares of each value? We see it instantly (in 40 epochs) has a solid answer!

So by creating a feature input of squares the model can derive a smooth non-linear classification. However, when your model is so simple that squaring your input would be the answer, you're limited on ways you can add features. This is why I'm fine with my small example being "kind of" accurate. It's not a real-world problem. But if you really want to make a model that squares numbers without adding that as a feature, you'll have to derive features as input that the model could determine answers to.